import numpy as np

class Simulation:
    def __init__(self):
        self.node_ids = {}
        self.next_node_id = 0

        self.branches = []
        self.shunt_caps = []

        self.branch_count = 0
        self.cstate_count = 0

        self.dirichlet = []

    def add_node(self, label:str="") -> int:
        if label in self.node_ids:
            raise ValueError(f"Node {label} already exists")
        idx = self.next_node_id
        self.node_ids[label] = idx
        self.next_node_id += 1
        return idx
    
    def _get_node_idx(self, label: str) -> int:
        if label not in self.node_ids:
            raise KeyError(f"Unknown node '{label}'")
        return self.node_ids[label]


    def add_branch(self, node1:str, node2:str, R:float=0.0, L:float=0.0, C:float=0.0, label:str="") -> int:
        branch_idx = self.branch_count
        self.branch_count += 1

        if label == "":
            label = f"b_{branch_idx}"
            
        a, b = self._get_node_idx(node1), self._get_node_idx(node2)

        c_idx = None
        if C and C > 0.0:
            c_idx = self.cstate_count
            self.cstate_count += 1
        self.branches.append((a, b, R, L, C, branch_idx, c_idx, label))
        
        return branch_idx, c_idx
    
    def add_shunt_C(self, node1: str, node2: str, C: float):
        a, b = self._get_node_idx(node1), self._get_node_idx(node2)
        self.shunt_caps.append((a, b, float(C)))

    def build_matrices(self):
        Nv = self.next_node_id
        Ni = self.branch_count
        Nc = self.cstate_count
        N = Nv + Ni + Nc

        E = np.zeros((N,N), dtype=float)
        A = np.zeros((N,N), dtype=float)

        for (a, b, R, L, C, branch_id, cap_id, _) in self.branches:
            i = Nv + branch_id      # Index of this branch's current

            # KCL for the nodes that have this current. Current leaves a and enters b
            A[a, i] += 1.0
            A[b, i] -= 1.0

            # KVL for branch: v(a)-v(b) - R*i - L di/dt - v_c = 0
            A[i, a] += 1.0
            A[i, b] -= 1.0
            if R: A[i, i] += -R
            if L: E[i, i] += L
            if C and cap_id is not None:
                vc = Nv + Ni + cap_id
                A[i, vc] += 1.0        # subtract v_c in KVL
                E[vc, vc] += C         # C * d v_c/dt = i
                A[vc, i] -= 1.0

        # Shunt capacitors into E(v,v) block
        for (a, b, C) in self.shunt_caps:
            if a != b:
                E[a, a] -= C;  E[a, b] += C
                E[b, a] += C;  E[b, b] -= C

        return E, A, N
        

    def step_BE(self, E, A, x_n, h, dirichlet_now):
        # dirichlet_now: dict label->value at t_{n+1}
        Nv = self.next_node_id
        Ni = self.branch_count
        Nc = self.cstate_count
        N  = Nv + Ni + Nc

        # fixed voltage indices and values
        fixed_v_idx = np.array([self._get_node_idx(lbl) for lbl in dirichlet_now.keys()], dtype=int)
        v_d = np.array([float(dirichlet_now[lbl]) for lbl in dirichlet_now.keys()], dtype=float)

        # free sets
        all_v_idx = np.arange(Nv, dtype=int)
        free_v_idx = np.array(sorted(set(all_v_idx) - set(fixed_v_idx)), dtype=int)
        free_i_idx = np.arange(Nv, Nv+Ni+Nc, dtype=int)
        free_idx   = np.concatenate([free_v_idx, free_i_idx])

        # Partitioned matrices
        Ef_f = E[np.ix_(free_idx, free_idx)]
        Af_f = A[np.ix_(free_idx, free_idx)]
        Ef_d = E[np.ix_(free_idx, fixed_v_idx)]
        Af_d = A[np.ix_(free_idx, fixed_v_idx)]

        # Left matrix and RHS for BE
        LHS = Ef_f - h * Af_f
        # print(np.linalg.cond(LHS))
        rhs = (E[np.ix_(free_idx, np.arange(N))] @ x_n) - (Ef_d - h * Af_d) @ v_d

        x_free_next = np.linalg.solve(LHS, rhs)
        x_next = np.zeros_like(x_n)
        x_next[free_idx]   = x_free_next
        x_next[fixed_v_idx]= v_d
        return x_next

def build_trace_with_split(sim, N=20, Rseg=0.05, Lseg=8e-9, Csh0=1e-12, Cshi=1e-14,
                           Rsrc=10.0, Rs_gnd=120.0, Rload=1e5,
                           Rsplit=0.3, Lsplit=50e-9, Cgap=100e-12,Rsplit_parallel=0.1):
    # nodes: Vs, V0..V{N}, GNDL, GNDR
    sim.add_node("Vs")
    for k in range(N+1):
        sim.add_node(f"V{k}")
    sim.add_node("GNDL"); sim.add_node("GNDR")

    # source
    sim.add_branch("Vs","V0", R=Rsrc)
    sim.add_branch("Vs","GNDL", R=Rs_gnd)

    # series ladder
    for k in range(N):
        Lk = Lseg if (k==0 or k==N-1) else (Lseg*(N/(N-1)))  # tweak if you want end-sections different
        sim.add_branch(f"V{k}", f"V{k+1}", R=Rseg, L=Lk)

    # shunt caps to local grounds
    sim.add_shunt_C("V0","GNDL", Csh0)
    for k in range(1, N):
        g = "GNDL" if k <= N//2 else "GNDR"
        sim.add_shunt_C(f"V{k}", g, Cshi)
    sim.add_shunt_C(f"V{N}", "GNDR", Csh0)

    # load
    sim.add_branch(f"V{N}", "GNDR", R=Rload)

    # ground tie: lossy via in parallel with gap capacitance
    sim.add_branch("GNDL","GNDR", R=Rsplit, L=Lsplit)    # via
    sim.add_shunt_C("GNDL","GNDR", Cgap)                 # gap
    
    sim.add_branch("GNDL","GNDR", R=Rsplit_parallel) # Parallel across the gap

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]

    # local-ground spatial profile
    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) ---
    # pick indices: 0, ~N/3, ~2N/3, N
    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):
        # local ground reference by side
        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)