import plotly.graph_objects as go
import numpy as np
from guang.interesting import Lorenz
from guang.Utils.plotly import Scatter3d
from guang.Utils.animation.utils import ExtremeValue


def getRange(x, dx=0.02):
    delta = np.max(x) - np.min(x)
    delta = delta * dx
    return np.min(x) - delta, np.max(x) + delta


def lorenz_Traces(x0=1,
                  y0=3,
                  z0=5,
                  lambd=0.01,
                  k=3,
                  t_start=0,
                  t_end=10,
                  all_step=900):

    t = np.linspace(t_start, t_end, all_step)
    Traces = []
    for i in range(k):
        trace = Lorenz.Trace(t=t,
                             xyz=[
                                 x0 + lambd * np.random.rand(),
                                 y0 + lambd * np.random.rand(),
                                 z0 + lambd * np.random.rand()
                             ])
        Traces.append(trace)
    return Traces


class Draw3D:
    def __init__(self):
        self.Traces = []
        self.initial_data = {}

    def get_extrem(self):
        Value = ExtremeValue()
        for trace in self.Traces:
            x, y, z = trace[:, 0], trace[:, 1], trace[:, 2]
            xm, xM = getRange(x)
            ym, yM = getRange(y)
            zm, zM = getRange(z)
            xm, ym, zm = Value.get_min(x=xm, y=ym, z=zm)
            xM, yM, zM = Value.get_max(x=xM, y=yM, z=zM)
        return xm, xM, ym, yM, zm, zM

    def get_frames(self, total_frames=100, period=1):
        """
        :param total_frames: total frames
        :param period: value interval:(0, 1] means the percentage of the total points.
        """
        self.k = len(self.Traces)
        N = len(self.Traces[0][:, 0])
        step = N // total_frames
        i_start, i_end = 0, 0
        counts_period = int(N * period)
        frames = []
        for i in range(total_frames):
            i_end += step
            i_start = i_end - counts_period if i_end > counts_period else 0
            data = []

            for trace in self.Traces:
                x, y, z = trace[:, 0], trace[:, 1], trace[:, 2]
                data_x = x[i_start:i_end]
                data_y = y[i_start:i_end]
                data_z = z[i_start:i_end]
                data.append(
                    go.Scatter3d(x=data_x,
                                 y=data_y,
                                 z=data_z,
                                 mode="lines+markers",
                                 line=dict(color=None, width=1.5),
                                 marker=dict(size=1.5,
                                             color=z,
                                             colorscale='Viridis')))

            if i == 0:
                self.initial_data.setdefault('data', data)
            frames.append(
                go.Frame(
                    data=data,
                    traces=[i for i in range(self.k)
                            ]  ####THIS IS THE LINE THAT MUST BE INSERTED
                ))

        return frames

    def get_layout(self):
        xm, xM, ym, yM, zm, zM = self.get_extrem()
        layout = go.Layout(
            width=780,
            height=780,
            scene={
                'xaxis': {
                    'range': [xm, xM],
                    "nticks": 10,
                    'autorange': False,
                    "zeroline": False
                },
                # , 'rangemode': 'tozero', 'tickmode': "linear", 'tick0': -5, 'dtick': 1},
                'yaxis': {
                    'range': [ym, yM],
                    "nticks": 10,
                    'autorange': False,
                    "zeroline": False
                },
                # , 'rangemode': 'tozero', 'tickmode': "linear", 'tick0': -5, 'dtick': 1},
                'zaxis': {
                    'range': [zm, zM],
                    "nticks": 10,
                    'autorange': False,
                    "zeroline": False
                },
                # , 'rangemode': 'tozero', 'tickmode': "linear", 'tick0': -5, 'dtick': 1},
                'aspectmode': 'cube',
            },
            title="Lorenz Curve",
            hovermode="closest",
            updatemenus=[
                dict(type="buttons",
                     buttons=[
                         dict(label="Play",
                              method="animate",
                              args=[
                                  None,
                                  dict(frame=dict(duration=20, redraw=True),
                                       transition=dict(duration=0),
                                       fromcurrent=True,
                                       mode='immediate')
                              ])
                     ])
            ])
        return layout

    def anima(self, total_frames=200, period=0.05):
        frames = self.get_frames(total_frames=total_frames, period=period)
        data = self.initial_data['data']
        fig = go.Figure(data=data, layout=self.get_layout(), frames=frames)
        return fig


class Draw2D:
    def __init__(self):
        self.Traces = []
        self.initial_data = {}

    def get_extrem(self):
        Value = ExtremeValue()
        for trace in self.Traces:
            x, y = trace[:, 0], trace[:, 1]
            xm, xM = getRange(x)
            ym, yM = getRange(y)

            xm, ym = Value.get_min(x=xm, y=ym)
            xM, yM = Value.get_max(x=xM, y=yM)
        return xm, xM, ym, yM

    def get_frames(self, total_frames=100, period=1):
        """
        :param total_frames: total frames
        :param period: value interval:(0, 1] means the percentage of the total points.
        """
        self.k = len(self.Traces)

        N = len(self.Traces[0][:, 0])
        step = N // total_frames
        i_start, i_end = 0, 0
        counts_period = int(N * period)
        frames = []
        for i in range(total_frames):
            i_end += step
            i_start = i_end - counts_period if i_end > counts_period else 0
            data = []

            for trace in self.Traces:
                x, y = trace[:, 0], trace[:, 1]
                data_x = x[i_start:i_end]
                data_y = y[i_start:i_end]
                data.append(
                    go.Scatter(x=data_x,
                               y=data_y,
                               mode="lines",
                               line=dict(color=None, width=2)))

            if i == 0:
                self.initial_data.setdefault('data', data)
            frames.append(
                go.Frame(
                    data=data,
                    traces=[i for i in range(self.k)
                            ]  #THIS IS THE LINE THAT MUST BE INSERTED
                ))
        return frames

    def get_layout(self,
                   width=980,
                   height=880,
                   duration=10,
                   title="Lorenz X-Y 2D Curve",
                   redraw=False,
                   mode='immediate'):
        xm, xM, ym, yM = self.get_extrem()

        # height = int(width* (yM-ym)/(xM-xm)) + 100
        layout = go.Layout(
            width=width,
            height=height,
            xaxis=dict(range=[xm, xM], autorange=False, zeroline=False),
            yaxis=dict(range=[ym, yM], autorange=False, zeroline=False),
            title=title,
            hovermode="closest",
            updatemenus=[
                dict(type="buttons",
                     buttons=[
                         dict(label="Play",
                              method="animate",
                              args=[
                                  None,
                                  dict(frame=dict(duration=duration,
                                                  redraw=redraw),
                                       transition=dict(duration=0),
                                       fromcurrent=True,
                                       mode=mode)
                              ])
                     ])
            ])
        return layout

    def anima(self, total_frames=200, period=0.05, width=990, height=900):
        frames = self.get_frames(total_frames=total_frames, period=period)
        data = self.initial_data['data']
        layout = self.get_layout(width=width, height=height)
        fig = go.Figure(data=data, layout=layout, frames=frames)
        return fig


def Traces2d(r=10, all_step=1000):

    theta = np.linspace(0, 2 * np.pi, all_step)[..., None]
    x = r * np.cos(theta)
    y = r * np.sin(theta)
    Trace = []
    Trace.append(np.concatenate((x, y), axis=1))
    Trace.append(np.concatenate((theta + r, x), axis=1))
    Trace.append(np.concatenate((theta - r, y), axis=1))
    return Trace


if __name__ == "__main__":
    draw2d = Draw2D()
    draw2d.Traces = Traces2d(r=10)
    fig = draw2d.anima(total_frames=200, period=1)
    fig.show()

    # draw3d = Draw3D()
    # draw3d.Traces = lorenz_Traces(x0=5, y0=6, z0=10, lambd=10,
    #                               k=5, all_step=900)
    # ---------------------------------------------------- option start
    # xm, xM, ym, yM, zm, zM = draw3d.get_extrem()
    # draw3d.initial_data =  [go.Scatter3d(x=[xm], y=[ym], z=[zm]),
    #                     go.Scatter3d(x=[xM], y=[yM], z=[zM]),
    #                     go.Scatter3d(x=[xm], y=[yM], z=[zm]), ]
    # ---------------------------------------------------- option end

    # fig = draw3d.anima(total_frames=200, period=0.05)
    # fig.show()
