import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from pathlib import Path

# -------------------------------------------------
# files produced by your C code
# -------------------------------------------------

variables = {
    "alpha": "alpha.dat",
    "Abar":  "Abar.dat",
    "Bbar":  "Bbar.dat",
    "KA":    "KA.dat",
    "KB":    "KB.dat"
}

output_dir = Path("animations")
output_dir.mkdir(exist_ok=True)

# -------------------------------------------------
# loader
# -------------------------------------------------

def load_data(filename):

    data = np.loadtxt(filename)

    t = data[:,0]
    r = data[:,1]
    v = data[:,2]

    times = np.unique(t)
    Nr = np.sum(t == times[0])

    r = r[:Nr]

    v = v.reshape(len(times), Nr)

    return times, r, v

# -------------------------------------------------
# animation
# -------------------------------------------------

def animate_variable(name, filename):

    print(f"Loading {filename}")

    t, r, v = load_data(filename)

    fig, ax = plt.subplots()

    line, = ax.plot([], [])

    ax.set_xlim(r.min(), r.max())
    ax.set_ylim(np.min(v), np.max(v))

    ax.set_xlabel("r")
    ax.set_ylabel(name)

    title = ax.set_title("")

    def update(frame):

        line.set_data(r, v[frame])
        title.set_text(f"{name}    t = {t[frame]:.3f}")

        return line,

    ani = animation.FuncAnimation(
        fig,
        update,
        frames=len(t),
        interval=30
    )

    outfile = output_dir / f"{name}.mp4"

    print("Saving", outfile)

    ani.save(outfile, dpi=200)

    plt.close()

# -------------------------------------------------
# run
# -------------------------------------------------

for name, file in variables.items():

    if Path(file).exists():
        animate_variable(name, file)