"""Bounded frozen-RHS Picard projection diagnostic for one AP1-R2c p level."""
from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np

from ap1_r2c_implicit_dae_multiresolution import (
    balanced_algebraic_projection,
    select_p_lift_projection_bank,
)
from ap1_r2c_implicit_dae_solver import ImplicitDAEResidual


def _accepted_parent(report: dict, nodes: int) -> np.ndarray:
    matches = [
        level for level in report["levels"]
        if level.get("accepted_as_checkpoint")
        and level.get("candidate_nodes") == nodes
        and level.get("solution_vector")
    ]
    if not matches:
        raise ValueError("validated immediate parent required")
    return np.asarray(matches[-1]["solution_vector"], dtype=float)


def _frozen_rhs_background_direction(system: ImplicitDAEResidual,
                                     vector: np.ndarray) -> np.ndarray:
    """Solve the four linear Lobatto defect equations at the current RHS."""
    n = system.cfg.candidate_nodes
    residual = system.evaluate(vector)
    matrix = system.cfg.delta_N * system.D_N.copy()
    matrix[0, :] = 0.0
    matrix[0, 0] = 1.0
    direction = np.zeros_like(vector)
    for block in range(4):
        defect = residual[block*n:(block+1)*n].copy()
        correction_rhs = -defect
        correction_rhs[0] = -defect[0]
        direction[block*n:(block+1)*n] = np.linalg.solve(
            matrix, correction_rhs
        )
    return direction


def scan(state: Path, checkpoint: Path, target_nodes: int) -> dict:
    report = json.loads(checkpoint.read_text(encoding="utf-8"))
    parent_nodes = target_nodes - 1
    parent_vector = _accepted_parent(report, parent_nodes)
    source = ImplicitDAEResidual(state, nodes=parent_nodes, uv_homotopy=1.0)
    target = ImplicitDAEResidual(state, nodes=target_nodes, uv_homotopy=1.0)
    base, lift_bank = select_p_lift_projection_bank(source, parent_vector, target)
    direction = _frozen_rhs_background_direction(target, base)
    n = target_nodes
    trials = []
    excluded_trials = []
    magnitudes = (1.0e-6, 3.0e-6, 1.0e-5, 3.0e-5,
                  5.0e-5, 1.0e-4, 3.0e-4)
    damping = (0.0, *magnitudes, *(-value for value in magnitudes))
    pairs = [(value, 0.0) for value in damping]
    pairs += [(0.0, value) for value in damping if value != 0.0]
    pairs += [(value, value) for value in damping if value != 0.0]
    for background_alpha, rho_alpha in pairs:
        background = base + background_alpha * direction
        residual = target.evaluate(background)
        rho_defect = residual[4*n:5*n]
        candidate = background.copy()
        candidate[4*n:5*n] -= rho_alpha * rho_defect
        try:
            projected, projection = balanced_algebraic_projection(
                target, candidate
            )
            final = target.evaluate(projected)
        except (ValueError, RuntimeError, FloatingPointError) as exc:
            excluded_trials.append({
                "background_alpha": background_alpha,
                "rho_alpha": rho_alpha,
                "reason": type(exc).__name__,
            })
            continue
        trials.append({
            "background_alpha": background_alpha,
            "rho_alpha": rho_alpha,
            "dae_max_abs_scaled_residual": float(np.max(np.abs(final))),
            "background_block": target.last_blocks["background"],
            "source_block": target.last_blocks["source"],
            "junction_block": target.last_blocks["junction"],
            "projection_alpha": projection["selected_alpha"],
        })
    trials.sort(key=lambda item: item["dae_max_abs_scaled_residual"])
    return {
        "target_nodes": target_nodes,
        "source": f"validated_{parent_nodes}_PASS_only",
        "base_lift": lift_bank["selected"],
        "method": "bounded_frozen_RHS_Lobatto_background_plus_rho_projection",
        "damping_bounds": [-3.0e-4, 3.0e-4],
        "trial_vectors_stored_or_used_as_seed": False,
        "excluded_trial_count": len(excluded_trials),
        "best": trials[:20],
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("state", type=Path)
    parser.add_argument("checkpoint", type=Path)
    parser.add_argument("--target-nodes", type=int, required=True)
    args = parser.parse_args()
    print(json.dumps(scan(args.state, args.checkpoint, args.target_nodes), indent=2))


if __name__ == "__main__":
    main()