"""Sparse digital coupled-learning grids for controlled scaling experiments.

The network is a linear resistor lattice.  Its local learning signal is the
same free-minus-clamped voltage-square difference used by coupled learning.
Component imperfections reuse the per-edge measurement model used by the
hardware-realistic simulator in :mod:`sdil.physical_grid`.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np
from scipy.sparse import coo_matrix
from scipy.sparse.linalg import spsolve

from sdil.physical_grid import (
    GridCircuit,
    GridSquareLawImperfection,
    RingClassificationDataset,
    edge_voltage_drops,
    output_difference,
)


Array = np.ndarray


def make_scaled_grid(side: int) -> GridCircuit:
    """Return a square grid whose side-4 boundary layout matches Figure 5."""
    if side < 4 or side % 4:
        raise ValueError("grid side must be a positive multiple of four")

    def node(row: int, column: int) -> int:
        return row * side + column

    return GridCircuit(
        rows=side,
        columns=side,
        source_nodes=(
            node(3 * side // 4, 3 * side // 4),
            node(3 * side // 4, side // 4),
            node(side // 4, 3 * side // 4),
            node(side // 4, side // 4),
        ),
        target_nodes=(
            node(side // 2, side // 2),
            node(side // 2, 0),
        ),
    )


def tile_figure5_gates(base_gates: Array, side: int) -> Array:
    """Tile a released 4-by-4 horizontal/vertical gate pattern."""
    gates = np.asarray(base_gates, dtype=float)
    if gates.shape != (32,):
        raise ValueError("the released base gate vector must have 32 entries")
    if side < 4 or side % 4:
        raise ValueError("grid side must be a positive multiple of four")
    horizontal = gates[:16].reshape(4, 4)
    vertical = gates[16:].reshape(4, 4)
    tiled_horizontal = np.asarray([
        horizontal[row % 4, column % 4]
        for row in range(side)
        for column in range(side)
    ])
    tiled_vertical = np.asarray([
        vertical[row % 4, column % 4]
        for row in range(side)
        for column in range(side)
    ])
    return np.concatenate((tiled_horizontal, tiled_vertical))


def solve_linear_grid_state(
    circuit: GridCircuit,
    gates: Array,
    source_values: Array,
    *,
    target_values: Array | None = None,
) -> Array:
    """Solve the sparse linear Kirchhoff system with fixed boundary nodes."""
    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")
    conductances = circuit.conductance_scale * (
        gates - circuit.threshold_voltage)
    if np.any(conductances <= 0.0):
        raise ValueError("all digital conductances must be positive")

    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_nodes = [
        node for node in range(circuit.node_count) if node not in fixed
    ]
    unknown_index = {node: index for index, node in enumerate(unknown_nodes)}
    matrix_rows: list[int] = []
    matrix_columns: list[int] = []
    matrix_values: list[float] = []
    rhs = np.zeros(len(unknown_nodes), dtype=float)

    for conductance, (first, second) in zip(
        conductances, circuit.edge_pairs
    ):
        for node, neighbour in ((first, second), (second, first)):
            if node in fixed:
                continue
            row = unknown_index[node]
            matrix_rows.append(row)
            matrix_columns.append(row)
            matrix_values.append(float(conductance))
            if neighbour in fixed:
                rhs[row] += conductance * fixed[neighbour]
            else:
                matrix_rows.append(row)
                matrix_columns.append(unknown_index[neighbour])
                matrix_values.append(float(-conductance))

    matrix = coo_matrix(
        (matrix_values, (matrix_rows, matrix_columns)),
        shape=(len(unknown_nodes), len(unknown_nodes)),
    ).tocsr()
    unknown_voltages = np.asarray(spsolve(matrix, rhs), dtype=float)
    if not np.all(np.isfinite(unknown_voltages)):
        raise RuntimeError("linear grid solve returned a nonfinite state")

    voltages = np.empty(circuit.node_count, dtype=float)
    for node, voltage in fixed.items():
        voltages[node] = voltage
    voltages[unknown_nodes] = unknown_voltages
    return voltages


def evaluate_digital_grid(
    circuit: GridCircuit,
    gates: Array,
    dataset: RingClassificationDataset,
) -> dict:
    outputs = []
    for inputs in dataset.inputs_v:
        state = solve_linear_grid_state(
            circuit, gates, circuit.source_values(*inputs))
        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(),
    }


@dataclass(frozen=True)
class DigitalTrainingConfig:
    epochs: int = 600
    record_every: int = 10
    standard_nudging: float = 128.0 / 129.0
    learning_time_seconds: float = 1.0e-3
    overclamp_nudging: float = 32.0 / 129.0
    overclamp_time_seconds_per_v: float = 2.5e-3
    overclamp_target_magnitude_v: float | None = None

    def __post_init__(self) -> None:
        if self.epochs < 1 or self.record_every < 1:
            raise ValueError("epochs and record interval must be positive")
        if self.learning_time_seconds <= 0.0:
            raise ValueError("learning time must be positive")


def train_digital_grid(
    circuit: GridCircuit,
    initial_gates: Array,
    dataset: RingClassificationDataset,
    imperfection: GridSquareLawImperfection,
    *,
    method: str,
    config: DigitalTrainingConfig,
    constant_bias_v_per_s: Array | None = None,
    noise_seed: int = 0,
) -> dict:
    """Train with a hand-written local coupled-learning update."""
    allowed = {
        "clean",
        "matched_noise",
        "raw",
        "constant",
        "sdil",
        "overclamp_clean",
        "overclamp",
        "overclamp_sdil",
    }
    if method not in allowed:
        raise ValueError(f"unrecognized method {method}")
    if method == "constant" and constant_bias_v_per_s is None:
        raise ValueError("constant calibration requires a per-edge baseline")

    gates = np.asarray(initial_gates, dtype=float).copy()
    if gates.shape != (circuit.edge_count,):
        raise ValueError("initial gate vector has the wrong shape")
    if constant_bias_v_per_s is not None:
        constant_bias = np.asarray(constant_bias_v_per_s, dtype=float)
        if constant_bias.shape != gates.shape:
            raise ValueError("constant baseline has the wrong shape")
    else:
        constant_bias = None

    rng = np.random.default_rng(noise_seed)
    local_updates = 0
    neutral_observations = 0
    clipped_updates = 0
    cumulative_learning_time = 0.0
    trace = [{
        "epoch": 0,
        "local_updates": 0,
        "cumulative_learning_time_seconds": 0.0,
        **evaluate_digital_grid(circuit, gates, dataset),
    }]

    for epoch in range(1, config.epochs + 1):
        for inputs, label in zip(dataset.inputs_v, dataset.labels_v):
            sources = circuit.source_values(*inputs)
            free_state = solve_linear_grid_state(circuit, gates, sources)
            output_free = output_difference(circuit, free_state)
            error = label - output_free
            if label * error <= 0.0:
                continue

            free_drops = edge_voltage_drops(circuit, free_state)
            is_overclamp = method.startswith("overclamp")
            if is_overclamp:
                target_magnitude = (
                    circuit.high_voltage
                    if config.overclamp_target_magnitude_v is None
                    else config.overclamp_target_magnitude_v
                )
                output_clamped = output_free + config.overclamp_nudging * (
                    target_magnitude * np.sign(error) - output_free)
                duration = (
                    config.overclamp_time_seconds_per_v * abs(error))
            else:
                output_clamped = output_free + (
                    config.standard_nudging * error)
                duration = config.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_linear_grid_state(
                circuit, gates, sources, target_values=target_values)
            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)

            if method in {"clean", "overclamp_clean"}:
                applied_rate = ideal_rate
            elif method == "matched_noise":
                measurement_error = observed_rate - ideal_rate
                random_sign = rng.choice((-1.0, 1.0), size=len(gates))
                applied_rate = ideal_rate + random_sign * np.abs(
                    measurement_error)
            elif method == "constant":
                applied_rate = observed_rate - constant_bias
            elif method in {"sdil", "overclamp_sdil"}:
                neutral_rate = imperfection.neutral_bias(
                    circuit.measured_learning_rate, free_drops)
                applied_rate = observed_rate - neutral_rate
                neutral_observations += 1
            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

        if epoch % config.record_every == 0 or epoch == config.epochs:
            trace.append({
                "epoch": epoch,
                "local_updates": local_updates,
                "cumulative_learning_time_seconds": float(
                    cumulative_learning_time),
                **evaluate_digital_grid(circuit, gates, dataset),
            })

    stable_zero_index = next((
        index for index, record in enumerate(trace)
        if record["classification_error"] == 0.0
        and all(
            later["classification_error"] == 0.0
            for later in trace[index:]
        )
    ), None)
    if stable_zero_index is not None:
        stable_record = trace[stable_zero_index]
        epochs_to_stable_zero = int(stable_record["epoch"])
        updates_to_stable_zero = int(stable_record["local_updates"])
        exposure_to_stable_zero = float(
            stable_record["cumulative_learning_time_seconds"])
        reached_stable_zero = True
    else:
        epochs_to_stable_zero = config.epochs
        updates_to_stable_zero = local_updates
        exposure_to_stable_zero = float(cumulative_learning_time)
        reached_stable_zero = False
    epoch_axis = np.asarray([record["epoch"] for record in trace])
    error_axis = np.asarray([
        record["classification_error"] for record in trace])
    error_auc = float(np.trapezoid(error_axis, epoch_axis) / config.epochs)
    final = trace[-1]
    return {
        "method": method,
        "classification_error": final["classification_error"],
        "hinge_loss_v2": final["hinge_loss_v2"],
        "margin_success_fraction": final["margin_success_fraction"],
        "outputs_v": final["outputs_v"],
        "reached_stable_zero_error": reached_stable_zero,
        "restricted_epochs_to_stable_zero_error": epochs_to_stable_zero,
        "restricted_updates_to_stable_zero_error": updates_to_stable_zero,
        "restricted_edge_updates_to_stable_zero_error": int(
            updates_to_stable_zero * circuit.edge_count),
        "restricted_learning_time_to_stable_zero_seconds": (
            exposure_to_stable_zero),
        "classification_error_auc": error_auc,
        "local_updates": local_updates,
        "local_edge_updates": int(local_updates * circuit.edge_count),
        "neutral_observations": neutral_observations,
        "neutral_scalar_observations": int(
            neutral_observations * circuit.edge_count),
        "cumulative_learning_time_seconds": float(cumulative_learning_time),
        "clipped_updates": clipped_updates,
        "final_gates_v": gates.tolist(),
        "trace": trace,
    }