import numpy as np
import matplotlib
from matplotlib import pyplot as plt
from matplotlib.animation import FuncAnimation


def ss_var_aniplot(theta, K, phi, R, tau, interval=500):
    
    #Initialise a figure with five subplots
    fig, (ax1, ax2, ax3, ax4, ax5) = plt.subplots(5,1)
    fig.subplots_adjust(hspace=0.8)
    
    nt, ns = K.shape[0], K.shape[1]
    xs = np.linspace(0, ns, ns)


    #Initialise line objects with independent axes
    line1, = ax1.plot([], [], label='theta')
    line2, = ax2.plot([], [], label='K')
    line3, = ax3.plot([], [], label='phi')
    line4, = ax4.plot([], [], label='R')
    line5, = ax5.plot([], [], label='tau')

    line = [line1, line2, line3, line4, line5]
    
    ax1.set(ylim=(np.nanmin(theta),np.nanmax(theta)))
    ax2.set(ylim=(np.nanmin(K),np.nanmax(K)))
    ax3.set(ylim=(np.nanmin(phi),np.nanmax(phi)))
    ax4.set(ylim=(np.nanmin(R[-1,100:700]),np.nanmax(R[-1,100:700])))
    ax5.set(ylim=(np.nanmin(tau),np.nanmax(tau)))
    
    for ax in [ax1, ax2, ax3, ax4, ax5]:
        ax.legend(loc='upper right')
        ax.set(xlim=(0,ns))
        
    def animate(i):
        line[0].set_data(xs, theta[i, :])
        line[1].set_data(xs, K[i, :])
        line[2].set_data(xs, phi[i, :])
        line[3].set_data(xs, R[i, :])
        line[4].set_data(xs, tau[i, :])    
        ax1.set_title('Time step: ' + str(i))
        return line
    anim = FuncAnimation(fig, animate, interval=interval, frames=nt-1, repeat=True, blit=False)
    return anim


def anim_2Dplot(data, ymin=-1000, ymax=1000, interval=200):
    
    #Initialise figure and set limits
    fig, ax = plt.subplots(figsize=(4, 3))
    ax.set(ylim=(ymin, ymax))
    
    #For a single line
    if data.ndim == 2:
        #Extract the number of lines to plot, number of time iterations and the number of samples in each lines
        nt, ns = data.shape[0], data.shape[1]
        
        #Initialise 2Dline first iteration
        xs = np.linspace(0, ns, ns)
        line = ax.plot(xs, data[0, :], color='k')[0]
        
        def animate(i):
            line.set_ydata(data[i, :])
            ax.set_title('Time step ' + str(i))
    
    #For a set of lines       
    else:        
        #Extract the number of lines to plot, number of time iterations and the number of samples in each lines
        n_channel, nt, ns = len(data), data.shape[1], data.shape[2]
    
        #Initialise iterable 2Dline instances
        xs = np.linspace(0, ns, ns)
        lines = [ax.plot(xs, data[i, 0, :])[0] for i in range(n_channel)]
        
        #Create animation function to cycle through the data
        def animate(t):
            ax.set_title('Time step ' + str(t))
            for i, line in enumerate(lines):
                line.set_ydata(data[i, t, :])
            
    #Create the animation
    anim = FuncAnimation(fig, animate, interval=interval, frames=nt-1, repeat=True, blit=True)
    return anim

def xyzdif_aniplots(curve1, curve2, interval=500):
    
    #Initialise a figure with five subplots
    fig, (ax1, ax2, ax3) = plt.subplots(3,1)
    fig.subplots_adjust(hspace=1.2)
    
    nt, ns, ax = curve1.shape[0], curve1.shape[1], curve1.shape[2]
    xs = np.linspace(0, ns, ns)

    #Initialise line objects with independent axes
    x_FS, = ax1.plot(xs, xs, label='Frenet-Serret')
    x_MT, = ax1.plot(xs, xs, label='Transformation Matrix')

    y_FS, = ax2.plot(xs, xs, label='Frenet-Serret')
    y_MT, = ax2.plot(xs, xs, label='Transformation Matrix')
    
    z_FS, = ax3.plot(xs, xs, label='Frenet-Serret')
    z_MT, = ax3.plot(xs, xs, label='Transformation Matrix')
    
    line = [x_FS, x_MT, y_FS, y_MT, z_FS, z_MT]
    
    ax1.set(ylim=(0,210))
    ax2.set(ylim=(-1,2))
    ax3.set(ylim=(0,2.5))
    
    ax1.legend(loc='center left', bbox_to_anchor=(-0.1, 1))
    ax1.set_title('x-values')
    ax2.set_title('y-values')
    ax3.set_title('z-values')
    ax3.set_yticks((0,2.5))
   
    for ax in [ax1, ax2, ax3]:
        ax.set(xlim=(0,ns))
        ax.set_xlabel('gauge number')
        ax.set_ylabel('cm')

        
    def animate(i):
        line[0].set_ydata(curve1[i, :, 0])
        line[1].set_ydata(curve2[i, :, 0])

        line[2].set_ydata(curve1[i, :, 1])
        line[3].set_ydata(curve2[i, :, 1])

        line[4].set_ydata(curve1[i, :, 2])
        line[5].set_ydata(curve2[i, :, 2])
       
        ax1.set_title('Time step: ' + str(i))
        return line
    
    anim = FuncAnimation(fig, animate, interval=interval, frames=nt-1, repeat=True, blit=False)
    return anim




################################# TRANSFORM TO 3D ANIMATED PLOTS ##############################
def plot_curve(P, P1=None, P2=None, TNB=None, name=None, name1=None, name2=None):
    if TNB == 'TNB':
        r_0, T, N, B = get_rTNB(P)

        N_0 = N[0, :]
        T_0 = T[0, :]
        B_0 = B[0, :]

        Path = pd.DataFrame({
            "x": P[:, 0],
            "y": P[:, 1],
            "z": P[:, 2]})

        Tpd = pd.DataFrame({
            "x": [r_0[0], T_0[0]],
            "y": [r_0[1], T_0[1]],
            "z": [r_0[2], T_0[2]]})

        Npd = pd.DataFrame({
            "x": [r_0[0], N_0[0]],
            "y": [r_0[1], N_0[1]],
            "z": [r_0[2], N_0[2]]})

        Bpd = pd.DataFrame({
            "x": [r_0[0], B_0[0]],
            "y": [r_0[1], B_0[1]],
            "z": [r_0[2], B_0[2]]})

        Path['Path'] = name
        Tpd['Path'] = 'T'
        Npd['Path'] = 'N'
        Bpd['Path'] = 'B'

        df = pd.concat([Tpd, Npd, Bpd, Path])

    else:
        P = np.abs(P)
        Path = pd.DataFrame({
            "x": P[:, 0],
            "y": P[:, 1],
            "z": P[:, 2]})
        Path['Path'] = name
        df = Path

    if P1 is not None:
        P1 = np.abs(P1)
        Path1 = pd.DataFrame({
            "x": P1[:, 0],
            "y": P1[:, 1],
            "z": P1[:, 2]})
        Path1['Path'] = name1
        df = pd.concat([df, Path1])

    if P2 is not None:
        P2 = np.abs(P2)
        Path2 = pd.DataFrame({
            "x": P2[:, 0],
            "y": P2[:, 1],
            "z": P2[:, 2]})
        Path2['Path'] = name2
        df = pd.concat([df, Path2])
    fig = px.line_3d(df, x="x", y="y", z="z", color="Path")
    #fig.update_layout(
    #    scene=dict(
    #        xaxis=dict(nticks=4, range=[0, 1], ),
    #        yaxis=dict(nticks=4, range=[0, 1], ),
    #        zaxis=dict(nticks=4, range=[0, 50], ), ),
    #    width=700,
    #    margin=dict(r=20, l=10, b=10, t=10))
    fig.show()

    return

def dif_xyz(P0, P1, P2, name):
    n = int(P1.shape[0])
    Coumpound_dif_1 = np.zeros((n))
    Coumpound_dif_2 = np.zeros((n))

    for i in np.arange(n - 1):
        Coumpound_dif_1[i + 1] = Coumpound_dif_1[i] + np.abs(P0[i, 0] - P1[i, 0]) + np.abs(
            P0[i, 1] - P1[i, 1]) + np.abs(P0[i, 2] - P1[i, 2])
        Coumpound_dif_2[i + 1] = Coumpound_dif_2[i] + np.abs(P0[i, 0] - P2[i, 0]) + np.abs(
            P0[i, 1] - P2[i, 1]) + np.abs(P0[i, 2] - P2[i, 2])

    fig, axs = plt.subplots(2, 2)
    fig.suptitle('Difference from true path')
    axs[0, 0].plot(P0[:, 0] - P1[:-1, 0])
    axs[0, 0].plot(P0[:, 0] - P2[:, 0])
    axs[0, 0].set_title('x-axis')
    #  axs[0,0].set_ylabel('mm')
    #  axs[0,0].set_xlabel('measurement #')

    axs[1, 0].plot(P0[:, 1] - P1[:-1, 1])
    axs[1, 0].plot(P0[:, 0] - P2[:, 0])
    axs[1, 0].set_title('y-axis')
    # axs[1,0].set_ylabel('mm')
    # axs[1,0].set_xlabel('measurement #')
    #    axs[1,0].set_ylim([-0.25,0.25])

    axs[0, 1].plot(P0[:, 2] - P1[:-1, 2])
    axs[0, 1].plot(P0[:, 2] - P2[:, 2])
    axs[0, 1].set_title('z-axis')
    # axs[0,1].set_ylabel('mm')
    # axs[0,1].set_xlabel('measurement #')
    #  axs[0,1].set_ylim([-0.25,0.25])

    axs[1, 1].plot(Coumpound_dif_1[:], label='Frenet_serret')
    axs[1, 1].plot(Coumpound_dif_2[:], label='Transformation Matrix')
    axs[1, 1].set_title('Compound')
    # axs[1,1].set_ylabel('mm')
    # axs[1,1].set_xlabel('measurement #')

    fig.legend(loc=4)
    fig.tight_layout()