"""Bounded p42 minimax with exact Friedmann-density compensation."""
from __future__ import annotations

import json
from pathlib import Path

import numpy as np
from scipy.optimize import linprog

from ap1_r2c_implicit_dae_multiresolution import (
    _physical_diagnostics,
    balanced_algebraic_projection,
)
from ap1_r2c_implicit_dae_preflight import positive_friedmann_root
from ap1_r2c_implicit_dae_solver import X_KEYS
from ap1_r2c_self_consistent_candidate import (
    classical_terms,
    portal_terms,
    standard_stress,
)
from p42_minimax_recheck import reconstruct_candidate


def solve_fixed_h_minimax(
    system,
    candidate: np.ndarray,
    radii: tuple[float, ...] = (1.0e-7, 3.0e-7, 1.0e-6, 3.0e-6),
    full_radius_diagnostics: bool = False,
    objective_mode: str = "global_minimax",
    background_residual_cap: float = 8.0e-5,
) -> tuple[dict, np.ndarray]:
    """Return a fail-closed fixed-H minimax result and its checked vector."""
    if objective_mode not in {"global_minimax", "rho_lexicographic"}:
        raise ValueError("recognized fixed-H minimax objective required")
    if not 0.0 < background_residual_cap <= 1.0e-4:
        raise ValueError("background residual cap must preserve the DAE gate")
    baseline = system.evaluate(candidate)
    n = system.cfg.candidate_nodes
    X0, Q0, _N_seed0 = system.unpack(candidate)
    states0 = np.column_stack([X0[key] for key in X_KEYS])
    rho_cl0 = np.array(
        [classical_terms(np.r_[row, system.seed[4]])["rho"] for row in states0]
    )
    rho_std0, _ = standard_stress(_N_seed0 + system.Nrel, system.pilot)
    target_total_density = rho_cl0 + rho_std0 + Q0["rho"]
    fresh_rho = Q0["rho"] - baseline[4 * n : 5 * n] * system.q_scales["rho"]

    def vector_from_u(u: np.ndarray) -> tuple[np.ndarray, float]:
        X, Q, N_seed = system.unpack(candidate)
        X["sigma_dot"][1:] += u[: n - 1] * system.x_scales["sigma_dot"]
        X["theta_dot"][1:] += u[n - 1 :] * system.x_scales["theta_dot"]
        states = np.column_stack([X[key] for key in X_KEYS])
        rho_cl = np.array(
            [
                classical_terms(np.r_[row, system.seed[4]])["rho"]
                for row in states
            ]
        )
        Q["rho"] += rho_cl0 - rho_cl
        # Close the last floating-point cancellation ulps in the exact same
        # summation order used by the Friedmann root.  This is a numerical
        # projection onto the unchanged total-density manifold.
        for _ in range(3):
            current_total = rho_cl + rho_std0 + Q["rho"]
            Q["rho"] += target_total_density - current_total
        vector = system.pack(X, Q, N_seed)
        # The scaled pack/unpack itself can reintroduce one cancellation ulp.
        # Project through that representation so the evaluated vector, not
        # merely its pre-packed fields, remains on the same density manifold.
        for _ in range(5):
            packed_X, packed_Q, packed_N = system.unpack(vector)
            packed_states = np.column_stack([
                packed_X[key] for key in X_KEYS
            ])
            packed_rho_cl = np.array([
                classical_terms(np.r_[row, system.seed[4]])["rho"]
                for row in packed_states
            ])
            current_total = packed_rho_cl + rho_std0 + packed_Q["rho"]
            packed_Q["rho"] += target_total_density - current_total
            vector = system.pack(packed_X, packed_Q, packed_N)
        packed_X, packed_Q, _ = system.unpack(vector)
        packed_states = np.column_stack([packed_X[key] for key in X_KEYS])
        packed_rho_cl = np.array([
            classical_terms(np.r_[row, system.seed[4]])["rho"]
            for row in packed_states
        ])
        total_error = float(np.max(np.abs(
            packed_rho_cl + rho_std0 + packed_Q["rho"]
            - target_total_density
        )))
        return vector, total_error

    def reduced_residual(u: np.ndarray) -> np.ndarray:
        vector, _ = vector_from_u(u)
        X, Q, N_seed = system.unpack(vector)
        states_without_H = np.column_stack([X[key] for key in X_KEYS])
        rho_cl = np.array(
            [
                classical_terms(np.r_[row, system.seed[4]])["rho"]
                for row in states_without_H
            ]
        )
        rho_std, _ = standard_stress(N_seed + system.Nrel, system.pilot)
        H = positive_friedmann_root(rho_cl + rho_std + Q["rho"])
        states = np.column_stack([states_without_H, H])
        dlist = [classical_terms(row) for row in states]
        mass = [portal_terms(row[0], row[2], system.chi_p) for row in states]
        rhs = {key: np.empty(n) for key in X_KEYS}
        rhs["sigma"] = X["sigma_dot"] / H
        rhs["theta"] = X["theta_dot"] / H
        rhs["sigma_dot"] = np.array(
            [
                (
                    -3.0 * H[i] * X["sigma_dot"][i]
                    + dlist[i]["P_sigma"]
                    - 0.5 * float(mass[i][1]) * Q["chi2"][i]
                )
                / H[i]
                for i in range(n)
            ]
        )
        rhs["theta_dot"] = np.array(
            [
                (
                    dlist[i]["P_theta"]
                    - 0.5 * float(mass[i][2]) * Q["chi2"][i]
                    - 3.0 * H[i] * dlist[i]["PX"] * X["theta_dot"][i]
                    - dlist[i]["PX_sigma"]
                    * X["sigma_dot"][i]
                    * X["theta_dot"][i]
                )
                / (dlist[i]["K"] * H[i])
                for i in range(n)
            ]
        )
        blocks = []
        seed_values = dict(zip(X_KEYS, system.seed[:4]))
        for key in X_KEYS:
            residual = (
                system.cfg.delta_N * (system.D_N @ X[key] - rhs[key])
                / system.x_scales[key]
            )
            residual[0] = (X[key][0] - seed_values[key]) / system.x_scales[key]
            blocks.append(residual)
        rho_closure = (Q["rho"] - fresh_rho) / system.q_scales["rho"]
        return np.concatenate([*blocks, rho_closure])

    u0 = np.zeros(2 * (n - 1))
    residual0 = reduced_residual(u0)
    h = 1.0e-9
    jacobian = np.empty((len(residual0), len(u0)))
    for column in range(len(u0)):
        offset = np.zeros_like(u0)
        offset[column] = h
        jacobian[:, column] = (
            reduced_residual(offset) - reduced_residual(-offset)
        ) / (2.0 * h)

    objective = np.r_[np.zeros(len(u0)), 1.0]
    radius_scan = []
    candidates = []
    for radius in radii:
        scaled = radius * jacobian
        if objective_mode == "global_minimax":
            constraints = np.vstack(
                [
                    np.c_[scaled, -np.ones(len(residual0))],
                    np.c_[-scaled, -np.ones(len(residual0))],
                ]
            )
            limits = np.r_[-residual0, residual0]
        else:
            background = slice(0, 4 * n)
            rho = slice(4 * n, 5 * n)
            constraints = np.vstack(
                [
                    np.c_[scaled[background], np.zeros(4 * n)],
                    np.c_[-scaled[background], np.zeros(4 * n)],
                    np.c_[scaled[rho], -np.ones(n)],
                    np.c_[-scaled[rho], -np.ones(n)],
                ]
            )
            limits = np.r_[
                background_residual_cap - residual0[background],
                background_residual_cap + residual0[background],
                -residual0[rho],
                residual0[rho],
            ]
        optimum = linprog(
            objective,
            A_ub=constraints,
            b_ub=limits,
            bounds=[(-1.0, 1.0)] * len(u0) + [(0.0, None)],
            method="highs",
        )
        if not optimum.success:
            raise RuntimeError(optimum.message)
        step = radius * optimum.x[:-1]
        verified_residual = reduced_residual(step)
        verified = float(np.max(np.abs(verified_residual)))
        verified_background = float(
            np.max(np.abs(verified_residual[: 4 * n]))
        )
        verified_rho = float(np.max(np.abs(verified_residual[4 * n :])))
        item = {
            "trust_radius": radius,
            "predicted_minimax": float(optimum.x[-1]),
            "verified_reduced_max_abs": verified,
            "verified_background_max_abs": verified_background,
            "verified_rho_closure_max_abs": verified_rho,
            "verified_background_within_cap": bool(
                verified_background <= background_residual_cap
            ),
            "max_abs_normalized_velocity_step": float(np.max(np.abs(step))),
        }
        radius_scan.append(item)
        selection_key = (
            (0 if verified_background <= background_residual_cap else 1),
            verified_rho,
            verified,
        ) if objective_mode == "rho_lexicographic" else (verified,)
        candidates.append((selection_key, step, item))

    _selection_key, step, selected = min(candidates, key=lambda item: item[0])
    unprojected, total_density_error = vector_from_u(step)
    projected, projection = balanced_algebraic_projection(system, unprojected)
    diagnostics = [_physical_diagnostics(system, projected) for _ in range(3)]
    accepted = all(all(item["gate_pass"].values()) for item in diagnostics)
    payload = {
        "schema": "apeiron-ap1-r2c-p42-fixed-h-minimax-v1.0",
        "source": f"validated_p{n - 1}_PASS_only",
        "target_nodes": n,
        "baseline_dae_max_abs_scaled_residual": float(np.max(np.abs(baseline))),
        "finite_difference_step": h,
        "objective_mode": objective_mode,
        "background_residual_cap": background_residual_cap,
        "radius_scan": radius_scan,
        "selected_trust_radius": selected["trust_radius"],
        "max_abs_total_density_compensation_error": total_density_error,
        "algebraic_projection": projection,
        "repeated_full_diagnostics": diagnostics,
        "all_repeats_identical_dae": len(
            {item["dae_max_abs_scaled_residual"] for item in diagnostics}
        ) == 1,
        "accepted_as_checkpoint": accepted,
        "solution_vector": projected.tolist() if accepted else None,
        "trial_vector_stored_or_used_as_seed": accepted,
        "equations_changed": False,
        "physics_changed": False,
        "gates_changed": False,
        "production_curves_released": False,
    }
    if full_radius_diagnostics:
        full_scan = []
        for _radius_key, radius_step, radius_item in candidates:
            radius_unprojected, radius_density_error = vector_from_u(radius_step)
            radius_projected, radius_projection = balanced_algebraic_projection(
                system, radius_unprojected
            )
            radius_diagnostics = _physical_diagnostics(system, radius_projected)
            radius_pass = all(radius_diagnostics["gate_pass"].values())
            full_scan.append({
                **radius_item,
                "max_abs_total_density_compensation_error": radius_density_error,
                "algebraic_projection": radius_projection,
                "full_diagnostics": radius_diagnostics,
                "all_absolute_gates_pass": radius_pass,
                "solution_vector": radius_projected.tolist() if radius_pass else None,
            })
        payload["full_radius_diagnostics"] = full_scan
    return payload, projected


def main() -> None:
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--output", type=Path)
    args = parser.parse_args()
    system, candidate = reconstruct_candidate()
    payload, _projected = solve_fixed_h_minimax(system, candidate)
    rendered = json.dumps(payload, indent=2) + "\n"
    if args.output:
        args.output.write_text(rendered, encoding="utf-8")
    else:
        print(rendered, end="")


if __name__ == "__main__":
    main()