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