"""Circuit-faithful 4x4 physical learning network used by Dillavou et al.

The nonlinear conductance, periodic topology, and local voltage-square update
follow Eqs. (2)--(3) of arXiv:2505.22887v2.  Source and target node locations
match the released Figure-5 experiment objects.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np


Array = np.ndarray


@dataclass(frozen=True)
class GridCircuit:
    rows: int = 4
    columns: int = 4
    conductance_scale: float = 8.0e-4
    threshold_voltage: float = 0.7
    measured_learning_rate: float = 2.5e3
    low_voltage: float = 0.0181
    high_voltage: float = 0.4351
    gate_minimum: float = 1.0
    gate_maximum: float = 5.2
    source_nodes: tuple[int, ...] = (15, 13, 7, 5)
    target_nodes: tuple[int, int] = (10, 8)

    @property
    def node_count(self) -> int:
        return self.rows * self.columns

    @property
    def edge_pairs(self) -> tuple[tuple[int, int], ...]:
        horizontal = []
        vertical = []
        for row in range(self.rows):
            for column in range(self.columns):
                node = row * self.columns + column
                horizontal.append((
                    node,
                    row * self.columns + (column + 1) % self.columns,
                ))
                vertical.append((
                    node,
                    ((row + 1) % self.rows) * self.columns + column,
                ))
        return tuple(horizontal + vertical)

    @property
    def edge_count(self) -> int:
        return len(self.edge_pairs)

    def source_values(self, input_one: float, input_two: float) -> Array:
        return np.asarray((
            input_one,
            input_two,
            self.low_voltage,
            self.high_voltage,
        ), dtype=float)


def edge_voltage_drops(circuit: GridCircuit, node_voltages: Array) -> Array:
    voltages = np.asarray(node_voltages, dtype=float)
    if voltages.shape != (circuit.node_count,):
        raise ValueError("node voltage vector has the wrong shape")
    return np.asarray([
        voltages[first] - voltages[second]
        for first, second in circuit.edge_pairs
    ])


def output_difference(circuit: GridCircuit, node_voltages: Array) -> float:
    positive, negative = circuit.target_nodes
    return float(node_voltages[positive] - node_voltages[negative])


def _residual_and_jacobian(
    circuit: GridCircuit, gates: Array, voltages: Array
) -> tuple[Array, Array]:
    residual = np.zeros(circuit.node_count, dtype=float)
    jacobian = np.zeros(
        (circuit.node_count, circuit.node_count), dtype=float)
    scale = circuit.conductance_scale
    threshold = circuit.threshold_voltage
    for gate, (first, second) in zip(gates, circuit.edge_pairs):
        voltage_first = voltages[first]
        voltage_second = voltages[second]
        conductance = scale * (
            gate - threshold - 0.5 * (voltage_first + voltage_second))
        current = conductance * (voltage_first - voltage_second)
        residual[first] += current
        residual[second] -= current
        derivative_first = scale * (gate - threshold - voltage_first)
        derivative_second = scale * (-gate + threshold + voltage_second)
        jacobian[first, first] += derivative_first
        jacobian[first, second] += derivative_second
        jacobian[second, first] -= derivative_first
        jacobian[second, second] -= derivative_second
    return residual, jacobian


def solve_grid_state(
    circuit: GridCircuit,
    gates: Array,
    source_values: Array,
    *,
    target_values: Array | None = None,
    initial_state: Array | None = None,
    tolerance: float = 1e-11,
    maximum_iterations: int = 20,
) -> Array:
    """Solve Kirchhoff's laws by Newton iteration with an analytic Jacobian."""
    gates = np.asarray(gates, dtype=float)
    sources = np.asarray(source_values, dtype=float)
    if gates.shape != (circuit.edge_count,):
        raise ValueError("gate vector has the wrong shape")
    if sources.shape != (len(circuit.source_nodes),):
        raise ValueError("source voltage vector has the wrong shape")
    fixed = dict(zip(circuit.source_nodes, sources))
    if target_values is not None:
        targets = np.asarray(target_values, dtype=float)
        if targets.shape != (2,):
            raise ValueError("target voltage vector must have shape (2,)")
        fixed.update(zip(circuit.target_nodes, targets))
    unknown = np.asarray([
        node for node in range(circuit.node_count) if node not in fixed
    ])
    voltages = np.full(
        circuit.node_count, float(np.mean(sources)), dtype=float)
    if initial_state is not None:
        initial = np.asarray(initial_state, dtype=float)
        if initial.shape != (circuit.node_count,):
            raise ValueError("initial state vector has the wrong shape")
        voltages[:] = initial
    for node, value in fixed.items():
        voltages[node] = value

    for _ in range(maximum_iterations):
        residual, jacobian = _residual_and_jacobian(
            circuit, gates, voltages)
        unknown_residual = residual[unknown]
        if np.linalg.norm(unknown_residual, ord=np.inf) <= tolerance:
            return voltages
        unknown_jacobian = jacobian[np.ix_(unknown, unknown)]
        step = np.linalg.solve(unknown_jacobian, unknown_residual)
        voltages[unknown] -= step
    residual, _ = _residual_and_jacobian(circuit, gates, voltages)
    raise RuntimeError(
        "grid state did not converge; residual="
        f"{np.linalg.norm(residual[unknown], ord=np.inf):.3e}")


@dataclass(frozen=True)
class GridSquareLawImperfection:
    free_gain: Array
    clamped_gain: Array
    free_input_offset_v: Array
    clamped_input_offset_v: Array
    multiplier_output_offset_v_per_s: Array

    def __post_init__(self) -> None:
        shapes = {
            np.asarray(value).shape
            for value in (
                self.free_gain,
                self.clamped_gain,
                self.free_input_offset_v,
                self.clamped_input_offset_v,
                self.multiplier_output_offset_v_per_s,
            )
        }
        if len(shapes) != 1:
            raise ValueError("grid imperfection arrays disagree")
        shape = next(iter(shapes))
        if len(shape) != 1 or shape[0] < 1:
            raise ValueError("grid imperfection arrays must be nonempty vectors")

    @classmethod
    def ideal(cls, edge_count: int) -> "GridSquareLawImperfection":
        return cls(
            free_gain=np.ones(edge_count),
            clamped_gain=np.ones(edge_count),
            free_input_offset_v=np.zeros(edge_count),
            clamped_input_offset_v=np.zeros(edge_count),
            multiplier_output_offset_v_per_s=np.zeros(edge_count),
        )

    @classmethod
    def sample_appendix_c(
        cls,
        edge_count: int,
        seed: int,
        *,
        gain_standard_deviation: float = 0.01,
        twin_mismatch_standard_deviation_v: float = 0.001,
        multiplier_offset_standard_deviation_v_per_s: float = 2.3,
    ) -> "GridSquareLawImperfection":
        rng = np.random.default_rng(seed)
        common_gain = 1.0 + rng.normal(
            0.0, gain_standard_deviation, edge_count)
        differential_gain = rng.normal(
            0.0, gain_standard_deviation, edge_count)
        common_offset = rng.normal(
            0.0, twin_mismatch_standard_deviation_v, edge_count)
        differential_offset = rng.normal(
            0.0, twin_mismatch_standard_deviation_v, edge_count)
        return cls(
            free_gain=common_gain + 0.5 * differential_gain,
            clamped_gain=common_gain - 0.5 * differential_gain,
            free_input_offset_v=common_offset + 0.5 * differential_offset,
            clamped_input_offset_v=common_offset - 0.5 * differential_offset,
            multiplier_output_offset_v_per_s=rng.normal(
                0.0, multiplier_offset_standard_deviation_v_per_s,
                edge_count),
        )

    def observed_rate(
        self,
        learning_rate: float,
        free_drops: Array,
        clamped_drops: Array,
    ) -> Array:
        measured_free = (
            self.free_gain * free_drops + self.free_input_offset_v)
        measured_clamped = (
            self.clamped_gain * clamped_drops + self.clamped_input_offset_v)
        return (
            learning_rate
            * (np.square(measured_free) - np.square(measured_clamped))
            + self.multiplier_output_offset_v_per_s
        )

    @staticmethod
    def ideal_rate(
        learning_rate: float, free_drops: Array, clamped_drops: Array
    ) -> Array:
        return learning_rate * (
            np.square(free_drops) - np.square(clamped_drops))

    def neutral_bias(self, learning_rate: float, free_drops: Array) -> Array:
        return self.observed_rate(learning_rate, free_drops, free_drops)


@dataclass
class EdgePolynomialPredictor:
    feature_center: Array
    feature_scale: Array
    coefficients: Array

    @classmethod
    def zeros(
        cls, feature_center: Array, feature_scale: Array, *, degree: int
    ) -> "EdgePolynomialPredictor":
        center = np.asarray(feature_center, dtype=float)
        scale = np.asarray(feature_scale, dtype=float)
        if center.ndim != 1 or scale.shape != center.shape:
            raise ValueError("edge feature metadata disagree")
        if degree < 0 or np.any(scale <= 0.0):
            raise ValueError("invalid polynomial degree or feature scale")
        return cls(
            feature_center=center.copy(),
            feature_scale=scale.copy(),
            coefficients=np.zeros((len(center), degree + 1), dtype=float),
        )

    @property
    def degree(self) -> int:
        return int(self.coefficients.shape[1] - 1)

    def features(self, local_state: Array) -> Array:
        state = np.asarray(local_state, dtype=float)
        if state.shape != self.feature_center.shape:
            raise ValueError("edge local state has the wrong shape")
        normalized = (state - self.feature_center) / self.feature_scale
        return np.stack([
            normalized ** power for power in range(self.degree + 1)
        ], axis=1)

    def predict(self, local_state: Array) -> Array:
        return np.sum(self.coefficients * self.features(local_state), axis=1)

    def copy(self) -> "EdgePolynomialPredictor":
        return EdgePolynomialPredictor(
            feature_center=self.feature_center.copy(),
            feature_scale=self.feature_scale.copy(),
            coefficients=self.coefficients.copy(),
        )


def fit_edge_predictor(
    predictor: EdgePolynomialPredictor,
    local_states: Array,
    neutral_measurements: Array,
    *,
    ridge: float = 1e-12,
) -> int:
    states = np.asarray(local_states, dtype=float)
    measurements = np.asarray(neutral_measurements, dtype=float)
    if states.ndim != 2 or measurements.shape != states.shape:
        raise ValueError("edge calibration matrices disagree")
    if states.shape[1] != len(predictor.feature_center):
        raise ValueError("edge calibration width changed")
    features = np.asarray([predictor.features(state) for state in states])
    for edge in range(states.shape[1]):
        design = features[:, edge, :]
        gram = design.T @ design
        rhs = design.T @ measurements[:, edge]
        predictor.coefficients[edge] = np.linalg.solve(
            gram + ridge * np.eye(gram.shape[0]), rhs)
    return int(len(states))