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

def generate_gif(x,U_t,real_time,filename="soliton.gif",
                 interval=50,U_a=[],title='', y_min=-7, y_max =  7,
                 msr = False,Umsr=[],xmsr=[]):
    if filename[-4:]!=".gif":
        filename += ".gif"

    fig, ax = plt.subplots()
    line, = ax.plot(x, U_t[0],label='numerical soultion')
    if msr:
        ax.scatter(xmsr,Umsr,label='measuring points',s=6,color='red')
    if len(U_a)!=0:
        line_a, = ax.plot(x,U_a[0],label='analytic solution',linestyle = 'dotted') 
    ax.set_ylim(y_min, y_max)
    ax.set_xlabel("x")
    ax.set_ylabel("U(x)")
    fig.suptitle(title)
    ax.legend()

    def update(frame):
        line.set_ydata(U_t[frame])
        if len(U_a)>0:
            line_a.set_ydata(U_a[frame])
        ax.set_title(f"Time {real_time[frame]:.2f}")
        return line,

    anim = FuncAnimation(fig, update, frames=U_t.shape[0]-1, interval=interval, blit=True)


    anim.save(filename, writer=PillowWriter(fps=1000//interval))
    plt.close(fig)

    print(f"GIF saved as '{filename}'")


def giff_comp(x,U_t,real_time,labels,title,filename="soliton.gif",interval=50,):
    if filename[-4:]!=".gif":
        filename += ".gif"

    y_min, y_max = -7, 7

    fig, ax = plt.subplots()
    lines = []
    for U,label in zip (U_t,labels):
        line, = ax.plot(x, U[0],label=label)
        lines.append(line)
    ax.set_ylim(y_min, y_max)
    ax.set_xlabel("x")
    ax.set_ylabel("U(x)")
    fig.suptitle(title)
    ax.legend(loc='upper right')

    def update(frame):
        for U, line in zip(U_t,lines): 
            line.set_ydata(U[frame])
        return line,

    anim = FuncAnimation(fig, update, frames=U_t[0].shape[0]-1, interval=interval, blit=True)


    anim.save(filename, writer=PillowWriter(fps=1000//interval))
    plt.close(fig)

    print(f"GIF saved as '{filename}'")
