"""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))


@dataclass(frozen=True)
class RingClassificationDataset:
    inputs_v: Array
    labels_v: Array

    def __post_init__(self) -> None:
        inputs = np.asarray(self.inputs_v)
        labels = np.asarray(self.labels_v)
        if inputs.ndim != 2 or inputs.shape[1] != 2:
            raise ValueError("ring inputs must have shape (samples, 2)")
        if labels.shape != (len(inputs),):
            raise ValueError("ring labels must match the sample count")
        if np.any(labels == 0.0):
            raise ValueError("classification labels must be signed")


def evaluate_grid_classifier(
    circuit: GridCircuit,
    gates: Array,
    dataset: RingClassificationDataset,
    *,
    initial_states: list[Array | None] | None = None,
) -> tuple[dict, list[Array]]:
    if initial_states is None:
        initial_states = [None] * len(dataset.labels_v)
    states = []
    outputs = []
    for index, inputs in enumerate(dataset.inputs_v):
        state = solve_grid_state(
            circuit,
            gates,
            circuit.source_values(*inputs),
            initial_state=initial_states[index],
        )
        states.append(state)
        outputs.append(output_difference(circuit, state))
    outputs_array = np.asarray(outputs)
    labels = np.asarray(dataset.labels_v)
    errors = labels - outputs_array
    active = labels * errors > 0.0
    return {
        "classification_error": float(np.mean(
            np.sign(outputs_array) != np.sign(labels))),
        "hinge_loss_v2": float(np.mean(np.where(
            active, np.square(errors), 0.0))),
        "margin_success_fraction": float(np.mean(~active)),
        "outputs_v": outputs_array.tolist(),
    }, states


def train_grid_classifier(
    circuit: GridCircuit,
    initial_gates: Array,
    dataset: RingClassificationDataset,
    imperfection: GridSquareLawImperfection,
    *,
    method: str,
    epochs: int,
    predictor: EdgePolynomialPredictor | None = None,
    standard_nudging: float = 128.0 / 129.0,
    standard_learning_time_seconds: float = 1.0e-3,
    overclamp_nudging: float = 32.0 / 129.0,
    overclamp_target_magnitude_v: float | None = None,
    overclamp_time_seconds_per_v: float = 0.05,
    record_every: int = 50,
) -> dict:
    """Train the physical grid with explicit local voltage-square updates."""
    allowed = {
        "clean",
        "raw",
        "constant",
        "sdil",
        "oracle_neutral",
        "overclamp_clean",
        "overclamp",
        "overclamp_constant",
        "overclamp_sdil",
        "overclamp_oracle_neutral",
    }
    if method not in allowed:
        raise ValueError(f"unrecognized method {method}")
    if epochs < 1:
        raise ValueError("epochs must be positive")
    if method in {
        "constant", "sdil", "overclamp_constant", "overclamp_sdil"
    } and predictor is None:
        raise ValueError(f"{method} requires a predictor")
    gates = np.asarray(initial_gates, dtype=float).copy()
    if gates.shape != (circuit.edge_count,):
        raise ValueError("initial gate vector has the wrong shape")
    active_predictor = predictor.copy() if predictor is not None else None
    target_magnitude = (
        circuit.high_voltage
        if overclamp_target_magnitude_v is None
        else overclamp_target_magnitude_v)
    is_overclamp = method.startswith("overclamp")
    free_cache: list[Array | None] = [None] * len(dataset.labels_v)
    clamped_cache: list[Array | None] = [None] * len(dataset.labels_v)
    trace = []
    cumulative_learning_time = 0.0
    clamp_l2_time = 0.0
    max_clamp_displacement = 0.0
    local_updates = 0
    clipped_updates = 0

    for epoch in range(epochs):
        for sample, (inputs, label) in enumerate(zip(
            dataset.inputs_v, dataset.labels_v
        )):
            sources = circuit.source_values(*inputs)
            free_state = solve_grid_state(
                circuit,
                gates,
                sources,
                initial_state=free_cache[sample],
            )
            free_cache[sample] = free_state
            output_free = output_difference(circuit, free_state)
            error = label - output_free
            if label * error <= 0.0:
                continue
            if is_overclamp:
                output_clamped = output_free + overclamp_nudging * (
                    target_magnitude * np.sign(error) - output_free)
                duration = overclamp_time_seconds_per_v * abs(error)
            else:
                output_clamped = output_free + standard_nudging * error
                duration = standard_learning_time_seconds
            target_mean = float(np.mean(
                free_state[list(circuit.target_nodes)]))
            target_values = np.asarray((
                target_mean + 0.5 * output_clamped,
                target_mean - 0.5 * output_clamped,
            ))
            clamped_state = solve_grid_state(
                circuit,
                gates,
                sources,
                target_values=target_values,
                initial_state=(
                    free_state if clamped_cache[sample] is None
                    else clamped_cache[sample]),
            )
            clamped_cache[sample] = clamped_state
            free_drops = edge_voltage_drops(circuit, free_state)
            clamped_drops = edge_voltage_drops(circuit, clamped_state)
            ideal_rate = imperfection.ideal_rate(
                circuit.measured_learning_rate, free_drops, clamped_drops)
            observed_rate = imperfection.observed_rate(
                circuit.measured_learning_rate, free_drops, clamped_drops)
            neutral_bias = imperfection.neutral_bias(
                circuit.measured_learning_rate, free_drops)
            if method in {"clean", "overclamp_clean"}:
                applied_rate = ideal_rate
            elif method in {
                "oracle_neutral", "overclamp_oracle_neutral"
            }:
                applied_rate = observed_rate - neutral_bias
            elif method in {
                "constant", "sdil", "overclamp_constant", "overclamp_sdil"
            }:
                applied_rate = observed_rate - active_predictor.predict(
                    free_drops)
            else:
                applied_rate = observed_rate
            proposed = gates + duration * applied_rate
            clipped = np.clip(
                proposed, circuit.gate_minimum, circuit.gate_maximum)
            clipped_updates += int(np.any(clipped != proposed))
            gates = clipped
            local_updates += 1
            cumulative_learning_time += duration
            displacement = abs(output_clamped - output_free)
            clamp_l2_time += duration * displacement * displacement
            max_clamp_displacement = max(
                max_clamp_displacement, displacement)
        if epoch % record_every == 0 or epoch == epochs - 1:
            metrics, free_cache = evaluate_grid_classifier(
                circuit, gates, dataset, initial_states=free_cache)
            trace.append({"epoch": epoch, **metrics})

    final = trace[-1]
    return {
        "method": method,
        "epochs": epochs,
        "initial_gates_v": np.asarray(initial_gates).tolist(),
        "final_gates_v": gates.tolist(),
        "classification_error": final["classification_error"],
        "hinge_loss_v2": final["hinge_loss_v2"],
        "margin_success_fraction": final["margin_success_fraction"],
        "outputs_v": final["outputs_v"],
        "local_updates": local_updates,
        "cumulative_learning_time_seconds": float(cumulative_learning_time),
        "clamp_displacement_l2_time_v2_s": float(clamp_l2_time),
        "max_abs_clamp_displacement_v": float(max_clamp_displacement),
        "clipped_updates": clipped_updates,
        "trace": trace,
    }