"""Local square-law imperfection model for physical coupled learning.

The model follows the voltage-square learning rule and the three hardware
error sources discussed in Appendix C of Dillavou et al.
(arXiv:2505.22887v2): differential gain error, twin-state voltage mismatch,
and multiplier output offset.  All predictors and updates are explicit local
NumPy operations with no autodiff path.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np

from sdil.physical_coupled import (
    Array,
    Circuit,
    Task,
    free_output,
    overclamped_clean_rate,
    standard_clean_rate,
    task_errors,
)


def edge_voltage_drops(circuit: Circuit, output: float) -> Array:
    """Return the two locally measured positive edge-voltage drops."""
    return np.asarray(
        (output - circuit.low, circuit.high - output), dtype=float)


@dataclass(frozen=True)
class DifferentialSquareLawImperfection:
    """Fixed per-edge imperfections in the free/clamped measurement paths."""

    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:
        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 np.asarray(value).shape != (2,):
                raise ValueError("two-edge imperfection arrays must have shape (2,)")

    @classmethod
    def ideal(cls) -> "DifferentialSquareLawImperfection":
        return cls(
            free_gain=np.ones(2, dtype=float),
            clamped_gain=np.ones(2, dtype=float),
            free_input_offset_v=np.zeros(2, dtype=float),
            clamped_input_offset_v=np.zeros(2, dtype=float),
            multiplier_output_offset_v_per_s=np.zeros(2, dtype=float),
        )

    @classmethod
    def sample_appendix_c(
        cls,
        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,
    ) -> "DifferentialSquareLawImperfection":
        """Sample a device using the component scales stated in Appendix C."""
        if gain_standard_deviation < 0.0:
            raise ValueError("gain standard deviation must be nonnegative")
        if twin_mismatch_standard_deviation_v < 0.0:
            raise ValueError("twin mismatch standard deviation must be nonnegative")
        if multiplier_offset_standard_deviation_v_per_s < 0.0:
            raise ValueError("multiplier offset scale must be nonnegative")
        rng = np.random.default_rng(seed)
        common_gain = 1.0 + rng.normal(0.0, gain_standard_deviation, size=2)
        differential_gain = rng.normal(
            0.0, gain_standard_deviation, size=2)
        free_gain = common_gain + 0.5 * differential_gain
        clamped_gain = common_gain - 0.5 * differential_gain
        common_input_offset = rng.normal(
            0.0, twin_mismatch_standard_deviation_v, size=2)
        differential_input_offset = rng.normal(
            0.0, twin_mismatch_standard_deviation_v, size=2)
        free_input_offset = common_input_offset + 0.5 * differential_input_offset
        clamped_input_offset = (
            common_input_offset - 0.5 * differential_input_offset)
        multiplier_offset = rng.normal(
            0.0, multiplier_offset_standard_deviation_v_per_s, size=2)
        return cls(
            free_gain=free_gain,
            clamped_gain=clamped_gain,
            free_input_offset_v=free_input_offset,
            clamped_input_offset_v=clamped_input_offset,
            multiplier_output_offset_v_per_s=multiplier_offset,
        )

    def observed_rate(
        self,
        circuit: Circuit,
        output_free: float,
        output_clamped: float,
    ) -> Array:
        free_drop = edge_voltage_drops(circuit, output_free)
        clamped_drop = edge_voltage_drops(circuit, output_clamped)
        measured_free = (
            self.free_gain * free_drop + self.free_input_offset_v)
        measured_clamped = (
            self.clamped_gain * clamped_drop + self.clamped_input_offset_v)
        return (
            circuit.measured_learning_rate
            * (np.square(measured_free) - np.square(measured_clamped))
            + self.multiplier_output_offset_v_per_s
        )

    def ideal_rate(
        self,
        circuit: Circuit,
        output_free: float,
        output_clamped: float,
    ) -> Array:
        free_drop = edge_voltage_drops(circuit, output_free)
        clamped_drop = edge_voltage_drops(circuit, output_clamped)
        return circuit.measured_learning_rate * (
            np.square(free_drop) - np.square(clamped_drop))

    def neutral_bias(self, circuit: Circuit, output_free: float) -> Array:
        return self.observed_rate(circuit, output_free, output_free)

    def as_dict(self) -> dict:
        return {
            "free_gain": np.asarray(self.free_gain).tolist(),
            "clamped_gain": np.asarray(self.clamped_gain).tolist(),
            "free_input_offset_v": np.asarray(
                self.free_input_offset_v).tolist(),
            "clamped_input_offset_v": np.asarray(
                self.clamped_input_offset_v).tolist(),
            "multiplier_output_offset_v_per_s": np.asarray(
                self.multiplier_output_offset_v_per_s).tolist(),
        }


@dataclass
class LocalPolynomialPredictor:
    """Independent per-edge polynomial filters trained by normalized LMS."""

    feature_center: Array
    feature_scale: Array
    coefficients: Array

    @classmethod
    def zeros(
        cls,
        feature_center: Array,
        feature_scale: Array,
        *,
        degree: int,
    ) -> "LocalPolynomialPredictor":
        if degree < 0:
            raise ValueError("polynomial degree must be nonnegative")
        center = np.asarray(feature_center, dtype=float)
        scale = np.asarray(feature_scale, dtype=float)
        if center.shape != (2,) or scale.shape != (2,):
            raise ValueError("feature metadata must have shape (2,)")
        if np.any(scale <= 0.0):
            raise ValueError("feature scales must be positive")
        return cls(
            feature_center=center.copy(),
            feature_scale=scale.copy(),
            coefficients=np.zeros((2, 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 != (2,):
            raise ValueError("local state must have shape (2,)")
        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 update(
        self,
        local_state: Array,
        neutral_measurement: Array,
        learning_rate: float,
    ) -> Array:
        measurement = np.asarray(neutral_measurement, dtype=float)
        if measurement.shape != (2,):
            raise ValueError("neutral measurement must have shape (2,)")
        features = self.features(local_state)
        prediction = np.sum(self.coefficients * features, axis=1)
        residual = measurement - prediction
        normalization = np.sum(features * features, axis=1, keepdims=True)
        self.coefficients += (
            learning_rate * residual[:, None] * features
            / np.maximum(normalization, 1e-12)
        )
        return residual

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


def calibrate_polynomial_predictor(
    predictor: LocalPolynomialPredictor,
    local_states: Array,
    neutral_measurements: Array,
    *,
    epochs: int,
    learning_rate: float,
) -> int:
    states = np.asarray(local_states, dtype=float)
    measurements = np.asarray(neutral_measurements, dtype=float)
    if states.ndim != 2 or states.shape[1] != 2:
        raise ValueError("local states must have shape (observations, 2)")
    if measurements.shape != states.shape:
        raise ValueError("neutral measurements must match local states")
    if epochs < 1:
        raise ValueError("epochs must be positive")
    observations = 0
    for _ in range(epochs):
        for state, measurement in zip(states, measurements):
            predictor.update(state, measurement, learning_rate)
            observations += 1
    return observations


def fit_polynomial_predictor(
    predictor: LocalPolynomialPredictor,
    local_states: Array,
    neutral_measurements: Array,
    *,
    ridge: float = 1e-12,
) -> int:
    """Fit independent per-edge sufficient statistics by local least squares."""
    states = np.asarray(local_states, dtype=float)
    measurements = np.asarray(neutral_measurements, dtype=float)
    if states.ndim != 2 or states.shape[1] != 2:
        raise ValueError("local states must have shape (observations, 2)")
    if measurements.shape != states.shape:
        raise ValueError("neutral measurements must match local states")
    if ridge < 0.0:
        raise ValueError("ridge must be nonnegative")
    features = np.asarray([predictor.features(state) for state in states])
    for edge in range(2):
        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))


def simulate_imperfect_alternating_tasks(
    circuit: Circuit,
    tasks: tuple[Task, Task],
    imperfection: DifferentialSquareLawImperfection,
    *,
    method: str,
    period_seconds: float,
    cycles: int,
    initial_gates: Array,
    predictor: LocalPolynomialPredictor | None = None,
    overclamp_nudging: float = 0.25,
    overclamp_magnitude: float | None = None,
    summary_cycles: int = 20,
    record_history: bool = False,
) -> dict:
    """Alternate two tasks under a fixed local hardware imperfection."""
    allowed = {
        "raw",
        "constant",
        "sdil",
        "oracle_neutral",
        "clean",
        "overclamp",
        "overclamp_constant",
        "overclamp_sdil",
        "overclamp_oracle_neutral",
        "overclamp_clean",
    }
    if method not in allowed:
        raise ValueError(f"unrecognized method {method}")
    if method in {
        "constant", "sdil", "overclamp_constant", "overclamp_sdil"
    } and predictor is None:
        raise ValueError(f"{method} requires a predictor")
    if len(tasks) != 2:
        raise ValueError("exactly two alternating tasks are required")
    if period_seconds <= 0.0 or cycles < 1:
        raise ValueError("period and cycles must be positive")
    gates = np.asarray(initial_gates, dtype=float).copy()
    if gates.shape != (2,):
        raise ValueError("initial gates must have shape (2,)")
    active_predictor = predictor.copy() if predictor is not None else None
    nominal_step = circuit.integration_step_seconds
    half_steps = max(1, int(round(period_seconds / (2.0 * nominal_step))))
    initial_error_scale = float(np.mean([
        abs(task.label_voltage - free_output(
            circuit, gates, task.input_voltage))
        for task in tasks
    ]))
    initial_error_scale = max(initial_error_scale, 1e-6)

    combined_error_history = []
    cycle_span_history = []
    residual_neutral_history = []
    rate_distortion_history = []
    learning_on_time = 0.0
    clamp_l2_time = 0.0
    max_clamp_displacement = 0.0
    clipped_updates = 0

    is_overclamp = method.startswith("overclamp")
    for _ in range(cycles):
        half_endpoints = []
        half_errors = []
        for task in tasks:
            for _ in range(half_steps):
                if is_overclamp:
                    _, output_free, output_clamped = overclamped_clean_rate(
                        circuit,
                        gates,
                        task,
                        nudging=overclamp_nudging,
                        clamp_magnitude=overclamp_magnitude,
                        exact_target=True,
                    )
                    duration = nominal_step * abs(
                        task.label_voltage - output_free) / initial_error_scale
                else:
                    _, output_free, output_clamped = standard_clean_rate(
                        circuit, gates, task)
                    duration = nominal_step
                ideal_rate = imperfection.ideal_rate(
                    circuit, output_free, output_clamped)
                observed_rate = imperfection.observed_rate(
                    circuit, output_free, output_clamped)
                local_state = edge_voltage_drops(circuit, output_free)
                neutral_bias = imperfection.neutral_bias(circuit, output_free)

                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(
                        local_state)
                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
                learning_on_time += duration
                displacement = abs(output_clamped - output_free)
                clamp_l2_time += duration * displacement * displacement
                max_clamp_displacement = max(
                    max_clamp_displacement, displacement)
                predicted_neutral = (
                    np.zeros(2, dtype=float)
                    if active_predictor is None
                    else active_predictor.predict(local_state)
                )
                residual_neutral_history.append(float(np.mean(
                    np.square(neutral_bias - predicted_neutral))))
                rate_distortion_history.append(float(np.mean(
                    np.square(applied_rate - ideal_rate))))
            half_endpoints.append(gates.copy())
            half_errors.append(task_errors(circuit, gates, tasks))
        combined_error_history.append(float(np.mean(half_errors)))
        cycle_span_history.append(float(np.linalg.norm(
            half_endpoints[1] - half_endpoints[0])))

    summary_count = min(summary_cycles, cycles)
    combined = np.asarray(combined_error_history[-summary_count:])
    spans = np.asarray(cycle_span_history[-summary_count:])
    result = {
        "method": method,
        "period_seconds": period_seconds,
        "cycles": cycles,
        "half_steps": half_steps,
        "initial_gates": np.asarray(initial_gates, dtype=float).tolist(),
        "final_gates": gates.tolist(),
        "mean_combined_error": float(np.mean(combined)),
        "std_combined_error": float(np.std(combined)),
        "mean_cycle_span": float(np.mean(spans)),
        "std_cycle_span": float(np.std(spans)),
        "learning_on_time_seconds": float(learning_on_time),
        "clamp_displacement_l2_time_v2_s": float(clamp_l2_time),
        "max_abs_clamp_displacement_v": float(max_clamp_displacement),
        "mean_residual_neutral_bias_mse": float(np.mean(
            residual_neutral_history)),
        "mean_applied_rate_distortion_mse": float(np.mean(
            rate_distortion_history)),
        "clipped_updates": clipped_updates,
    }
    if record_history:
        result.update({
            "combined_error_history": combined_error_history,
            "cycle_span_history": cycle_span_history,
            "residual_neutral_mse_history": residual_neutral_history,
            "rate_distortion_mse_history": rate_distortion_history,
        })
    return result