#!/usr/bin/env python3
"""
5-Clock Divergence Navigation Demo.

Demonstrates the core theory: use 5 offset clocks sampling adjacent
timelines to track how a quantum state drifts through KK-compactified
extra dimensions, then compute corrections to navigate back to origin.

Sections:
    1. Clock Divergence Verification
    2. Per-Dimension KK Displacement
    3. Clock Velocity Estimates
    4. Return-to-Origin Correction
    5. Honest Assessment

Usage:
    python examples/demo_clock_navigation.py
"""

import sys
import os
import numpy as np

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

from fce.multi_clock import (
    MultiClockAnalyzer, MultiClockConfig,
    compute_return_unitary, estimate_return_displacement,
)
from fce.kk_hamiltonian import KKTowerBuilder
from fce.fidelity import uhlmann_fidelity


def separator(title: str):
    print(f"\n{'=' * 70}")
    print(f"  {title}")
    print(f"{'=' * 70}\n")


def main():
    print("=" * 70)
    print("  5-Clock Divergence Navigation -- Return to Origin Demo")
    print("=" * 70)

    # Build KK system
    builder = KKTowerBuilder(dim_q=32)
    H = builder.build_hamiltonian()
    rho_init = builder.build_initial_state('superposition')
    params = {'T1': 1e-3, 'T2': 5e-4, 'gate_fidelity': 0.999}
    dt = 1e-6
    n_steps = 50

    P_ops = builder.build_momentum_operators()
    radii = builder.radii
    n_dims = len(P_ops)

    # Spectrum info
    eigenvalues = np.linalg.eigvalsh(H)
    E_range = eigenvalues[-1] - eigenvalues[0]

    print(f"\n  System: KK tower on T^7, dim_q = {builder.dim_q}")
    print(f"  Energy range: [{eigenvalues[0]:.3f}, {eigenvalues[-1]:.3f}]"
          f"  (E_range = {E_range:.3f})")
    print(f"  Radii: {', '.join(f'{r:.4f}' for r in radii)}")
    print(f"  dt = {dt:.1e}, n_steps = {n_steps}")
    print(f"  Total evolution time: {dt * n_steps:.1e}")

    # ----------------------------------------------------------------
    # 1. Clock Divergence Verification
    # ----------------------------------------------------------------
    separator("1. Clock Divergence Verification")

    # Show the problem with default multiplier
    default_multiplier = 1.0
    auto_multiplier = MultiClockAnalyzer.compute_offset_multiplier(H, dt)

    print(f"  Default offset_multiplier: {default_multiplier}")
    print(f"  Default phase rotation:    "
          f"{E_range * dt * default_multiplier:.2e} rad")
    print()
    print(f"  Auto-scaled multiplier:    {auto_multiplier:.1f}")
    print(f"  Auto phase rotation:       "
          f"{E_range * dt * auto_multiplier:.4f} rad (~0.1 rad target)")

    # Run with auto-scaling
    analyzer = MultiClockAnalyzer(H, system_params=params)
    result = analyzer.evolve_clocks_with_displacement(
        rho_init, dt, n_steps, builder, target_rotation=0.1
    )

    print(f"\n  Clock Fidelities to Central (final step):")
    print(f"  {'Clock':>8}  {'Offset':>8}  {'Fidelity':>10}")
    print(f"  {'-' * 30}")
    offsets = result.clock_result.offsets
    for k in range(len(result.clock_fidelities)):
        marker = " <-- central" if k == result.clock_result.central_index else ""
        print(f"  {k:>8d}  {offsets[k]:>+8.0f}  "
              f"{result.clock_fidelities[k]:>10.6f}{marker}")

    outer_fid = min(result.clock_fidelities[0], result.clock_fidelities[-1])
    print(f"\n  Outer clock fidelity: {outer_fid:.6f}")
    if outer_fid < 0.99:
        print("  --> Clocks show VISIBLE divergence (theory requirement met)")
    else:
        print("  --> Clocks still too similar (increase target_rotation)")

    # ----------------------------------------------------------------
    # 2. Per-Dimension KK Displacement
    # ----------------------------------------------------------------
    separator("2. Per-Dimension KK Displacement")

    print(f"  {'Dim':>4}  {'R_d':>8}  {'x_d(0)':>12}  {'x_d(t_f)':>12}  "
          f"{'delta_d':>12}  {'|delta|':>10}")
    print(f"  {'-' * 62}")

    for d in range(n_dims):
        dim_label = f"d={d+4}"
        print(f"  {dim_label:>4}  {radii[d]:>8.4f}  "
              f"{result.initial_coordinates[d]:>12.6e}  "
              f"{result.central_coordinates[d]:>12.6e}  "
              f"{result.displacements[d]:>12.6e}  "
              f"{abs(result.displacements[d]):>10.6e}")

    total_disp = np.linalg.norm(result.displacements)
    print(f"\n  Total displacement ||delta|| = {total_disp:.6e}")

    # Identify most-displaced dimension
    max_dim = np.argmax(np.abs(result.displacements))
    print(f"  Most displaced: dimension {max_dim + 4} "
          f"(|delta| = {abs(result.displacements[max_dim]):.6e})")

    # ----------------------------------------------------------------
    # 3. Clock Velocity Estimates
    # ----------------------------------------------------------------
    separator("3. Clock Velocity Estimates (dx/dt from finite differences)")

    offset_dt = dt * auto_multiplier

    print(f"  Offset dt between adjacent clocks: {offset_dt:.6e}")
    print()
    print(f"  {'Dim':>4}  {'v_d (dx/dt)':>14}  {'Direction':>10}")
    print(f"  {'-' * 32}")

    for d in range(n_dims):
        v = result.clock_velocity_estimates[d]
        direction = "+" if v > 0 else ("-" if v < 0 else "0")
        print(f"  d={d+4:>2}  {v:>14.6e}  {direction:>10}")

    # ----------------------------------------------------------------
    # 4. Return-to-Origin Correction
    # ----------------------------------------------------------------
    separator("4a. Gain Calibration")

    central_idx = result.clock_result.central_index
    rho_final = result.clock_result.clocks[central_idx].density_matrices[-1]
    before_norm = np.linalg.norm(result.displacements)

    # Find optimal gain by scanning
    print("  The position operator X_d generates translations in momentum")
    print("  space. The gain parameter scales the correction strength.")
    print("  Too small = undershoot, too large = overshoot.")
    print()
    print(f"  {'Gain':>6}  {'Residual':>12}  {'Reduction':>10}")
    print(f"  {'-' * 34}")

    best_gain = 1.0
    best_reduction = 0.0

    for gain in [1, 5, 10, 15, 20, 25, 30, 40, 50]:
        ret = estimate_return_displacement(result)
        U = compute_return_unitary(ret, builder, gain=float(gain))
        rho_c = U @ rho_final @ U.conj().T
        c_coords = np.array([
            np.real(np.trace(P_ops[d] @ rho_c)) * radii[d]
            for d in range(n_dims)
        ])
        c_disp = np.linalg.norm(c_coords - result.initial_coordinates)
        pct = (1 - c_disp / before_norm) * 100 if before_norm > 1e-15 else 0
        marker = ""
        if pct > best_reduction:
            best_reduction = pct
            best_gain = float(gain)
            marker = "  <-- best"
        print(f"  {gain:>6}  {c_disp:>12.6e}  {pct:>+8.1f}%{marker}")

    print(f"\n  Optimal gain: {best_gain:.0f} ({best_reduction:+.1f}% reduction)")

    separator("4b. Single-Step Return Correction")

    ret_disp = estimate_return_displacement(result)
    U_return = compute_return_unitary(ret_disp, builder, gain=best_gain)

    # Verify unitarity
    identity = np.eye(U_return.shape[0], dtype=complex)
    unitarity_error = np.linalg.norm(U_return @ U_return.conj().T - identity)
    print(f"  Return unitary shape: {U_return.shape}")
    print(f"  Unitarity check ||UU† - I||: {unitarity_error:.2e}")
    print(f"  Gain used: {best_gain:.0f}")

    # Apply correction
    rho_corrected = U_return @ rho_final @ U_return.conj().T

    corrected_coords = np.array([
        np.real(np.trace(P_ops[d] @ rho_corrected)) * radii[d]
        for d in range(n_dims)
    ])
    corrected_disp = corrected_coords - result.initial_coordinates

    print(f"\n  {'Dim':>4}  {'Before':>12}  {'After':>12}  {'Reduction':>10}")
    print(f"  {'-' * 50}")
    for d in range(n_dims):
        before = abs(result.displacements[d])
        after = abs(corrected_disp[d])
        if before > 1e-15:
            pct = (1 - after / before) * 100
            reduction = f"{pct:>+8.1f}%"
        else:
            reduction = "    n/a"
        print(f"  d={d+4:>2}  {before:>12.6e}  {after:>12.6e}  {reduction}")

    after_norm = np.linalg.norm(corrected_disp)
    pct_total = (1 - after_norm / before_norm) * 100 if before_norm > 1e-15 else 0

    print(f"\n  Total displacement: {before_norm:.6e} -> {after_norm:.6e}"
          f"  ({pct_total:+.1f}%)")

    separator("4c. Iterative Correction (feedback loop)")

    print("  Apply correction repeatedly, re-measuring displacement each time.")
    print("  This mimics continuous feedback: measure -> correct -> measure -> ...")
    print()
    print(f"  {'Iter':>5}  {'Residual':>12}  {'Cumulative':>12}  {'Reduction':>10}")
    print(f"  {'-' * 46}")
    print(f"  {'init':>5}  {before_norm:>12.6e}  {'---':>12}  {'---':>10}")

    rho_curr = rho_final.copy()
    iter_gain = 10.0
    for iteration in range(8):
        curr_coords = np.array([
            np.real(np.trace(P_ops[d] @ rho_curr)) * radii[d]
            for d in range(n_dims)
        ])
        curr_disp = curr_coords - result.initial_coordinates
        ret = -curr_disp
        U = compute_return_unitary(ret, builder, gain=iter_gain)
        rho_curr = U @ rho_curr @ U.conj().T

        new_coords = np.array([
            np.real(np.trace(P_ops[d] @ rho_curr)) * radii[d]
            for d in range(n_dims)
        ])
        new_disp_norm = np.linalg.norm(new_coords - result.initial_coordinates)
        pct = (1 - new_disp_norm / before_norm) * 100
        print(f"  {iteration + 1:>5}  {new_disp_norm:>12.6e}  "
              f"{pct:>+10.1f}%  {'converging' if pct > 0 else 'diverging':>10}")

    final_iter_norm = new_disp_norm
    final_iter_pct = (1 - final_iter_norm / before_norm) * 100

    # Fidelity check
    fid_to_init_before = uhlmann_fidelity(rho_final, rho_init)
    fid_to_init_after = uhlmann_fidelity(rho_corrected, rho_init)
    print(f"\n  Fidelity to initial state:")
    print(f"    Before correction: {fid_to_init_before:.6f}")
    print(f"    After correction:  {fid_to_init_after:.6f}")

    # Purity check
    purity_before = np.real(np.trace(rho_final @ rho_final))
    purity_after = np.real(np.trace(rho_corrected @ rho_corrected))
    purity_init = np.real(np.trace(rho_init @ rho_init))
    print(f"\n  Purity:")
    print(f"    Initial:           {purity_init:.6f}")
    print(f"    Before correction: {purity_before:.6f}")
    print(f"    After correction:  {purity_after:.6f}")

    # ----------------------------------------------------------------
    # 5. Honest Assessment
    # ----------------------------------------------------------------
    separator("5. Honest Assessment")

    print("  What the 5-clock theory demonstrates:")
    print()
    print("  [+] Auto-scaling makes clock divergence VISIBLE")
    print(f"      Outer clock fidelity: {outer_fid:.6f} (< 1.0)")
    print()
    print("  [+] Per-dimension displacement is MEASURABLE")
    print(f"      Total drift: {before_norm:.6e} across {n_dims} KK dimensions")
    print()
    print("  [+] Clock stencil gives finite-difference VELOCITY estimates")
    n_nonzero_v = np.sum(np.abs(result.clock_velocity_estimates) > 1e-15)
    print(f"      {n_nonzero_v}/{n_dims} dimensions have non-zero velocity")
    print()
    print("  [+] Return unitary REDUCES coordinate displacement")
    print(f"      Single-step (gain={best_gain:.0f}): {pct_total:+.1f}% reduction")
    print(f"      Iterative (8 steps, gain=10): {final_iter_pct:+.1f}% reduction")

    print()
    print("  Limitations (honest physics):")
    print()
    print("  [-] Unitary correction CANNOT restore purity lost to decoherence")
    print(f"      Purity dropped from {purity_init:.6f} to {purity_before:.6f}")
    print(f"      After correction: {purity_after:.6f} (same -- unitary preserves purity)")
    print()
    print("  [-] Return-to-origin corrects COHERENT displacement only")
    print("      Incoherent mixing from T1/T2 processes is irreversible")
    print("      (requires actual quantum error correction, not just unitary rotation)")
    print()
    print("  [-] Finite-difference velocity is a LOCAL estimate")
    print("      Accurate for short times; longer evolution needs repeated sampling")

    # ----------------------------------------------------------------
    # Summary
    # ----------------------------------------------------------------
    separator("Summary")

    print(f"  System:              KK tower T^7, dim_q={builder.dim_q}")
    print(f"  Clocks:              5 (offsets [-2,-1,0,+1,+2])")
    print(f"  Offset multiplier:   {auto_multiplier:.1f} (auto-scaled)")
    print(f"  Phase rotation:      {E_range * dt * auto_multiplier:.4f} rad")
    print(f"  Outer clock F:       {outer_fid:.6f}")
    print(f"  Total displacement:  {before_norm:.6e}")
    print(f"  Single-step corr:   {after_norm:.6e} ({pct_total:+.1f}%)")
    print(f"  Iterative corr:     {final_iter_norm:.6e} ({final_iter_pct:+.1f}%)")
    print(f"  Fidelity (init):     {fid_to_init_before:.6f} -> {fid_to_init_after:.6f}")
    print(f"  Purity:              {purity_init:.6f} -> {purity_before:.6f}"
          f" -> {purity_after:.6f}")
    print()

    if final_iter_pct > 0:
        print("  CONCLUSION: 5-clock theory WORKS for tracking and correcting")
        print("  coherent displacement through KK-compactified dimensions.")
    else:
        print("  CONCLUSION: Further tuning needed.")


if __name__ == "__main__":
    main()
