import numpy as np

from Simulation import Simulation
from helpers import build_trace_with_split

if __name__ == "__main__":
    import matplotlib.pyplot as plt

    N = 50
    sim = Simulation()
    # build_trace_with_split(sim,
    #                        N=N,
    #                        Rseg=0.02,
    #                        Lseg=8e-9,
    #                        Csh0=1e-12,
    #                        Cshi=10e-15,
    #                        Rsrc=50.0,
    #                        Rs_gnd=120.0,
    #                        Rload=120,
    #                        Rsplit=0.5,
    #                        Lsplit=50e-9,
    #                        Cgap=50e-12,
    #                        Rsplit_parallel=0.1)
    build_trace_with_split(sim,
                           N=N,
                           Rseg=0.02,
                           Lseg=8e-9,
                           Csh0=1e-12,
                           Cshi=1e-12,
                           Rsrc=50.0,
                           Rs_gnd=120.0,
                           Rload=120,
                           Rsplit=0.0,
                           Lsplit=0.0,
                           Cgap=0.0,
                           Rsplit_parallel=0.1)

    E,A,NN = sim.build_matrices()
    dt = 1e-12              # 1 ps
    t_final = 50e-9         # 50 ns
    steps = int(t_final / dt)
    x = np.zeros(NN)
    drive = {"Vs":3.0, "GNDL":0.0}

    nodes = ["Vs"]+[f"V{k}" for k in range(N+1)]+["GNDL","GNDR"]
    idx   = [sim._get_node_idx(s) for s in nodes]
    tvec  = np.arange(steps+1)*dt
    traces = np.zeros((len(nodes), steps+1))
    traces[:,0] = x[idx]

    for n in range(steps):
        x = sim.step_BE(E, A, x, dt, drive)
        traces[:,n+1] = x[idx]

    # Node voltages referenced to their local ground
    def spatial_profile(n):
        v = np.zeros(N+1)
        for k in range(N+1):
            if k <= N//2:
                v[k] = traces[nodes.index(f"V{k}"), n] - traces[nodes.index("GNDL"), n]
            else:
                v[k] = traces[nodes.index(f"V{k}"), n] - traces[nodes.index("GNDR"), n]
        return v
    
    # Time plots for 4 evenly spaced taps (include V0 and VN vs local ground) ---
    tap_idx = sorted(set([0, N//3, (2*N)//3, N]))
    tap_labels = [f"V{k}" for k in tap_idx]

    def v_local(k, n):
        if k <= N//2:
            return traces[nodes.index(f"V{k}"), n] - traces[nodes.index("GNDL"), n]
        else:
            return traces[nodes.index(f"V{k}"), n] - traces[nodes.index("GNDR"), n]

    Vt = {k: np.array([v_local(k, n) for n in range(traces.shape[1])]) for k in tap_idx}

    plt.figure(figsize=(8,4))
    for k in tap_idx:
        plt.plot(tvec*1e9, Vt[k], label=f"{'V0' if k==0 else ('V_N' if k==N else f'V{k}')}")
    plt.xlabel("Time [ns]")
    plt.ylabel("Voltage vs local ground [V]")
    plt.title("Selected nodes over time")
    plt.xlim(0, 10)
    plt.grid(True)
    plt.legend()
    plt.tight_layout()
    plt.savefig("reflection_wave.png")
    plt.show()

    # GIF
    from PIL import Image
    import io, numpy as np, matplotlib.pyplot as plt

    xpos = np.arange(N+1)
    Nt = traces.shape[1]
    target_fps = 5
    max_frames = 300
    stride = max(1, Nt // max_frames)
    frame_idx = np.arange(0, Nt, stride)
    vals = np.vstack([spatial_profile(i) for i in frame_idx])
    ymin, ymax = 1.1*vals.min(), 1.1*vals.max()
    duration_ms = int(1000/target_fps)

    frames = []
    for i in frame_idx:
        v = spatial_profile(i)
        fig, ax = plt.subplots(figsize=(6,3.5))
        ax.plot(xpos, v, marker=None)
        ax.set_xlim(0, N)
        ax.set_ylim(ymin, ymax)
        ax.set_xlabel("Position index (0 … N)")
        ax.set_ylabel("V vs local ground [V]")
        ax.set_title(f"t = {tvec[i]*1e9:.2f} ns")
        ax.grid(True)
        buf = io.BytesIO()
        plt.tight_layout()
        fig.savefig(buf, format='png', dpi=120)
        plt.close(fig)
        buf.seek(0)
        frames.append(Image.open(buf).convert("P"))

    out_path = "reflection_wave_N24.gif"
    frames[0].save(out_path, save_all=True, append_images=frames[1:], duration=duration_ms, loop=0, optimize=False)
    print("Saved", out_path)