import numpy as np
import gif_gen as gg
import sg_soliton_functions as ssf
import matplotlib.pyplot as plt
import two_soliton_start as two_s
import split_step_Fourier as split_f
def plot(error_vec,Nx_vec,Lx,time_vec,dt_vec,title,fname,labels):
    if len(labels)==0:
        labels = np.zeros_like(error_vec)
    fig, ax = plt.subplots()
    ax.set_xlabel("time")
    ax.set_ylabel("($U_a$ -  $U_n$)/grid size")
    for Nx, error, time,dt,lab in zip(Nx_vec,error_vec,time_vec,dt_vec,labels):
        ax.plot(time,error,label=lab)
    fig.legend()
    ax.set_title(title)
    fig.savefig(fname)
    print(f"Plot saved as \'{fname}\'")
def starting_point():
    #constant values
    T = 30
    Lx = 100
    speed = 0.1
    w = 0.1
    gamma = 1/np.sqrt(1-speed**2)

    # Grid
    Nx_vec = [512,1028,2048]
    dt_vec = [1e-2,1e-2,1e-2]
    x0=10
    
    # Time
    error_vec =[]
    time_vec =[]
    method = split_f.step_4th_order
    
    '''
    i=0
    for Nx, dt in zip(Nx_vec,dt_vec):
        x = np.linspace(-Lx/2, Lx/2, Nx, dtype=np.float64,endpoint=False)
        dx = x[1]-x[0]
        k = 2*np.pi*np.fft.fftfreq(Nx, d=dx)
        steps = int(T/dt)
        save_every = max(1, steps//90)
        # Starting point
        u0,v0,name,_,f_a,kfilter = two_s.moving_breather(x,speed,w,gamma)
        title = f'Moving breather for dx={dx:.2f} and dt=dt{dt:.2E}'
        # Integration
        
        timeline,time = split_f.integration_loop(u0,v0,steps,save_every,dt,k,kfilter,method)
        #method  =split_f.strang_split_step
        time_vec.append(time)
        # Analitic part
        U_a = np.array([f_a(x,t,w,speed,gamma) for t in time])
        delta = split_f.sum_error(U_a,timeline)
        error_vec.append(delta.copy())
        gg.generate_gif(x,timeline,time,f'images/breathers/{name[:-3]}_dx{dx}_dt{dt}_{i}.gif',U_a=U_a,title=title)
        i+=1
    fname = "images/breathers/err_mov_breather.png"
    plot(error_vec,Nx_vec,Lx,time_vec,dt_vec,title,fname,labels=['4th order','2nd order'])
    print(np.sum(error_vec[0]-error_vec[1]))
    error_vec =[]
    time_vec =[]
    '''
    Nx = 1028
    dt = 1e-2
    x = np.linspace(-Lx/2, Lx/2, Nx, dtype=np.float64,endpoint=False)
    dx = x[1]-x[0]
    k = 2*np.pi*np.fft.fftfreq(Nx, d=dx)
    v_vec = [0.1,0.5,0.5,0.8]
    w_vec = [0.1,0.5,0.7,0.5]
    steps = int(T/dt)
    save_every = max(1, steps//90)
    for speed,w in zip(v_vec,w_vec):
        gamma = 1/np.sqrt(1-speed**2)
        
        # Starting point
        u0,v0,name,_,f_a,kfilter = two_s.kink_start(x,speed,x0)
        title = f'Moving breather for v={speed:.2f} and $\omega$={w:.2E}'
        # Integration
        
        timeline,time = split_f.integration_loop(u0,v0,steps,save_every,dt,k,kfilter,method)
        #method  =split_f.strang_split_step
        time_vec.append(time)
        # Analitic part
        U_a = np.array([f_a(x,t,x0,speed) for t in time])
        delta = split_f.sum_error(U_a,timeline)
        error_vec.append(delta.copy())
        gg.generate_gif(x,timeline,time,f'images/{name[:-3]}_v{speed}_w{w}.gif',U_a=U_a,title=title)
    fname = "images/kink_err.png"
    plot(error_vec,Nx_vec,Lx,time_vec,dt_vec,title,fname,labels=[f'v{vlab} $\omega$={wlab}'for vlab,wlab in  zip(v_vec,w_vec)])
    print(np.sum(error_vec[0]-error_vec[1]))
   
    '''
    method = split_f.step_4th_order
  
    speed = 0.7

    for Nx, dt in zip(Nx_vec,dt_vec):
        x = np.linspace(-Lx/2, Lx/2, Nx, endpoint=False)
        dx = x[1]-x[0]
        k = 2*np.pi*np.fft.fftfreq(Nx, d=dx)
        steps = int(T/dt)
        save_every = max(1, steps//90)
        # Starting point
        u0,v0,name,_,f_a,kfilter = two_s.kink_antikink_start(x,speed,x0)
        title = f'Kink and antikink for dx={dx:.2f} and dt=dt{dt:.2E}'
        # Integration
        timeline,time = split_f.integration_loop(u0,v0,steps,save_every,dt,k,kfilter,method)
        time_vec.append(time)
        #method  =split_f.strang_split_step
        # Analitic part
        U_a = np.array([f_a(x,t,speed,x0) for t in time])
        delta = split_f.sum_error(U_a,timeline)
        error_vec.append(delta)
        gg.generate_gif(x,timeline,time,f"images/kink_antikink/{name[:-3]}_dx{dx}_dt{dt}.gif",U_a=U_a,title=title)
    fname = "images/kink_antikink/err_kink_antikink_dx.png"
    plot(error_vec,Nx_vec,Lx,time_vec,dt_vec,title,fname,labels=[f'dx = {Lx/2/Nx_vec[0]}',f'dx = {Lx/2/Nx_vec[1]}',f'dx = {Lx/2/Nx_vec[2]}'])
    ''' 
    
    #print(timeline[0:4])
starting_point()