"""
Demo: KK Tower Hamiltonian -- Genuine 11D Navigation

Uses a physically motivated Kaluza-Klein tower Hamiltonian with 64-level
Hilbert space to demonstrate genuine per-dimension navigation. Unlike the
toy 4-level system, this Hamiltonian encodes actual KK mass physics:

  E_d(n_d) = n_d^2 / (2 R_d^2)

Key improvements over the toy system:
  - 64 eigenvalues populate multiple spectral bands -> genuine curvature
  - Navigation coordinates are quantum observables: Tr(P_d @ rho) * R_d
  - Energy bands derived from physical KK mass thresholds
  - Coupling between modes creates non-trivial spectral structure

Answers the same 5 core questions as demo_11d_navigation.py, now with
physical KK structure backing every calculation.
"""

import sys
import os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

import numpy as np
from fce.kk_hamiltonian import KKTowerBuilder
from fce.engine import EngineConfig
from fce.navigation import (
    NavigationEngine, NavigationConfig, SpectralDimensionMapper,
    DimensionalScaling, build_dimension_info,
)
from fce.multi_clock import MultiClockAnalyzer, MultiClockConfig
from fce.frenet_serret import QuantumFrenetSerret


def main():
    print("=" * 70)
    print("  Fractal Correction Engine -- KK Tower 11D Navigation Demo")
    print("=" * 70)

    # ---- Build the KK Hamiltonian ----
    print("\n  Building KK tower Hamiltonian (dim_q=64, 7 compact dims)...")
    builder = KKTowerBuilder(dim_q=64, coupling=0.02)
    H = builder.build_hamiltonian()
    dim_q = H.shape[0]

    # Print spectrum summary
    summary = builder.spectrum_summary()
    print(f"\n  Hilbert space dimension:  {summary['dim_q']}")
    print(f"  Energy range:             [{summary['E_min']:.4f}, {summary['E_max']:.4f}]")
    print(f"  Off-diagonal elements:    {summary['n_off_diagonal']}")

    print(f"\n  Compactification radii and first-mode energies:")
    print(f"  {'Dim':>4s}  {'Radius (nat.)':>14s}  {'E1 (1st mode)':>14s}")
    print("  " + "-" * 36)
    for d in range(7):
        print(f"  {d+4:>4d}  {builder.radii[d]:>14.6f}  "
              f"{builder.first_mode_energies[d]:>14.4f}")

    print(f"\n  Band populations (eigenvalues per band):")
    print(f"  {'Band':>5s}  {'Type':>10s}  {'Count':>6s}")
    print("  " + "-" * 25)
    band_labels = ['time', 'x', 'y', 'z'] + [f'compact_{d}' for d in range(4, 11)]
    for b, count in enumerate(summary['band_populations']):
        btype = 'spacetime' if b < 4 else 'compact'
        print(f"  {b:>5d}  {btype:>10s}  {count:>6d}")

    # ---- Initial state ----
    print("\n  Initial state: superposition of 8 lowest eigenstates")
    rho_init = builder.build_initial_state('superposition')
    purity = float(np.real(np.trace(rho_init @ rho_init)))
    E_init = float(np.real(np.trace(H @ rho_init)))
    print(f"  Purity:  {purity:.6f}")
    print(f"  Energy:  {E_init:.4f}")

    # Initial KK momenta
    P_ops = builder.build_momentum_operators()
    print(f"\n  Initial KK momentum observables:")
    print(f"  {'Dim':>4s}  {'<P_d>':>12s}  {'coord (R*<P>)':>14s}")
    print("  " + "-" * 34)
    for d in range(7):
        p_exp = float(np.real(np.trace(P_ops[d] @ rho_init)))
        coord = p_exp * builder.radii[d]
        print(f"  {d+4:>4d}  {p_exp:>12.6f}  {coord:>14.6e}")

    # ---- System parameters ----
    system_params = {'T1': 1e-3, 'T2': 5e-4, 'gate_fidelity': 0.999}

    engine_config = EngineConfig(
        feedback_gain=0.1,
        correction_threshold=0.001,
        qec_enabled=True,
        enable_thermal_monitoring=False,
    )

    nav_config = NavigationConfig(
        kk_activation_sharpness=10.0,
        enable_boundary_effects=True,
        navigation_temperature=1.0,
    )

    dt = 1e-6
    n_steps = 100

    # ================================================================
    # QUESTION 1: KK Activation Thresholds
    # ================================================================
    print("\n" + "=" * 70)
    print("  Q1: KK Activation Thresholds (from physical radii)")
    print("=" * 70)

    scaler = DimensionalScaling()
    dim_info = build_dimension_info()
    mapper = SpectralDimensionMapper(H, scaler, dim_info, kk_builder=builder)

    print(f"\n  System eigenvalue range: [{mapper.E_min:.4f}, {mapper.E_max:.4f}]")
    print()
    print(f"  {'Dim':>4s}  {'Name':>12s}  {'Type':>10s}  "
          f"{'Threshold':>12s}  {'Active?':>8s}")
    print("  " + "-" * 52)

    for d in range(11):
        di = dim_info[d]
        thresh = mapper.activation_energies[d]
        act = mapper.activation_function(E_init, d)
        active_str = "YES" if act > 0.5 else f"{act:.3f}"
        thresh_str = f"{thresh:.4f}" if np.isfinite(thresh) else "-inf"
        print(f"  {d:>4d}  {di.name:>12s}  {di.dim_type:>10s}  "
              f"{thresh_str:>12s}  {active_str:>8s}")

    # ================================================================
    # QUESTION 2: Per-Dimension Curvature
    # ================================================================
    print("\n" + "=" * 70)
    print("  Q2: Per-Dimension Curvature (Spectral Decomposition)")
    print("=" * 70)

    fs = QuantumFrenetSerret()
    dim_curv = mapper.compute_dimensional_curvature(
        rho_init, fs, E_init, sharpness=nav_config.kk_activation_sharpness
    )

    total_kappa = fs.compute_curvature(rho_init, H)
    total_tau = fs.compute_torsion(rho_init, H)

    print(f"\n  Total scalar curvature:  kappa = {total_kappa:.6f}")
    print(f"  Total scalar torsion:    tau   = {total_tau:.6f}")
    print()

    genuine_count = 0
    print(f"  {'Dim':>4s}  {'kappa_d':>10s}  {'tau_d':>10s}  "
          f"{'activation':>12s}  {'weight':>10s}  {'source':>14s}")
    print("  " + "-" * 68)

    for d in range(11):
        k = dim_curv.curvatures[d]
        t = dim_curv.torsions[d]
        a = dim_curv.activation[d]
        w = dim_curv.subspace_weights[d]
        source = "spectral" if w > 1e-12 else "extrapolated"
        if source == "spectral":
            genuine_count += 1
        print(f"  {d:>4d}  {k:>10.6f}  {t:>10.6f}  "
              f"{a:>12.6f}  {w:>10.6f}  {source:>14s}")

    print(f"\n  Genuine (spectral) curvature in {genuine_count}/11 dimensions")

    # ================================================================
    # QUESTION 3: Multi-Clock Dimensional Resolution
    # ================================================================
    print("\n" + "=" * 70)
    print("  Q3: Multi-Clock KK Probing (64-dim projectors)")
    print("=" * 70)

    mc_config = MultiClockConfig(n_clocks=5, offset_multiplier=1.0)
    mc = MultiClockAnalyzer(
        hamiltonian=H, system_params=system_params, config=mc_config
    )

    kk_result = mc.evolve_kk_probing_clocks(
        rho_init, dt, min(n_steps, 50), mapper.projectors
    )

    print(f"\n  Occupied energy bands: {kk_result['occupied_bands']}")
    print()
    print(f"  {'Band':>5s}  {'Divergence':>12s}  {'Sensitivity':>12s}  {'Status':>10s}")
    print("  " + "-" * 45)

    for b in range(min(11, len(kk_result['divergence_per_band']))):
        div = kk_result['divergence_per_band'][b]
        sens = kk_result['sensitivity_per_dim'][b]
        status = "ACTIVE" if b in kk_result['occupied_bands'] else "empty"
        print(f"  {b:>5d}  {div:>12.6f}  {sens:>12.6f}  {status:>10s}")

    n_resolved = len(kk_result['occupied_bands'])
    print(f"\n  Dimensions resolved by multi-clock: {n_resolved}/11")

    # ================================================================
    # QUESTION 4: QFT Boundary Effects
    # ================================================================
    print("\n" + "=" * 70)
    print("  Q4: QFT Boundary Effects")
    print("=" * 70)

    # Run navigation
    print("\n  Running 11D navigation (100 steps, dim_q=64)...")
    nav = NavigationEngine(
        hamiltonian=H,
        system_params=system_params,
        engine_config=engine_config,
        nav_config=nav_config,
        kk_builder=builder,
    )
    nav_result = nav.navigate(rho_init, dt, n_steps)

    print(f"  Navigation complete. Fidelity = {nav_result.total_fidelity:.6f}")

    print()
    if nav_result.qft_boundary_effects:
        print(f"  {'Dim':>4s}  {'Casimir (J)':>14s}  {'Hawking T':>14s}  "
              f"{'Radius (m)':>14s}  {'KK mass (GeV)':>14s}")
        print("  " + "-" * 66)
        for d, effects in sorted(nav_result.qft_boundary_effects.items()):
            print(f"  {d:>4d}  {effects['casimir_energy']:>14.4e}  "
                  f"{effects['hawking_temperature']:>14.4e}  "
                  f"{effects['compactification_radius']:>14.4e}  "
                  f"{effects['kk_mass_GeV']:>14.4e}")
    else:
        print("  No boundary transitions. Static boundary analysis:")
        print()
        print(f"  {'Dim':>4s}  {'Casimir (J)':>14s}  {'Hawking T':>14s}  "
              f"{'Radius (m)':>14s}  {'KK mass (GeV)':>14s}")
        print("  " + "-" * 66)
        for d in range(4, 11):
            effects = nav._qft_boundary_analysis(d)
            print(f"  {d:>4d}  {effects['casimir_energy']:>14.4e}  "
                  f"{effects['hawking_temperature']:>14.4e}  "
                  f"{effects['compactification_radius']:>14.4e}  "
                  f"{effects['kk_mass_GeV']:>14.4e}")

    # ================================================================
    # QUESTION 5: Thermodynamic Cost
    # ================================================================
    print("\n" + "=" * 70)
    print("  Q5: Thermodynamic Navigation Cost")
    print("=" * 70)

    cost = nav_result.navigation_cost
    if cost is not None:
        print(f"\n  Free energy change (Jarzynski): {cost.free_energy_change:.6e}")
        print(f"  Mean work:                      {np.mean(cost.work_values):.6e}")
        print(f"  Dissipated work:                {cost.dissipated_work:.6e}")
        print(f"  Entropy production:             {cost.entropy_production:.6e}")
        print(f"  Landauer cost:                  {cost.landauer_cost:.6e}")

        print("\n  Physics consistency:")
        print(f"    Entropy production >= 0:  "
              f"{'PASS' if cost.entropy_production >= -1e-15 else 'FAIL'}")

        if cost.crooks_verification is not None:
            print(f"\n  Crooks fluctuation theorem:")
            for k, v in cost.crooks_verification.items():
                print(f"    {k}: {v:.6e}")

    # ================================================================
    # Observable coordinate dynamics
    # ================================================================
    print("\n" + "=" * 70)
    print("  Observable Coordinate Dynamics (Tr(P_d @ rho) * R_d)")
    print("=" * 70)

    # Show coordinates at 5 time snapshots
    snapshots = [0, n_steps // 4, n_steps // 2, 3 * n_steps // 4, n_steps]
    print()
    header = f"  {'Dim':>4s}"
    for s in snapshots:
        header += f"  {'t=' + str(s):>14s}"
    print(header)
    print("  " + "-" * (4 + 16 * len(snapshots)))

    for d in range(11):
        row = f"  {d:>4d}"
        for s in snapshots:
            c = nav_result.states[s].coordinates[d]
            row += f"  {c:>14.6e}"
        print(row)

    # Count dimensions with actual coordinate evolution
    initial_c = nav_result.states[0].coordinates
    final_c = nav_result.states[-1].coordinates
    active_dims = sum(
        1 for d in range(11)
        if abs(final_c[d] - initial_c[d]) > 1e-15
    )

    # ================================================================
    # Holographic & Chaos
    # ================================================================
    print("\n" + "=" * 70)
    print("  Holographic Metrics")
    print("=" * 70)

    holo = nav_result.holographic_metrics
    if holo:
        print(f"\n  BH entropy:             {holo.get('bh_entropy', 0):.6e}")
        print(f"  Scrambling time:        {holo.get('scrambling_time', 0):.6e}")
        print(f"  RT entanglement:        {holo.get('rt_entanglement_entropy', 0):.6e}")

    print("\n" + "=" * 70)
    print("  Chaos Analysis")
    print("=" * 70)

    chaos = nav_result.chaos_metrics
    if chaos:
        print(f"\n  Max Lyapunov:           {chaos.get('lyapunov_max', 0):.6e}")
        print(f"  Is chaotic:             {chaos.get('is_chaotic', False)}")
        print(f"  MSS bound satisfied:    {chaos.get('mss_bound_satisfied', True)}")

    # ================================================================
    # Summary
    # ================================================================
    print("\n" + "=" * 70)
    print("  Navigation Summary")
    print("=" * 70)

    print(f"\n  Hilbert space:              {dim_q} levels (vs 4 in toy system)")
    print(f"  FCE fidelity:               {nav_result.total_fidelity:.6f}")
    print(f"  Active dimensions:          {nav_result.transition_state.active_dimensions}")
    print(f"  Dims with coord evolution:  {active_dims}/11")
    print(f"  Genuine spectral curvature: {genuine_count}/11 dimensions")
    print(f"  Hausdorff dimension:        "
          f"{nav_result.evolution_result.hausdorff_dimension:.4f}")
    print(f"  Trajectory health:          "
          f"{nav_result.evolution_result.trajectory_health}")

    # KK coordinate comparison: initial vs final
    print(f"\n  KK momentum dynamics (compact dims):")
    print(f"  {'Dim':>4s}  {'Initial <P>':>14s}  {'Final <P>':>14s}  "
          f"{'Delta':>14s}  {'Active?':>8s}")
    print("  " + "-" * 60)

    for d in range(7):
        p_init = float(np.real(np.trace(
            P_ops[d] @ nav_result.states[0].quantum_state
        )))
        p_final = float(np.real(np.trace(
            P_ops[d] @ nav_result.states[-1].quantum_state
        )))
        delta = p_final - p_init
        active = "YES" if abs(delta) > 1e-10 else "no"
        print(f"  {d+4:>4d}  {p_init:>14.6e}  {p_final:>14.6e}  "
              f"{delta:>14.6e}  {active:>8s}")

    print("\n" + "=" * 70)
    print("  KK Tower Navigation Complete")
    print("  All 5 core questions answered with physical KK structure.")
    print("=" * 70)


if __name__ == '__main__':
    main()
