"""Module E: Closed-loop verification demo for HPI algorithm — NN edition.

Runs the full three-phase HPI pipeline and verifies the optimal
policy by simulating the closed-loop system from a perturbed initial state.
"""

import numpy as np
from scipy.integrate import solve_ivp

from .hpi_controller import HPIController


def run_demo():
    """Run full HPI demonstration and closed-loop verification.

    Returns:
        dict: All results from the HPI pipeline and verification.
    """
    print("=" * 60)
    print("  HPI: Homotopy-Based Policy Iteration (NN)")
    print("  Inverted Pendulum Adaptive Optimal Control")
    print("=" * 60)

    # ── Initialize controller ──
    controller = HPIController(J=1.0, mgl=1.0, R=1.0, pe_amplitude=0.5,
                               lr=1e-3, epochs=500, lambda_weight=1e-6)

    # ── Run three-phase HPI ──
    print("\n[1/3] Phase 1: Searching for admissible L0 ...")
    results = controller.run(
        x0=np.array([0.5, 0.0]),
        T=5.0,
        dt=0.05,
        L_start=0.1,
        L_step=0.5,
        L_max=20.0,
    )

    print(f"  L0 = {results['phase1_L0']:.2f}")
    print(f"  Phase 1 info: {results['phase1_info']}")
    print(f"  Positive definite: {controller.check_positive_definite()}")

    print(f"\n[2/3] Phase 2: Homotopy contraction ({len(results['phase2_history'])} iterations)")
    for i, (L, loss, pd) in enumerate(results["phase2_history"]):
        print(f"  iter {i}: L={L:.4f}, loss={loss:.4e}, PD={pd}")

    print(f"\n[3/3] Phase 3: Policy iteration (converged: {results['phase3_converged']})")
    print(f"  Iterations: {len(results['phase3_history'])}")

    # ── Closed-loop verification ──
    print("\n" + "=" * 60)
    print("  Closed-loop verification")
    print("=" * 60)

    x0_test = np.array([0.9, 0.1])

    def optimal_control(t, x):
        return controller.trainer.compute_u_hat(x)

    def closed_loop_dynamics(t, x):
        return controller.collector.dynamics(t, x, optimal_control)

    T_test = 5.0
    t_eval = np.linspace(0, T_test, 200)
    sol = solve_ivp(closed_loop_dynamics, (0, T_test), x0_test,
                    method="RK45", t_eval=t_eval, rtol=1e-8, atol=1e-10)

    X_cl = sol.y.T
    x_final = X_cl[-1]
    norm_final = np.linalg.norm(x_final)

    print(f"  Initial state:     {x0_test}")
    print(f"  Final state:       {np.array2string(x_final, precision=8)}")
    print(f"  ||x(T)||:          {norm_final:.2e}")

    if norm_final < 1e-3:
        print("  VERDICT: System converged to origin. [OK]")
    else:
        print(f"  VERDICT: Residual norm {norm_final:.2e} > 1e-3 [FAIL]")

    results["verification_x0"] = x0_test
    results["verification_x_final"] = x_final
    results["verification_norm"] = norm_final

    return results


if __name__ == "__main__":
    run_demo()