"""Scaled damped Anderson solve for the AP1-R2c candidate source fixed point."""
from __future__ import annotations

from dataclasses import asdict, dataclass
import json
from pathlib import Path

import numpy as np
from scipy.interpolate import CubicSpline
from scipy.optimize import NoConvergence, anderson

from ap1_m1_background_preflight import decode_hard_pass_state
from ap1_r2c_high_precision_modes import completed_history_decimal
from ap1_r2c_mode_inheritance import inherited_trajectory
from ap1_r2c_self_consistent_candidate import (
    CandidateConfig,
    classical_terms,
    completed_candidate_sources_decimal,
    integrate_segment,
    match_seed_N,
    standard_stress,
)
from chi_background_closure import ChiParameters, ChiTrajectory, physical_shells, portal_terms
from planck2018_neutrino_closure import Planck2018Pilot


KEYS = ("rho", "pressure", "chi2")


@dataclass(frozen=True)
class AndersonConfig:
    candidate: CandidateConfig = CandidateConfig(candidate_nodes=33, fixed_point_steps=1)
    alpha: float = 0.05
    memory: int = 4
    regularization: float = 0.05
    max_iterations: int = 16
    residual_tolerance: float = 1.0e-4
    basis_degree: int = 8


class CandidateMap:
    def __init__(self, state_path: Path, cfg: CandidateConfig):
        self.cfg = cfg
        self.pilot = Planck2018Pilot()
        self.chi_p = ChiParameters()
        frozen = decode_hard_pass_state(state_path)
        self.seed = frozen[-1, :5].copy()
        self.prefix = inherited_trajectory(state_path, cfg.inherited_nodes)
        self.k, self.weights = physical_shells(self.chi_p.Lambda, nodes=cfg.mode_nodes)
        inherited = completed_history_decimal(
            self.prefix, self.k, self.weights, cfg.digits
        )["completed_history"]
        self.Nrel = np.linspace(0.0, cfg.delta_N, cfg.candidate_nodes)
        initial = {
            key: np.full(cfg.candidate_nodes, float(inherited[key][-1]))
            for key in KEYS
        }
        self.scales = {
            key: max(float(np.max(np.abs(value))), 1.0e-300)
            for key, value in initial.items()
        }
        self.initial = self.pack(initial)
        self.evaluations = 0
        self.last_data: dict | None = None

    def pack(self, q: dict[str, np.ndarray]) -> np.ndarray:
        return np.concatenate([np.asarray(q[key]) / self.scales[key] for key in KEYS])

    def unpack(self, x: np.ndarray) -> dict[str, np.ndarray]:
        n = self.cfg.candidate_nodes
        return {
            key: np.asarray(x[i * n:(i + 1) * n]) * self.scales[key]
            for i, key in enumerate(KEYS)
        }

    def evaluate(self, x: np.ndarray) -> tuple[np.ndarray, dict]:
        self.evaluations += 1
        q = self.unpack(np.asarray(x))
        if not all(np.all(np.isfinite(value)) for value in q.values()):
            raise RuntimeError("non-finite Anderson candidate source")
        N_seed = match_seed_N(self.seed, float(q["rho"][0]), self.pilot)
        y, Hdot = integrate_segment(
            self.seed, self.Nrel, N_seed, q, self.pilot, self.cfg
        )
        full = ChiTrajectory(
            N=np.concatenate([self.prefix.N, self.prefix.N[-1] + self.Nrel[1:]]),
            H=np.concatenate([self.prefix.H, y[1:, 4]]),
            Hdot=np.concatenate([self.prefix.Hdot, Hdot[1:]]),
            sigma=np.concatenate([self.prefix.sigma, y[1:, 0]]),
            theta=np.concatenate([self.prefix.theta, y[1:, 2]]),
        )
        candidate = ChiTrajectory(
            N=self.prefix.N[-1] + self.Nrel, H=y[:, 4], Hdot=Hdot,
            sigma=y[:, 0], theta=y[:, 2],
        )
        fresh, modes = completed_candidate_sources_decimal(
            full, candidate, self.k, self.weights, self.cfg.digits
        )
        residual = self.pack({key: fresh[key] - q[key] for key in KEYS})
        data = {
            "q": q, "fresh": fresh, "N_seed": N_seed, "y": y,
            "Hdot": Hdot, "modes": modes,
            "max_scaled_residual": float(np.max(np.abs(residual))),
        }
        self.last_data = data
        return residual, data

    def residual(self, x: np.ndarray) -> np.ndarray:
        return self.evaluate(x)[0]


class SpectralCandidateMap:
    """Smooth Chebyshev representation of the three candidate sources."""

    def __init__(self, grid_model: CandidateMap, degree: int):
        if not (1 <= degree < grid_model.cfg.candidate_nodes):
            raise ValueError("invalid candidate source basis degree")
        self.grid = grid_model
        self.degree = degree
        self.xgrid = np.linspace(-1.0, 1.0, grid_model.cfg.candidate_nodes)
        initial_q = grid_model.unpack(grid_model.initial)
        self.initial = self.fit(initial_q)
        self.last_data: dict | None = None

    def fit(self, q: dict[str, np.ndarray]) -> np.ndarray:
        return np.concatenate([
            np.polynomial.chebyshev.chebfit(
                self.xgrid, np.asarray(q[key]) / self.grid.scales[key], self.degree
            )
            for key in KEYS
        ])

    def expand(self, coefficients: np.ndarray) -> dict[str, np.ndarray]:
        width = self.degree + 1
        return {
            key: self.grid.scales[key] * np.polynomial.chebyshev.chebval(
                self.xgrid, coefficients[i * width:(i + 1) * width]
            )
            for i, key in enumerate(KEYS)
        }

    def evaluate(self, coefficients: np.ndarray) -> tuple[np.ndarray, dict]:
        q = self.expand(np.asarray(coefficients))
        _grid_residual, data = self.grid.evaluate(self.grid.pack(q))
        coefficient_residual = self.fit(data["fresh"]) - np.asarray(coefficients)
        data["coefficient_residual"] = float(np.max(np.abs(coefficient_residual)))
        projected_fresh = self.expand(self.fit(data["fresh"]))
        data["spectral_projection_error"] = max(
            float(np.max(np.abs(projected_fresh[key] - data["fresh"][key])) /
                  self.grid.scales[key])
            for key in KEYS
        )
        self.last_data = data
        return coefficient_residual, data

    def residual(self, coefficients: np.ndarray) -> np.ndarray:
        return self.evaluate(coefficients)[0]


def diagnostics(model: CandidateMap, data: dict) -> dict:
    q = data["q"]; fresh = data["fresh"]; y = data["y"]
    dlist = [classical_terms(row) for row in y]
    rho_s, _p_s = standard_stress(data["N_seed"] + model.Nrel, model.pilot)
    total = np.array([d["rho"] for d in dlist]) + rho_s + q["rho"]
    friedmann = (3.0 * y[:, 4]**2 - total) / max(float(np.max(np.abs(total))), 1.0e-300)
    mass2 = np.asarray(portal_terms(y[:, 0], y[:, 2], model.chi_p)[0])
    rho_dot = y[:, 4] * CubicSpline(model.Nrel, fresh["rho"])(model.Nrel, 1)
    mass2_dot = y[:, 4] * CubicSpline(model.Nrel, mass2)(model.Nrel, 1)
    ward_terms = (
        rho_dot,
        3.0 * y[:, 4] * (fresh["rho"] + fresh["pressure"]),
        -0.5 * mass2_dot * fresh["chi2"],
    )
    ward = sum(ward_terms)
    ward_scale = max(*(float(np.max(np.abs(term))) for term in ward_terms), 1.0e-300)
    return {
        "finite": bool(np.all(np.isfinite(y))),
        "positive_H": bool(np.all(y[:, 4] > 0.0)),
        "min_PX": float(min(d["PX"] for d in dlist)),
        "min_K": float(min(d["K"] for d in dlist)),
        "max_abs_friedmann_residual": float(np.max(np.abs(friedmann))),
        "ward_normalized": float(np.max(np.abs(ward)) / ward_scale),
        "validation_source_relative_change": data["max_scaled_residual"],
        "spectral_coefficient_residual": data["coefficient_residual"],
        "spectral_projection_error": data["spectral_projection_error"],
        "validation_wronskian_relative_error": data["modes"]["wronskian_relative_error"],
    }


def run_anderson(state_path: Path, cfg: AndersonConfig | None = None) -> dict:
    cfg = AndersonConfig() if cfg is None else cfg
    model = CandidateMap(state_path, cfg.candidate)
    spectral = SpectralCandidateMap(model, cfg.basis_degree)
    log: list[dict] = []
    best = {"norm": float("inf"), "x": spectral.initial.copy()}

    def callback(x: np.ndarray, f: np.ndarray) -> None:
        norm = float(np.max(np.abs(f)))
        log.append({"iteration": len(log) + 1, "scaled_residual": norm})
        if norm < best["norm"]:
            best["norm"] = norm
            best["x"] = np.asarray(x).copy()

    converged = True
    try:
        solution = anderson(
            spectral.residual, spectral.initial, alpha=cfg.alpha, M=cfg.memory,
            w0=cfg.regularization, maxiter=cfg.max_iterations,
            f_tol=cfg.residual_tolerance, line_search="armijo", callback=callback,
        )
    except NoConvergence:
        converged = False
        solution = best["x"]
    residual, final = spectral.evaluate(solution)
    final_norm = float(np.max(np.abs(residual)))
    converged = converged and final_norm <= cfg.residual_tolerance
    return {
        "schema": "apeiron-ap1-r2c-candidate-anderson-v1.0",
        "classification": "ANDERSON_CANDIDATE_CONVERGED" if converged else "ANDERSON_CANDIDATE_NOT_CONVERGED",
        "config": {**asdict(cfg), "candidate": asdict(cfg.candidate)},
        "function_evaluations": model.evaluations,
        "iterations": log,
        "diagnostics": diagnostics(model, final),
        "old_solver_or_physical_map_called": False,
        "claim_boundary": "numerical short-candidate gate only; no production observable or empirical claim",
    }


def main() -> None:
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("state", type=Path)
    parser.add_argument("--output", type=Path)
    args = parser.parse_args()
    rendered = json.dumps(run_anderson(args.state), indent=2) + "\n"
    if args.output:
        args.output.write_text(rendered, encoding="utf-8")
    else:
        print(rendered, end="")


if __name__ == "__main__":
    main()