"""Bounded safeguarded Newton-Krylov pilot for the AP1-R2c implicit DAE."""
from __future__ import annotations

import json
from pathlib import Path

import numpy as np
from scipy.optimize import NoConvergence, newton_krylov

from ap1_r2c_implicit_dae_solver import ImplicitDAEResidual


def run_newton(state_path: Path, max_iterations: int = 6) -> dict:
    system = ImplicitDAEResidual(state_path)
    log: list[dict] = []
    rejected = 0
    best = {"norm": float("inf"), "x": system.initial.copy()}

    def safe(vector: np.ndarray) -> np.ndarray:
        nonlocal rejected
        try:
            return system.evaluate(vector)
        except (ValueError, RuntimeError, FloatingPointError):
            rejected += 1
            # A deterministic outward penalty lets Armijo shorten the step;
            # it is never eligible for a physical accepted state.
            return np.full_like(vector, 1.0e6) + (vector - system.initial)

    initial_residual = safe(system.initial)
    best["norm"] = float(np.max(np.abs(initial_residual)))

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

    converged = True
    try:
        solution = newton_krylov(
            safe, system.initial, method="lgmres", inner_maxiter=12, outer_k=4,
            maxiter=max_iterations, f_tol=1.0e-4, line_search="armijo",
            rdiff=1.0e-6, callback=callback,
        )
    except NoConvergence:
        converged = False
        solution = best["x"]
    except (ValueError, RuntimeError, FloatingPointError) as exc:
        converged = False
        solution = best["x"]
        log.append({"termination": type(exc).__name__, "message": str(exc)})

    final = system.evaluate(solution)
    final_norm = float(np.max(np.abs(final)))
    converged = converged and final_norm <= 1.0e-4
    return {
        "schema": "apeiron-ap1-r2c-implicit-dae-newton-v1.0",
        "classification": "DAE_NEWTON_17_NODE_CONVERGED" if converged else "DAE_NEWTON_17_NODE_NOT_CONVERGED",
        "max_outer_iterations": max_iterations,
        "iterations": log,
        "rejected_nonphysical_trials": rejected,
        "initial_max_abs_scaled_residual": float(np.max(np.abs(initial_residual))),
        "best_max_abs_scaled_residual": best["norm"],
        "final_residual_blocks": system.last_blocks,
        "old_solver_or_physical_map_called": False,
        "claim_boundary": "bounded 17-node DAE pilot only; no multi-resolution PASS or observable released",
    }


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_newton(args.state),indent=2)+"\n"
    if args.output: args.output.write_text(rendered,encoding="utf-8")
    else: print(rendered,end="")


if __name__ == "__main__": main()