import plotly.graph_objects as go
import plotly.express as px
import numpy as np

def plot_3D_invfitness(trait, fitness, resident, range, color="RdBu"):
    X, Y = np.meshgrid(trait, trait)
    f_projection = (np.min(fitness) - 2) * np.ones(fitness.shape)
    axis = dict(
        showbackground=True,
        backgroundcolor="rgb(230, 230,230)",
        showgrid=False,
        zeroline=False,
        showline=False,
    )

    layout = go.Layout(
        autosize=False,
        width=600,
        height=500,
        scene=dict(
            xaxis=dict(axis),
            yaxis=dict(axis),
            zaxis=dict(axis, range=range),
            aspectratio=dict(x=1, y=1, z=1),
            xaxis_title="Resident trait",
            yaxis_title="Mutant trait",
            zaxis_title="Invasion fitness",
        ),
    )
    x_projection = resident * np.ones(fitness.shape)
    fitness_surface = go.Surface(x=X, y=Y, z=fitness, colorscale=color)
    PIP = go.Surface(
        x=X,
        y=Y,
        z=f_projection,
        surfacecolor=(fitness > 0),
        colorscale="Greens",
        showlegend=False,
        showscale=False,
    )
    slice = go.Surface(
        x=x_projection,
        y=Y,
        z=(fitness * 1e5),
        surfacecolor=x_projection,
        colorscale="Greys",
        opacity=0.5,
        showlegend=False,
        showscale=False,
    )
    fig = go.Figure(
        data=[
            fitness_surface,
            PIP,
            slice,
        ],
        layout=layout,
    )
    return fig


def plot_invasionfitness(zm, zlist, fitness_func, pars, range):
    inv_fitness = fitness_func(zm, zlist, pars)

    fig = px.line(
        x=zlist, y=inv_fitness, labels={
            "x": "Mutant trait value (z)", "y": "Invasion fitness"}
    )
    fig.add_vline(x=zm, line_dash="dashdot")
    fig.add_hline(y=0, line_dash="dash")
    fig.update_layout(
        title="Interactive invasion process",
        xaxis=dict(range=[0, zlist[-1]], autorange=False),
        yaxis=dict(range=range, autorange=False), autosize=False,
        width=450,
        height=400
    )
    return fig


def make_interactive_video(z_start, z_end, step, zlist, fitness_func, pars, range):
    inv_vid = []
    for z_val in np.arange(z_start, z_end, step):
        inv_vid.append(fitness_func(z_val, zlist, pars))
    vid = go.Figure(
        data=[
            go.Line(x=zlist, y=fitness_func(
                z_start, zlist, pars), name="invasion fitness"),
            go.Line(
                x=zlist,
                y=np.zeros(len(zlist)),
                line=dict(color="black", width=1, dash="dash"),
                name="Invasion threshold",
            ),
            go.Scatter(
                x=[z_start] * 10,
                y=np.linspace(-2, 2, 10),
                mode="lines",
                line=dict(color="black", dash="dashdot"),
                name="Resident trait value",
            ),
        ],
        layout=go.Layout(
            title="Invasion process video", autosize=False,
            width=550,
            height=500,
            xaxis=dict(range=[0, zlist[-1]], autorange=False),
            yaxis=dict(range=range, autorange=False),
            xaxis_title="Mutant trait value",
            updatemenus=[
                dict(type="buttons", buttons=[
                     dict(label="Play", method="animate", args=[None, {"frame": {"duration": 500, "redraw": False},
                                                                       "fromcurrent": True, "transition": {"duration": 300,
                                                                                                           "easing": "quadratic-in-out"}}]),
                     dict(label="Pause", method="animate", args=[[None], {"frame": {"duration": 0, "redraw": False},
                                                                          "mode": "immediate",
                                                                          "transition": {"duration": 0}}])]),

            ],
        ),
        frames=[
            go.Frame(
                data=[
                    go.Line(x=zlist, y=i),
                    go.Line(
                        x=zlist,
                        y=np.zeros(len(zlist)),
                        line=dict(color="black", dash="dash"),
                    ),
                    go.Scatter(
                        x=[z_val] * 10,
                        y=np.linspace(-2, 2, 10),
                        line=dict(color="black", dash="dashdot"),
                        mode="lines",
                    ),
                ]
            )
            for i, z_val in zip(inv_vid, np.arange(z_start, z_end, step))
        ],
    )
    return vid


def plot_PIP(zlist, fitness_func, pars):
    X, Y = np.meshgrid(zlist, zlist)
    inv_fitness3D = fitness_func(X, Y, pars)
    fig = go.Figure(
        data=go.Contour(
            x=zlist,
            y=zlist,
            z=inv_fitness3D,
            colorscale="PRGn",
            showscale=False,
            contours=dict(
                start=-20,
                end=0,
                size=10,
            ),
        )
    )
    fig.update_layout(
        autosize=False,
        width=400,
        height=500,
        xaxis_title="Resident trait",
        yaxis_title="Mutant trait",
    )
    return fig