from __future__ import annotations

import json
from pathlib import Path
import tempfile
import unittest

from ap1_r2c_implicit_dae_multiresolution import (
    DAE_RESIDUAL_GATE, G22_THRESHOLDS, balanced_algebraic_projection,
    classify, degree_pair_bspline_lift, endpoint_bspline_lift,
    load_lambda1_checkpoint,
    richardson_bspline_lift,
)


class ImplicitDAEMultiresolutionTests(unittest.TestCase):
    def test_gates_remain_frozen(self):
        self.assertEqual(DAE_RESIDUAL_GATE, 1.0e-4)
        self.assertEqual(G22_THRESHOLDS, {
            "max_abs_friedmann_residual": 1.0e-6,
            "ward_normalized": 2.0e-6,
            "validation_source_relative_change": 1.0e-4,
            "validation_wronskian_relative_error": 1.0e-11,
        })

    def test_nonpass_level_keeps_g22_orange(self):
        levels = []
        for nodes in (17, 33, 65):
            gates = {"dae_residual": True, "structural": True}
            gates.update({key: True for key in G22_THRESHOLDS})
            diagnostics = {key: 1.0e-12 for key in G22_THRESHOLDS}
            diagnostics["gate_pass"] = gates
            levels.append({"candidate_nodes": nodes, "diagnostics": diagnostics})
        levels[-1]["diagnostics"]["gate_pass"]["dae_residual"] = False
        verdict = classify(levels)
        self.assertFalse(verdict["pass"])
        self.assertEqual(verdict["classification"], "G22_ORANGE_NOT_PASS")

    def test_lambda1_checkpoint_is_required(self):
        with tempfile.TemporaryDirectory() as folder:
            path = Path(folder) / "report.json"
            path.write_text(json.dumps({"stages": [{
                "lambda": 0.1, "converged": True,
                "solution_vector": [0.0] * 120,
            }]}), encoding="utf-8")
            with self.assertRaises(ValueError):
                load_lambda1_checkpoint(path)

    def test_endpoint_bspline_requires_single_p_step(self):
        class Config:
            candidate_nodes = 3
        class Source:
            cfg = Config()
        class TargetConfig:
            candidate_nodes = 5
        class Target:
            cfg = TargetConfig()
        with self.assertRaises(ValueError):
            endpoint_bspline_lift(Source(), [0.0], Target())

    def test_balanced_projection_finds_affine_minimax(self):
        class Config:
            candidate_nodes = 1
        class AffineSystem:
            cfg = Config()
            def evaluate(self, vector):
                # Seven one-node blocks plus junction. Pressure is exactly
                # algebraic; chi2 couples with the opposite sign into block 0.
                value = float(vector[6])
                residual = [value - 0.25, 0.0, 0.0, 0.0,
                            0.0, float(vector[5]) - 2.0,
                            value - 1.0, 0.0]
                return __import__("numpy").asarray(residual)
        vector = __import__("numpy").zeros(8)
        chosen, report = balanced_algebraic_projection(AffineSystem(), vector)
        self.assertGreaterEqual(report["selected_alpha"], 0.0)
        self.assertLessEqual(report["selected_alpha"], 1.0)
        self.assertLess(report["selected_norm"], 0.38)
        self.assertLess(report["affine_verification_error"], 1.0e-12)
        self.assertAlmostEqual(chosen[5], 2.0)

    def test_richardson_factor_is_bounded(self):
        with self.assertRaises(ValueError):
            richardson_bspline_lift(None, None, None, 6.0)

    def test_degree_pair_requires_ordered_bounded_degrees(self):
        with self.assertRaises(ValueError):
            degree_pair_bspline_lift(None, None, None, 4, 3, 1.0)
        with self.assertRaises(ValueError):
            degree_pair_bspline_lift(None, None, None, 3, 4, 5.1)


if __name__ == "__main__":
    unittest.main()