import hashlib
import uuid
import schemdraw
import schemdraw.elements as elm
import numpy as np
from dataclasses import dataclass, field
from typing import Dict, List, Tuple, Optional
from abc import ABC

@dataclass
class ComponentSpec:
    """Immutable specification for a circuit component."""
    kind: str  # 'R', 'V', 'LED'
    value: float  # resistance in ohms, voltage in volts
    nodes: Tuple[int, int]  # node pair (preserve order for voltage sources)
    id: str = field(default_factory=lambda: uuid.uuid4().hex)
    
    def __post_init__(self):
        # Canonicalize nodes only for resistors and LEDs (not voltage sources)
        if self.kind in ('R', 'LED') and self.nodes[0] > self.nodes[1]:
            self.nodes = (self.nodes[1], self.nodes[0])
        
        # Validate
        if self.kind in ('R', 'LED') and self.value <= 0:
            raise ValueError(f"Resistance must be positive, got {self.value}")
        if self.kind == 'V' and self.value == 0:
            raise ValueError("Zero voltage sources are not supported")

@dataclass
class ComponentState:
    """Mutable simulation state for a component."""
    current: Optional[float] = None
    voltage_drop: Optional[float] = None
    
    def reset(self):
        """Reset simulation state."""
        self.current = None
        self.voltage_drop = None

class Component(ABC):
    """
    Base class for all circuit components.

    Separates immutable specification from mutable simulation state.
    """
    
    def __init__(self, spec: ComponentSpec, state: Optional[ComponentState] = None):
        self.spec = spec
        self.state = state or ComponentState()
    
    @property
    def connections(self) -> List[int]:
        """Backward compatibility: return connections as list."""
        return list(self.spec.nodes)
    
    @property
    def id(self) -> str:
        """Component ID."""
        return self.spec.id
    
    def reset_state(self):
        """Reset simulation state."""
        self.state.reset()

    def __repr__(self) -> str:
        return f"{type(self).__name__}(ID: {self.spec.id}, Nodes: {self.spec.nodes}, Value: {self.spec.value})"

class Resistor(Component):
    """Resistor component."""
    
    def __init__(self, resistance: float, connections: List[int], id: Optional[str] = None):
        spec = ComponentSpec('R', resistance, tuple(connections), id or uuid.uuid4().hex)
        super().__init__(spec)
    
    @property
    def resistance(self) -> float:
        return self.spec.value
    
    def calculate_voltage_drop(self) -> float:
        """Calculate voltage drop using Ohm's Law: V = IR."""
        if self.state.current is None:
            raise ValueError("Current not yet calculated for resistor")
        self.state.voltage_drop = self.state.current * self.resistance
        return self.state.voltage_drop
    
    def set_current(self, current: float):
        """Set the current flowing through the resistor."""
        self.state.current = current

    def calculate_power(self) -> float:
        """Calculate power dissipated: P = V * I."""
        if self.state.voltage_drop is None or self.state.current is None:
            raise ValueError("Voltage or current is not set for power calculation")
        return self.state.voltage_drop * self.state.current

class LED(Component):
    """LED component with forward voltage model."""
    
    def __init__(self, resistance: float, connections: List[int], forward_voltage: float = 2.0, id: Optional[str] = None):
        spec = ComponentSpec('LED', resistance, tuple(connections), id or uuid.uuid4().hex)
        super().__init__(spec)
        self.forward_voltage = forward_voltage
    
    @property
    def resistance(self) -> float:
        return self.spec.value
    
    def calculate_voltage_drop(self) -> float:
        """Calculate voltage drop including forward voltage."""
        if self.state.current is None:
            raise ValueError("Current not yet calculated for LED")
        # For now, treat as resistor + forward voltage
        # In a more sophisticated model, this would be non-linear
        if self.state.current > 0:
            self.state.voltage_drop = self.state.current * self.resistance + self.forward_voltage
        else:
            self.state.voltage_drop = 0.0  # LED is off
        return self.state.voltage_drop
    
    def get_effective_resistance(self) -> float:
        """Get effective resistance for MNA solver."""
        # For LED, we need to account for forward voltage
        # This is a simplified model - in reality, LED is non-linear
        return self.resistance
    
    def set_current(self, current: float):
        """Set the current flowing through the LED."""
        self.state.current = current

    def calculate_power(self) -> float:
        """Calculate power dissipated: P = V * I."""
        if self.state.voltage_drop is None or self.state.current is None:
            raise ValueError("Voltage or current is not set for power calculation")
        return self.state.voltage_drop * self.state.current

class PowerSource(Component):
    """Power source component."""
    
    def __init__(self, voltage: float, connections: List[int], id: Optional[str] = None):
        spec = ComponentSpec('V', voltage, tuple(connections), id or uuid.uuid4().hex)
        super().__init__(spec)
    
    @property
    def voltage(self) -> float:
        return self.spec.value
    
    def calculate_voltage_drop(self) -> float:
        """Power sources provide constant voltage."""
        return self.voltage

class Capacitor(Component):
    """Capacitor with simple stored voltage state for transient simulation."""
    
    def __init__(self, capacitance: float, connections: List[int], initial_voltage: float = 0.0, id: Optional[str] = None):
        spec = ComponentSpec('C', capacitance, tuple(connections), id or uuid.uuid4().hex)
        super().__init__(spec)
        # voltage across capacitor at previous timestep (V_n-1)
        self.voltage_prev = float(initial_voltage)
    
    @property
    def capacitance(self) -> float:
        return self.spec.value
    
    def update_prev_voltage(self, v: float):
        """Update the previous voltage for the next time step."""
        self.voltage_prev = float(v)
    
    def get_prev_voltage(self) -> float:
        """Get the previous voltage across the capacitor."""
        return self.voltage_prev

def solve_mna(components: List[Component], ground_node: int = 0) -> Tuple[Dict[int, float], Dict[str, float]]:
    """
    Solve circuit using Modified Nodal Analysis (MNA).
    
    Args:
        components: List of Component objects
        ground_node: Node index to use as ground (default 0)
    
    Returns:
        node_voltage: Dict mapping node -> voltage
        comp_currents: Dict mapping component.id -> current
    
    Raises:
        RuntimeError: If circuit matrix is singular or no power source found
    """
    # Validate circuit
    power_sources = [c for c in components if isinstance(c, PowerSource)]
    if not power_sources:
        raise RuntimeError("No power source found in circuit")
    if len(power_sources) > 1:
        raise RuntimeError("Multiple power sources not yet supported")
    
    # Collect all nodes
    nodes = set()
    v_sources = []
    for c in components:
        nodes.update(c.spec.nodes)
        if isinstance(c, PowerSource):
            v_sources.append(c)
    
    # Map nodes to matrix indices (excluding ground)
    node_list = sorted(n for n in nodes if n != ground_node)
    n = len(node_list)
    m = len(v_sources)
    node_index = {node_list[i]: i for i in range(n)}
    
    # System matrix: [V_nodes (exclude ground)] [I_vsources]
    dim = n + m
    A = np.zeros((dim, dim), dtype=float)
    z = np.zeros((dim,), dtype=float)
    
    # Stamp resistors and LEDs
    for c in components:
        if isinstance(c, (Resistor, LED)):
            n1, n2 = c.spec.nodes
            g = 1.0 / c.spec.value  # conductance
            
            if n1 != ground_node and n2 != ground_node:
                i = node_index[n1]
                j = node_index[n2]
                A[i, i] += g
                A[j, j] += g
                A[i, j] -= g
                A[j, i] -= g
            elif n1 != ground_node:
                i = node_index[n1]
                A[i, i] += g
            elif n2 != ground_node:
                j = node_index[n2]
                A[j, j] += g
    
    # Note: Capacitors are not stamped in DC analysis (open circuits)
    
    # For LEDs, we need to add forward voltage to the RHS
    # This is a simplified approach - in reality, LED is non-linear
    for c in components:
        if isinstance(c, LED):
            n1, n2 = c.spec.nodes
            if n1 != ground_node:
                i = node_index[n1]
                z[i] += c.forward_voltage / c.spec.value
            if n2 != ground_node:
                j = node_index[n2]
                z[j] -= c.forward_voltage / c.spec.value
    
    # Stamp voltage sources
    for k, vs in enumerate(v_sources):
        n1, n2 = vs.spec.nodes
        
        # Connect to node equations
        if n1 != ground_node:
            i = node_index[n1]
            A[i, n + k] += 1.0
            A[n + k, i] += 1.0
        if n2 != ground_node:
            j = node_index[n2]
            A[j, n + k] -= 1.0
            A[n + k, j] -= 1.0
        
        # RHS = source voltage (V(n1) - V(n2) = Vs)
        # For voltage source from n1 to n2, V(n1) - V(n2) = Vs
        z[n + k] = vs.spec.value
    
    # Solve system
    try:
        x = np.linalg.solve(A, z)
    except np.linalg.LinAlgError as e:
        raise RuntimeError("Circuit matrix is singular; check connectivity and floating nodes") from e
    
    # Extract node voltages (including ground = 0)
    node_voltage = {ground_node: 0.0}
    for node, idx in node_index.items():
        node_voltage[node] = float(x[idx])
    
    # Calculate component currents and voltage drops
    comp_currents = {}
    for c in components:
        n1, n2 = c.spec.nodes
        v1 = node_voltage.get(n1, 0.0)
        v2 = node_voltage.get(n2, 0.0)
        
        if isinstance(c, (Resistor, LED)):
            i_comp = (v1 - v2) / c.spec.value
            comp_currents[c.spec.id] = float(i_comp)
        elif isinstance(c, Capacitor):
            # Capacitors have zero current in DC analysis (open circuit)
            comp_currents[c.spec.id] = 0.0
        elif isinstance(c, PowerSource):
            # Current through source is stored in x[n + k]
            k = v_sources.index(c)
            comp_currents[c.spec.id] = float(x[n + k])
    
    return node_voltage, comp_currents

def solve_mna_time_step(components: List[Component], prev_node_voltage: Dict[int, float], dt: float, ground_node: int = 0,
                       integration_method: str = "backward_euler", prev_capacitor_currents: Optional[Dict[str, float]] = None
                       ) -> Tuple[Dict[int, float], Dict[str, float]]:
    """
    Solve MNA for one time step using companion model for capacitors.

    Integration methods:
    - "backward_euler": First-order accurate, unconditionally stable
    - "trapezoidal": Second-order accurate, stable for most circuits

    Capacitor between n1,n2:
      Backward-Euler: G = C/dt, I_eq = G * (V_prev_n1 - V_prev_n2)
      Trapezoidal: G = 2C/dt, I_eq = G * (V_prev_n1 - V_prev_n2) + I_prev

    Args:
        components: list of Component instances
        prev_node_voltage: mapping node -> voltage at previous timestep (t - dt)
        dt: timestep in seconds
        ground_node: node index used as ground
        integration_method: "backward_euler" or "trapezoidal"
        prev_capacitor_currents: dict capacitor_id -> current at previous timestep (required for trapezoidal)

        Returns:
        node_voltage: dict node -> voltage at current timestep
        comp_currents: dict component_id -> current (positive from node1 -> node2)
    """
    if dt <= 0:
        raise ValueError("dt must be positive")

    # collect nodes and voltage sources
    nodes = set()
    v_sources = []
    for c in components:
        nodes.update(c.spec.nodes)
        if isinstance(c, PowerSource):
            v_sources.append(c)

    node_list = sorted(n for n in nodes if n != ground_node)
    n = len(node_list)
    m = len(v_sources)
    node_index = {node_list[i]: i for i in range(n)}

    dim = n + m
    A = np.zeros((dim, dim), dtype=float)
    z = np.zeros((dim,), dtype=float)

    # stamp resistors and LEDs as before
    for c in components:
        if isinstance(c, (Resistor, LED)):
            n1, n2 = c.spec.nodes
            g = 1.0 / c.spec.value
            if n1 != ground_node:
                i = node_index[n1]
                A[i, i] += g
            if n2 != ground_node:
                j = node_index[n2]
                A[j, j] += g
            if n1 != ground_node and n2 != ground_node:
                i = node_index[n1]
                j = node_index[n2]
                A[i, j] -= g
                A[j, i] -= g

    # stamp capacitors using companion model
    for c in components:
        if isinstance(c, Capacitor):
            n1, n2 = c.spec.nodes
            C = c.capacitance
            
            # Choose integration method
            if integration_method == "backward_euler":
                G = C / dt  # conductance in BE companion model
                # compute RHS injection from previous capacitor voltage
                Vprev_n1 = prev_node_voltage.get(n1, 0.0)
                Vprev_n2 = prev_node_voltage.get(n2, 0.0)
                Ieq = G * (Vprev_n1 - Vprev_n2)  # inject at n1, remove at n2
            elif integration_method == "trapezoidal":
                G = 2.0 * C / dt  # conductance in trapezoidal companion model
                # compute RHS injection from previous capacitor voltage and current
                Vprev_n1 = prev_node_voltage.get(n1, 0.0)
                Vprev_n2 = prev_node_voltage.get(n2, 0.0)
                I_prev = prev_capacitor_currents.get(c.spec.id, 0.0) if prev_capacitor_currents else 0.0
                Ieq = G * (Vprev_n1 - Vprev_n2) + I_prev  # inject at n1, remove at n2
            else:
                raise ValueError(f"Unknown integration method: {integration_method}")
            
            # stamp conductance like a resistor between n1 and n2
            if n1 != ground_node:
                i = node_index[n1]
                A[i, i] += G
            if n2 != ground_node:
                j = node_index[n2]
                A[j, j] += G
            if n1 != ground_node and n2 != ground_node:
                i = node_index[n1]
                j = node_index[n2]
                A[i, j] -= G
                A[j, i] -= G
            
            # inject equivalent current
            if n1 != ground_node:
                i = node_index[n1]
                z[i] += Ieq
            if n2 != ground_node:
                j = node_index[n2]
                z[j] -= Ieq

    # stamp LED forward-voltage as earlier (simple linear approx)
    for c in components:
        if isinstance(c, LED):
            n1, n2 = c.spec.nodes
            # add forward-voltage equivalent injection: +Vf / R to n1 and - to n2
            Vf = c.forward_voltage
            g = 1.0 / c.spec.value
            if n1 != ground_node:
                i = node_index[n1]
                z[i] += Vf * g
            if n2 != ground_node:
                j = node_index[n2]
                z[j] -= Vf * g

    # stamp voltage sources (MNA)
    for k, vs in enumerate(v_sources):
        n1, n2 = vs.spec.nodes
        if n1 != ground_node:
            i = node_index[n1]
            A[i, n + k] += 1.0
            A[n + k, i] += 1.0
        if n2 != ground_node:
            j = node_index[n2]
            A[j, n + k] -= 1.0
            A[n + k, j] -= 1.0
        z[n + k] = vs.spec.value  # V(n1)-V(n2)=Vs

    # solve
    try:
        x = np.linalg.solve(A, z)
    except np.linalg.LinAlgError as e:
        raise RuntimeError("Circuit matrix is singular at this time step; check connectivity/floating nodes") from e

    node_voltage = {ground_node: 0.0}
    for node, idx in node_index.items():
        node_voltage[node] = float(x[idx])

    # compute component currents (positive: node1 -> node2)
    comp_currents = {}
    for c in components:
        n1, n2 = c.spec.nodes
        v1 = node_voltage.get(n1, 0.0)
        v2 = node_voltage.get(n2, 0.0)
        if isinstance(c, (Resistor, LED)):
            i_comp = (v1 - v2) / c.spec.value
            comp_currents[c.spec.id] = float(i_comp)
        elif isinstance(c, Capacitor):
            # exact current calculation depends on integration method
            Vprev = c.get_prev_voltage()
            if integration_method == "backward_euler":
                # BE: I = C*(Vt - Vprev)/dt
                i_comp = c.capacitance * ((v1 - v2) - Vprev) / dt
            elif integration_method == "trapezoidal":
                # Trapezoidal: I = 2C*(Vt - Vprev)/dt - I_prev
                I_prev = prev_capacitor_currents.get(c.spec.id, 0.0) if prev_capacitor_currents else 0.0
                i_comp = 2.0 * c.capacitance * ((v1 - v2) - Vprev) / dt - I_prev
            else:
                raise ValueError(f"Unknown integration method: {integration_method}")
            comp_currents[c.spec.id] = float(i_comp)
        elif isinstance(c, PowerSource):
            k = v_sources.index(c)
            comp_currents[c.spec.id] = float(x[n + k])

    return node_voltage, comp_currents

class Circuit:
    """
    Circuit simulator using Modified Nodal Analysis (MNA).
    
    Supports arbitrary topologies with resistors, LEDs, and voltage sources.
    """

    def __init__(self):
        """Initialize an empty circuit."""
        self.components: List[Component] = []
        self._node_voltage_cache: Optional[Dict[int, float]] = None
        self._comp_currents_cache: Optional[Dict[str, float]] = None
        self._topology_hash: Optional[str] = None
    
    def add(self, component: Component):
        """Add a component to the circuit."""
        self.components.append(component)
        self._invalidate_cache()
    
    def _invalidate_cache(self):
        """Invalidate cached simulation results."""
        self._node_voltage_cache = None
        self._comp_currents_cache = None
        self._topology_hash = None
    
    def _get_topology_hash(self) -> str:
        """Get hash of current topology for caching."""
        if self._topology_hash is None:
            # Create a hash based on component types and connections
            topology_str = ""
            for comp in sorted(self.components, key=lambda c: c.spec.id):
                topology_str += f"{comp.spec.kind}:{comp.spec.nodes}:{comp.spec.value};"
            self._topology_hash = hashlib.md5(topology_str.encode()).hexdigest()
        return self._topology_hash
    
    def _validate_circuit(self):
        """Validate circuit before simulation."""
        if not self.components:
            raise RuntimeError("Circuit is empty")
        
        # Check for power source
        power_sources = [c for c in self.components if isinstance(c, PowerSource)]
        if not power_sources:
            raise RuntimeError("No power source found in circuit")
        
        # Check for isolated nodes
        all_nodes = set()
        connected_nodes = set()
        for comp in self.components:
            all_nodes.update(comp.spec.nodes)
            connected_nodes.update(comp.spec.nodes)
        
        # All nodes should be connected (this is a simple check)
        if len(connected_nodes) < 2:
            raise RuntimeError("Circuit must have at least 2 connected nodes")
    
    def simulate(self) -> None:
        """
        Simulate the circuit using MNA.
        
        Updates all component states with calculated currents and voltage drops.
        """
        self._validate_circuit()
        
        # Reset all component states
        for comp in self.components:
            comp.reset_state()
        
        # Solve circuit
        node_voltage, comp_currents = solve_mna(self.components)
        
        # Update component states
        for comp in self.components:
            comp_id = comp.spec.id
            if comp_id in comp_currents:
                comp.state.current = comp_currents[comp_id]
                if isinstance(comp, (Resistor, LED)):
                    comp.calculate_voltage_drop()
        
        # Cache results
        self._node_voltage_cache = node_voltage
        self._comp_currents_cache = comp_currents
    
    def calculate_current(self) -> float:
        """
        Calculate total current in the circuit.

        Returns:
            Total current in amperes
        """
        if not self.components:
            raise RuntimeError("Circuit is empty")
        
        # Find power source current
        power_sources = [c for c in self.components if isinstance(c, PowerSource)]
        if not power_sources:
            raise RuntimeError("No power source found in circuit")
        
        # Simulate if needed
        if self._comp_currents_cache is None:
            self.simulate()
        
        # Return current through power source
        ps = power_sources[0]
        return self._comp_currents_cache[ps.spec.id]
    
    def calculate_total_resistance(self) -> float:
        """
        Calculate equivalent resistance seen by the power source.
        
        Note: This is only meaningful for simple series/parallel circuits.
        For complex topologies, use the MNA solver.
        """
        if not self.components:
            return 0.0
        
        # Simulate to get current
        try:
            total_current = self.calculate_current()
            power_sources = [c for c in self.components if isinstance(c, PowerSource)]
            if power_sources:
                source_voltage = power_sources[0].voltage
                return source_voltage / total_current if total_current != 0 else float('inf')
        except:
            pass
        
        # Fallback: sum series resistances (only works for series circuits)
        total_resistance = sum(
            comp.spec.value for comp in self.components 
            if isinstance(comp, (Resistor, LED))
        )
        return total_resistance
    
    def measure_voltage(self, connections: List[int]) -> float:
        """
        Measure voltage drop across specified connections.
        
        Args:
            connections: List of two node IDs to measure across
            
        Returns:
            Voltage drop in volts
        """
        if len(connections) != 2:
            raise ValueError("Must specify exactly 2 nodes for voltage measurement")
        
        # Simulate if needed
        if self._node_voltage_cache is None:
            self.simulate()
        
        n1, n2 = connections
        v1 = self._node_voltage_cache.get(n1, 0.0)
        v2 = self._node_voltage_cache.get(n2, 0.0)
        return v1 - v2
    
    def measure_current(self, connections: List[int]) -> float:
        """
        Measure current through specified connections.
        
        Args:
            connections: List of two node IDs to measure current through
            
        Returns:
            Current in amperes
        """
        if len(connections) != 2:
            raise ValueError("Must specify exactly 2 nodes for current measurement")
        
        # Simulate if needed
        if self._comp_currents_cache is None:
            self.simulate()
        
        # Find component connected to these nodes
        # Prefer resistors/LEDs over power sources
        for comp in self.components:
            if set(comp.spec.nodes) == set(connections) and isinstance(comp, (Resistor, LED)):
                return self._comp_currents_cache.get(comp.spec.id, 0.0)
        
        # If no resistor/LED found, look for any component
        for comp in self.components:
            if set(comp.spec.nodes) == set(connections):
                return self._comp_currents_cache.get(comp.spec.id, 0.0)
        
        raise ValueError(f"No component found between nodes {connections}")
    
    def calculate_power(self) -> float:
        """
        Calculate total power dissipated in the circuit.
        
        Returns:
            Total power in watts
        """
        if self._comp_currents_cache is None:
            self.simulate()
        
        total_power = 0.0
        for comp in self.components:
            if isinstance(comp, (Resistor, LED)) and comp.spec.id in self._comp_currents_cache:
                current = self._comp_currents_cache[comp.spec.id]
                voltage_drop = comp.calculate_voltage_drop()
                total_power += abs(current * voltage_drop)
            elif isinstance(comp, Capacitor) and comp.spec.id in self._comp_currents_cache:
                # Capacitors don't dissipate power in DC analysis
                pass
        
        return total_power
    
    def draw_circuit_diagram(self):
        """
        Draw a circuit diagram using schemdraw.
        
        This is a simplified version - for complex circuits, you might want
        a more sophisticated layout algorithm.
        """
        d = schemdraw.Drawing()
        
        # Simple layout: components in a line
        x_pos = 0
        for comp in self.components:
            if isinstance(comp, PowerSource):
                d += elm.SourceV().at([x_pos, 0]).label(f'{comp.voltage}V')
            elif isinstance(comp, Resistor):
                d += elm.Resistor().at([x_pos, 0]).label(f'{comp.resistance}Ω')
            elif isinstance(comp, LED):
                d += elm.LED().at([x_pos, 0]).label('LED')
            
            x_pos += 2
        
        d.draw()
    
    def get_component_by_id(self, component_id: str) -> Optional[Component]:
        """Get component by its ID."""
        for comp in self.components:
            if comp.spec.id == component_id:
                return comp
        return None
    
    def remove_component(self, component_id: str) -> bool:
        """
        Remove component by ID.

        Returns:
            True if component was found and removed, False otherwise
        """
        for i, comp in enumerate(self.components):
            if comp.spec.id == component_id:
                del self.components[i]
                self._invalidate_cache()
                return True
        return False
    
    def transient_simulate(self, t_stop: float, dt: float, record_nodes: Optional[List[int]] = None,
                          integration_method: str = "backward_euler"
                          ) -> Tuple[np.ndarray, Dict[int, np.ndarray], Dict[str, np.ndarray]]:
        """
        Run transient simulation from t=0 to t=t_stop (inclusive) with timestep dt.

        Args:
            t_stop: End time in seconds
            dt: Time step in seconds
            record_nodes: Optional list of nodes to record (if None, records all)
            integration_method: "backward_euler" or "trapezoidal"

        Returns:
          times, node_history, comp_current_history
          - times: numpy array of time points
          - node_history: dict node -> array(len(times)) of node voltages
          - comp_current_history: dict comp_id -> array(len(times)) of currents (A)
        """
        # Basic checks
        if dt <= 0 or t_stop <= 0:
            raise ValueError("t_stop and dt must be positive")
        self._validate_circuit()  # reuse earlier validator

        # collect nodes and initialize previous voltages
        nodes = sorted({n for comp in self.components for n in comp.spec.nodes})
        ground_node = 0
        prev_node_voltage = {n: 0.0 for n in nodes}
        # If any capacitor has an initial voltage, set accordingly
        for comp in self.components:
            if isinstance(comp, Capacitor):
                # Set initial voltage across capacitor
                n1, n2 = comp.spec.nodes
                Vprev = comp.get_prev_voltage()
                # Set the voltage difference between nodes
                prev_node_voltage[n1] = prev_node_voltage.get(n2, 0.0) + Vprev

        times = np.arange(0.0, t_stop + dt/2, dt)
        steps = len(times)

        # Prepare histories
        node_history = {n: np.zeros(steps, dtype=float) for n in nodes}
        comp_history = {comp.spec.id: np.zeros(steps, dtype=float) for comp in self.components}

        # For trapezoidal integration, we need to track previous capacitor currents
        prev_capacitor_currents = None
        if integration_method == "trapezoidal":
            prev_capacitor_currents = {}

        # time loop
        for k, t in enumerate(times):
            node_voltage, comp_currents = solve_mna_time_step(
                self.components, prev_node_voltage, dt, ground_node=ground_node,
                integration_method=integration_method, prev_capacitor_currents=prev_capacitor_currents
            )

            # store
            for n in nodes:
                node_history[n][k] = node_voltage.get(n, 0.0)
            for comp in self.components:
                cid = comp.spec.id
                comp_history[cid][k] = comp_currents.get(cid, 0.0)

            # update component states (currents / voltage drops) for user inspection
            for comp in self.components:
                cid = comp.spec.id
                if cid in comp_currents:
                    comp.state.current = comp_currents[cid]
                    if isinstance(comp, (Resistor, LED)):
                        comp.calculate_voltage_drop()
                    elif isinstance(comp, Capacitor):
                        # update capacitor prev voltage for next step
                        n1, n2 = comp.spec.nodes
                        Vc = node_voltage.get(n1, 0.0) - node_voltage.get(n2, 0.0)
                        comp.update_prev_voltage(Vc)
                        
                        # For trapezoidal integration, store current for next step
                        if integration_method == "trapezoidal":
                            prev_capacitor_currents[cid] = comp_currents[cid]

            # prepare prev_node_voltage for next step (use new node voltages)
            prev_node_voltage = node_voltage

        # optionally filter node_history to record_nodes, but returning all is useful
        return times, node_history, comp_history
    
    def compare_integration_methods(self, t_stop: float, dt: float, node_to_compare: int = 2) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """
        Compare backward-Euler and trapezoidal integration methods.
        
        Args:
            t_stop: End time in seconds
            dt: Time step in seconds
            node_to_compare: Node to compare voltages for

        Returns:
            times, be_voltage, trap_voltage
            - times: numpy array of time points
            - be_voltage: voltage history using backward-Euler
            - trap_voltage: voltage history using trapezoidal
        """
        # Run simulation with backward-Euler
        times_be, node_history_be, _ = self.transient_simulate(t_stop, dt, integration_method="backward_euler")
        
        # Run simulation with trapezoidal
        times_trap, node_history_trap, _ = self.transient_simulate(t_stop, dt, integration_method="trapezoidal")
        
        # Extract voltage history for the specified node
        be_voltage = node_history_be[node_to_compare]
        trap_voltage = node_history_trap[node_to_compare]
        
        return times_be, be_voltage, trap_voltage