"""
test_optimize.py -- End-to-End Optimization Tests

Tests the full pipeline: Finsler metric → QUBO → QAOA → Classical comparison.

Test Cases:
    1. "E = mc^2"       → expected Φ, approximation_ratio > 0.95
    2. "a^2 + b^2 = c^2" → expected Σ, approximation_ratio > 0.90
    3. "∀x. P(x) → Q(x)" → expected Λ, approximation_ratio > 0.90

Each test:
    1. Builds the Finsler metric on the Hachimoji simplex
    2. Encodes as QUBO with equation-specific bias
    3. Solves with QAOA (statevector simulation)
    4. Solves with classical methods (HiGHS + SA)
    5. Compares results and checks approximation ratio
"""

from __future__ import annotations

import math
import sys
import time
from typing import Any

import numpy as np

# Add stage4-optimize to path
sys.path.insert(0, "/mnt/agents/output/rebuild/stage4-optimize")

from finsler_metric import (
    make_uniform_hachimoji_states,
    compute_finsler_distance_matrix,
    build_full_finsler_pipeline,
)
from qubo_builder import (
    QUBO,
    build_equation_qubo,
    brute_force_qubo,
    extract_dominant_state,
    finsler_to_qubo,
    qubo_to_ising,
    ising_to_pauli,
)
from qaoa_circuit import qaoa_solve
from classical_solver import solve_classical, compare_solvers


# =========================================================================
# Test Configuration
# =========================================================================

TEST_CASES: list[dict] = [
    {
        "name": "E = mc^2",
        "equation": "E = mc^2",
        "expected_state": "\u03a6",  # Phi: energy-mass equivalence
        "min_approx_ratio": 0.95,
    },
    {
        "name": "a^2 + b^2 = c^2",
        "equation": "a^2 + b^2 = c^2",
        "expected_state": "\u03a3",  # Sigma: Pythagorean symmetric
        "min_approx_ratio": 0.90,
    },
    {
        "name": "\u2200x. P(x) \u2192 Q(x)",
        "equation": "\u2200x. P(x) \u2192 Q(x)",
        "expected_state": "\u039b",  # Lambda: universal implication
        "min_approx_ratio": 0.90,
    },
]


def run_single_test(test_case: dict, states: list, verbose: bool = True) -> dict:
    """Run a single test case through the full pipeline.

    Pipeline:
        Equation → QUBO → QAOA + Classical → Comparison
    """
    name = test_case["name"]
    equation = test_case["equation"]
    expected = test_case["expected_state"]
    min_ratio = test_case["min_approx_ratio"]

    if verbose:
        print(f"\n{'='*60}")
        print(f"TEST: {name}")
        print(f"Equation: {equation}")
        print(f"Expected state: {expected}")
        print(f"{'='*60}")

    # Step 1: Build QUBO with equation-specific bias
    t0 = time.time()
    qubo, target = build_equation_qubo(equation, states)
    qubo_time = time.time() - t0

    if verbose:
        print(f"\n[1] QUBO built in {qubo_time:.4f}s")
        print(f"    n={qubo.n}, terms={len(qubo.matrix)}, target={target}")

    # Step 2: Brute force ground truth (n=8 → 256 states)
    t0 = time.time()
    bf = brute_force_qubo(qubo)
    ground_energy = bf["energy"]
    ground_state = extract_dominant_state(bf["solution"])
    bf_time = time.time() - t0

    if verbose:
        print(f"\n[2] Ground truth (brute force) in {bf_time:.4f}s")
        print(f"    Ground energy: {ground_energy:.6f}")
        print(f"    Ground state:  {ground_state}")

    # Step 3: QAOA solve
    t0 = time.time()
    qaoa_result = qaoa_solve(qubo, p=2, shots=2048, optimize_params=True)
    qaoa_time = time.time() - t0

    qaoa_state = qaoa_result["optimal_state"]
    qaoa_energy = qaoa_result["energy"]
    qaoa_ratio = qaoa_result["approximation_ratio"]

    if verbose:
        print(f"\n[3] QAOA solve in {qaoa_time:.4f}s")
        print(f"    QAOA state:  {qaoa_state}")
        print(f"    QAOA energy: {qaoa_energy:.6f}")
        print(f"    Approx ratio: {qaoa_ratio:.4f}")
        print(f"    Circuit depth: {qaoa_result['circuit_depth']}")

    # Step 4: Classical solvers
    t0 = time.time()
    classical = compare_solvers(qubo, time_limit=1.0)
    classical_time = time.time() - t0

    highs_state = classical["highs"]["optimal_state"]
    highs_energy = classical["highs"]["energy"]
    sa_state = classical["sa"]["optimal_state"]
    sa_energy = classical["sa"]["energy"]

    if verbose:
        print(f"\n[4] Classical solvers in {classical_time:.4f}s")
        print(f"    HiGHS: state={highs_state}, energy={highs_energy:.6f}")
        print(f"    SA:    state={sa_state}, energy={sa_energy:.6f}")

    # Step 5: Check results
    qaoa_matches = qaoa_state == expected
    classical_matches = (highs_state == expected) or (sa_state == expected)
    all_agree = len(set([qaoa_state, highs_state, sa_state])) == 1
    ratio_ok = qaoa_ratio >= min_ratio

    passed = qaoa_matches and classical_matches and ratio_ok

    if verbose:
        print(f"\n[5] Results:")
        print(f"    QAOA matches expected:   {qaoa_matches} ({qaoa_state} == {expected})")
        print(f"    Classical matches:       {classical_matches}")
        print(f"    All solvers agree:       {all_agree}")
        print(f"    Approx ratio >= {min_ratio}:  {ratio_ok} ({qaoa_ratio:.4f})")
        print(f"    PASSED: {passed}")

    return {
        "name": name,
        "equation": equation,
        "expected": expected,
        "qaoa_state": qaoa_state,
        "qaoa_energy": qaoa_energy,
        "qaoa_ratio": qaoa_ratio,
        "highs_state": highs_state,
        "highs_energy": highs_energy,
        "sa_state": sa_state,
        "sa_energy": sa_energy,
        "ground_state": ground_state,
        "ground_energy": ground_energy,
        "qaoa_matches": qaoa_matches,
        "classical_matches": classical_matches,
        "all_agree": all_agree,
        "ratio_ok": ratio_ok,
        "passed": passed,
    }


def run_all_tests(verbose: bool = True) -> dict:
    """Run all 3 test cases and return summary."""
    print("="*60)
    print("Finsler → QUBO → QAOA Optimizer: End-to-End Tests")
    print("="*60)
    print(f"\nBuilding Hachimoji 8-state system...")

    t0 = time.time()
    states = make_uniform_hachimoji_states()
    finsler_pipeline = build_full_finsler_pipeline(states)
    setup_time = time.time() - t0

    print(f"Setup complete in {setup_time:.4f}s")
    print(f"States: {[s.symbol for s in states]}")
    print(f"Finsler anisotropic: {finsler_pipeline['is_anisotropic']}")

    results = []
    all_passed = True

    for tc in TEST_CASES:
        result = run_single_test(tc, states, verbose=verbose)
        results.append(result)
        if not result["passed"]:
            all_passed = False

    # Summary
    if verbose:
        print(f"\n{'='*60}")
        print("SUMMARY")
        print(f"{'='*60}")
        for r in results:
            status = "PASS" if r["passed"] else "FAIL"
            print(f"  [{status}] {r['name']:30s} → "
                  f"QAOA:{r['qaoa_state']} (ratio={r['qaoa_ratio']:.4f}) | "
                  f"HiGHS:{r['highs_state']} | SA:{r['sa_state']} | "
                  f"Expected:{r['expected']}")
        print(f"\nOverall: {'ALL PASSED' if all_passed else 'SOME FAILED'}")

    return {
        "all_passed": all_passed,
        "results": results,
        "states": [s.to_dict() for s in states],
        "finsler": finsler_pipeline,
    }


def test_finsler_metric_properties(verbose: bool = True) -> dict:
    """Test mathematical properties of the Finsler metric.

    Verifies:
        1. α is symmetric: α(a,b) = α(b,a)
        2. β is antisymmetric: β(a,b) = -β(b,a)
        3. F is positive: F(a,b) > 0
        4. Phase distances respect circular topology
    """
    from finsler_metric import (
        compute_alpha_component,
        compute_beta_component,
        compute_finsler_metric,
        phase_distance_s1,
    )

    states = make_uniform_hachimoji_states()

    checks = []

    # Check 1: α symmetry
    alpha_sym_ok = True
    for i in range(len(states)):
        for j in range(i + 1, len(states)):
            a_ij = compute_alpha_component(states[i], states[j])
            a_ji = compute_alpha_component(states[j], states[i])
            if abs(a_ij - a_ji) > 1e-9:
                alpha_sym_ok = False
                break
    checks.append(("α symmetry", alpha_sym_ok))

    # Check 2: β antisymmetry
    beta_antisym_ok = True
    for i in range(len(states)):
        for j in range(i + 1, len(states)):
            b_ij = compute_beta_component(states[i], states[j])
            b_ji = compute_beta_component(states[j], states[i])
            if abs(b_ij + b_ji) > 1e-9:
                beta_antisym_ok = False
                break
    checks.append(("β antisymmetry", beta_antisym_ok))

    # Check 3: F positivity
    f_pos_ok = True
    for i in range(len(states)):
        for j in range(len(states)):
            if i != j:
                F = compute_finsler_metric(states[i], states[j])
                if F <= 0:
                    f_pos_ok = False
                    break
    checks.append(("F positivity", f_pos_ok))

    # Check 4: Phase circular distance
    phase_ok = True
    d_0_180 = phase_distance_s1(0, 180)    # π (half circle)
    d_0_90 = phase_distance_s1(0, 90)      # π/2 (quarter circle)
    d_0_270 = phase_distance_s1(0, 270)    # π/2 (shortest arc is 90°)
    d_0_360 = phase_distance_s1(0, 360)    # 0 (same point on S¹)
    # 0→90 should equal 0→270 (both are π/2: shortest arc)
    if abs(d_0_90 - d_0_270) > 1e-9:
        phase_ok = False
    # 0→180 should be π
    if abs(d_0_180 - math.pi) > 1e-9:
        phase_ok = False
    # 0→360 should be 0 (same point)
    if d_0_360 > 1e-9:
        phase_ok = False
    # 0→90 should be π/2
    if abs(d_0_90 - math.pi / 2) > 1e-9:
        phase_ok = False
    checks.append(("Phase circular distance", phase_ok))

    if verbose:
        print(f"\nFinsler Metric Property Checks:")
        for name, ok in checks:
            print(f"  [{'PASS' if ok else 'FAIL'}] {name}")

    all_ok = all(ok for _, ok in checks)
    return {"all_ok": all_ok, "checks": checks}


# =========================================================================
# Main
# =========================================================================

if __name__ == "__main__":
    print("Stage 4: Finsler → QUBO → QAOA Optimizer")
    print("="*60)

    # First test the Finsler metric properties
    print("\n--- Mathematical Property Tests ---")
    prop_result = test_finsler_metric_properties(verbose=True)

    # Run main end-to-end tests
    print("\n--- End-to-End Optimization Tests ---")
    test_result = run_all_tests(verbose=True)

    # Final report
    print(f"\n{'='*60}")
    print("FINAL REPORT")
    print(f"{'='*60}")
    print(f"Finsler properties: {'ALL PASS' if prop_result['all_ok'] else 'SOME FAIL'}")
    print(f"End-to-end tests:   {'ALL PASS' if test_result['all_passed'] else 'SOME FAIL'}")

    if test_result["all_passed"] and prop_result["all_ok"]:
        print(f"\n✓ Stage 4: ALL TESTS PASSED")
    else:
        print(f"\n✗ Stage 4: SOME TESTS FAILED")
        sys.exit(1)