%reload_ext autoreload
%autoreload 2

from visualization_2D import *

# ============================================================================
# EXAMPLES (2D)
# ============================================================================

def example_2d_harmonic_oscillator():
    """Example: 2D Harmonic oscillator (separable)"""
    x, y = sp.symbols('x y', real=True)
    xi, eta = sp.symbols('xi eta', real=True)
    
    H = xi**2 + eta**2 + x**2 + y**2
    
    print("\n=== 2D HARMONIC OSCILLATOR ===")
    print("H = ξ² + η² + x² + y²")
    print("Integrable system with circular orbits\n")
    
    fig, geos, orbits, caust = visualize_symbol_2d(
        H,
        x_range=(-3, 3),
        y_range=(-3, 3),
        xi_range=(-3, 3),
        eta_range=(-3, 3),
        geodesics_params=[
            (1, 0, 0, 1, np.pi, 'red'),
            (0, 1, 1, 0, np.pi, 'blue'),
            (1, 1, 0.5, 0.5, np.pi, 'green'),
            (-1, 1, 1, 1, np.pi, 'purple'),
        ],
        E_range=(1, 8),
        hbar=1.0,
        resolution=40
    )
    
    plt.savefig('2d_harmonic_oscillator.png', dpi=150, bbox_inches='tight')
    plt.show()
    
    print(f"✓ {len(geos)} geodesics computed")
    print(f"✓ {len(orbits)} periodic orbits found")
    print(f"✓ {len(caust)} caustic curves detected\n")
    
    return fig, geos, orbits, caust


def example_2d_anisotropic():
    """Example: Anisotropic harmonic oscillator"""
    x, y = sp.symbols('x y', real=True)
    xi, eta = sp.symbols('xi eta', real=True)
    
    # Different frequencies: ω_x = 1, ω_y = 2
    H = xi**2 + eta**2 + x**2 + 4*y**2
    
    print("\n=== 2D ANISOTROPIC OSCILLATOR ===")
    print("H = ξ² + η² + x² + 4y²")
    print("Different frequencies in x and y\n")
    
    fig, geos, orbits, caust = visualize_symbol_2d(
        H,
        x_range=(-3, 3),
        y_range=(-2, 2),
        xi_range=(-3, 3),
        eta_range=(-4, 4),
        geodesics_params=[
            (2, 0, 0, 1, 4, 'red'),
            (0, 1, 1, 0, 4, 'blue'),
            (1.5, 0.5, 0.5, 1.5, 4, 'green'),
        ],
        E_range=(1, 10),
        hbar=1.0,
        resolution=40
    )
    
    plt.savefig('2d_anisotropic.png', dpi=150, bbox_inches='tight')
    plt.show()
    
    print(f"✓ Visualization complete\n")
    
    return fig, geos, orbits, caust


def example_2d_mexican_hat():
    """Example: Mexican hat potential"""
    x, y = sp.symbols('x y', real=True)
    xi, eta = sp.symbols('xi eta', real=True)
    
    # Mexican hat: V(x,y) = (x² + y² - r₀²)²
    r0 = 1.5
    H = xi**2 + eta**2 + (x**2 + y**2 - r0**2)**2
    
    print("\n=== 2D MEXICAN HAT POTENTIAL ===")
    print(f"H = ξ² + η² + (x² + y² - {r0}²)²")
    print("Circular symmetry with minimum at r=r₀\n")
    
    fig, geos, orbits, caust = visualize_symbol_2d(
        H,
        x_range=(-3, 3),
        y_range=(-3, 3),
        xi_range=(-2, 2),
        eta_range=(-2, 2),
        geodesics_params=[
            (r0, 0, 0, 0.5, 10, 'red'),
            (0, r0, 0.5, 0, 10, 'blue'),
            (r0/np.sqrt(2), r0/np.sqrt(2), 0.3, 0.3, 10, 'green'),
        ],
        E_range=(0.1, 3),
        hbar=0.5,
        resolution=40
    )
    
    plt.savefig('2d_mexican_hat.png', dpi=150, bbox_inches='tight')
    plt.show()
    
    print(f"✓ Visualization complete\n")
    
    return fig, geos, orbits, caust


def example_2d_coupled():
    """Example: Coupled oscillators (non-separable)"""
    x, y = sp.symbols('x y', real=True)
    xi, eta = sp.symbols('xi eta', real=True)
    
    # Coupling term x*y
    H = xi**2 + eta**2 + x**2 + y**2 + 0.5*x*y
    
    print("\n=== 2D COUPLED OSCILLATORS ===")
    print("H = ξ² + η² + x² + y² + 0.5xy")
    print("Non-separable, mixed mode dynamics\n")
    
    fig, geos, orbits, caust = visualize_symbol_2d(
        H,
        x_range=(-3, 3),
        y_range=(-3, 3),
        xi_range=(-3, 3),
        eta_range=(-3, 3),
        geodesics_params=[
            (1, 0, 0, 1, 5, 'red'),
            (0, 1, 1, 0, 5, 'blue'),
            (1, 1, 0.5, 0.5, 5, 'green'),
            (-1, 1, 1, -1, 5, 'orange'),
        ],
        E_range=(1, 8),
        hbar=1.0,
        resolution=40
    )
    
    plt.savefig('2d_coupled.png', dpi=150, bbox_inches='tight')
    plt.show()
    
    print(f"✓ Visualization complete\n")
    
    return fig, geos, orbits, caust


def example_2d_laplacian():
    """Example: 2D Laplacian (free particle)"""
    x, y = sp.symbols('x y', real=True)
    xi, eta = sp.symbols('xi eta', real=True)
    
    H = xi**2 + eta**2
    
    print("\n=== 2D LAPLACIAN (FREE PARTICLE) ===")
    print("H = ξ² + η²")
    print("Straight line geodesics\n")
    
    fig, geos, orbits, caust = visualize_symbol_2d(
        H,
        x_range=(-5, 15),
        y_range=(-5, 15),
        xi_range=(-2, 2),
        eta_range=(-2, 2),
        geodesics_params=[
            (0, 0, 1, 0, 10, 'red'),
            (0, 0, 0, 1, 10, 'blue'),
            (0, 0, 1, 1, 10, 'green'),
            (0, 0, 1, -1, 10, 'purple'),
        ],
        E_range=(0.5, 3),
        hbar=1.0,
        resolution=40
    )
    
    plt.savefig('2d_laplacian.png', dpi=150, bbox_inches='tight')
    plt.show()
    
    print(f"✓ Visualization complete\n")
    
    return fig, geos, orbits, caust

# ============================================================================
# ADVANCED 2D EXAMPLE: STADIUM BILLIARD (CHAOTIC)
# ============================================================================
def example_2d_stadium_billiard():
    """
    Example: Stadium billiard potential (chaotic system)
    
    The stadium billiard is a classic example of quantum chaos
    """
    x, y = sp.symbols('x y', real=True)
    xi, eta = sp.symbols('xi eta', real=True)
    
    # Approximate stadium shape with steep potential walls
    a, b = 2.0, 1.0  # Semi-axes
    epsilon = 10  # Steepness
    
    # Stadium-like potential (smooth approximation)
    # V = 0 inside stadium, ∞ outside
    # Approximated by: V ~ exp(ε * (distance from stadium))
    
    # For simplicity, use a "soft stadium" potential
    V = 0.1 * (sp.Max(0, x**2/a**2 + y**2/b**2 - 1)**2)
    H = xi**2 + eta**2 + V
    
    print("\n=== STADIUM BILLIARD (CHAOTIC) ===")
    print("Soft stadium potential")
    print("Expected: chaotic dynamics, Wigner statistics\n")
    
    fig, geos, orbits, caust = visualize_symbol_2d(
        H,
        x_range=(-3, 3),
        y_range=(-2, 2),
        xi_range=(-2, 2),
        eta_range=(-2, 2),
        geodesics_params=[
            (0.5, 0, 0.5, 0.5, 20, 'red'),
            (0, 0.5, 0.5, 0.3, 20, 'blue'),
            (1.0, 0.3, 0.4, 0.6, 20, 'green'),
            (-0.5, 0.5, 0.6, 0.4, 20, 'purple'),
        ],
        E_range=(1, 5),
        hbar=0.5,
        resolution=40
    )
    
    plt.savefig('2d_stadium_chaotic.png', dpi=150, bbox_inches='tight')
    plt.show()
    
    if orbits:
        # Analyze KAM structure
        Utilities2D.plot_kam_structure(orbits, 'stadium_kam.png')
    
    print(f"✓ Stadium billiard visualization complete\n")
    
    return fig, geos, orbits, caust


# ============================================================================
# COMPARISON: INTEGRABLE VS CHAOTIC (2D)
# ============================================================================

def compare_integrable_chaotic_2d():
    """
    Side-by-side comparison of integrable vs chaotic 2D systems
    """
    print("\n" + "="*70)
    print("COMPARISON: INTEGRABLE vs CHAOTIC (2D)")
    print("="*70)
    
    x, y = sp.symbols('x y', real=True)
    xi, eta = sp.symbols('xi eta', real=True)
    
    # Integrable: 2D harmonic oscillator
    H_integrable = xi**2 + eta**2 + x**2 + y**2
    
    # Chaotic: coupled with anharmonic term
    H_chaotic = xi**2 + eta**2 + x**2 + y**2 + 0.3*x**2*y**2
    
    fig = plt.figure(figsize=(20, 10))
    
    # INTEGRABLE
    print("\n[1/2] Analyzing integrable system...")
    geo_int = SymbolGeometry2D(H_integrable, x, y, xi, eta)
    
    # Compute some geodesics
    geos_int = []
    for params in [(1, 0, 0, 1, 2*np.pi), (0, 1, 1, 0, 2*np.pi)]:
        geo = geo_int.compute_geodesic(*params)
        geos_int.append(geo)
    
    # Find periodic orbits
    orbits_int = geo_int.find_periodic_orbits_2d(2.0, (-2,2), (-2,2), (-2,2), (-2,2), 15)
    
    # CHAOTIC
    print("\n[2/2] Analyzing chaotic system...")
    geo_chaos = SymbolGeometry2D(H_chaotic, x, y, xi, eta)
    
    geos_chaos = []
    for params in [(1, 0, 0, 1, 10), (0, 1, 1, 0, 10)]:
        geo = geo_chaos.compute_geodesic(*params)
        geos_chaos.append(geo)
    
    orbits_chaos = geo_chaos.find_periodic_orbits_2d(2.0, (-2,2), (-2,2), (-2,2), (-2,2), 15)
    
    # VISUALIZATION
    # Panel 1: Integrable - Configuration space
    ax1 = fig.add_subplot(2, 4, 1)
    for geo in geos_int:
        ax1.plot(geo.x, geo.y, linewidth=2)
    ax1.set_xlabel('x')
    ax1.set_ylabel('y')
    ax1.set_title('INTEGRABLE\nConfiguration space', fontweight='bold')
    ax1.grid(True, alpha=0.3)
    ax1.set_aspect('equal')
    
    # Panel 2: Chaotic - Configuration space
    ax2 = fig.add_subplot(2, 4, 5)
    for geo in geos_chaos:
        ax2.plot(geo.x, geo.y, linewidth=2)
    ax2.set_xlabel('x')
    ax2.set_ylabel('y')
    ax2.set_title('CHAOTIC\nConfiguration space', fontweight='bold')
    ax2.grid(True, alpha=0.3)
    ax2.set_aspect('equal')
    
    # Panel 3: Integrable - Poincaré section
    ax3 = fig.add_subplot(2, 4, 2)
    for geo in geos_int:
        for i in range(len(geo.y)-1):
            if geo.y[i] * geo.y[i+1] < 0:
                alpha = -geo.y[i] / (geo.y[i+1] - geo.y[i])
                x_cross = geo.x[i] + alpha * (geo.x[i+1] - geo.x[i])
                xi_cross = geo.xi[i] + alpha * (geo.xi[i+1] - geo.xi[i])
                ax3.plot(x_cross, xi_cross, 'bo', markersize=5)
    ax3.set_xlabel('x')
    ax3.set_ylabel('ξ')
    ax3.set_title('Poincaré section\n(regular structure)', fontweight='bold')
    ax3.grid(True, alpha=0.3)
    
    # Panel 4: Chaotic - Poincaré section
    ax4 = fig.add_subplot(2, 4, 6)
    for geo in geos_chaos:
        for i in range(len(geo.y)-1):
            if geo.y[i] * geo.y[i+1] < 0:
                alpha = -geo.y[i] / (geo.y[i+1] - geo.y[i])
                x_cross = geo.x[i] + alpha * (geo.x[i+1] - geo.x[i])
                xi_cross = geo.xi[i] + alpha * (geo.xi[i+1] - geo.xi[i])
                ax4.plot(x_cross, xi_cross, 'ro', markersize=5)
    ax4.set_xlabel('x')
    ax4.set_ylabel('ξ')
    ax4.set_title('Poincaré section\n(chaotic mixing)', fontweight='bold')
    ax4.grid(True, alpha=0.3)
    
    # Panel 5: Integrable - Periodic orbits
    ax5 = fig.add_subplot(2, 4, 3)
    if orbits_int:
        E_int = [orb.energy for orb in orbits_int]
        S_int = [orb.action for orb in orbits_int]
        ax5.scatter(E_int, S_int, s=150, c='blue', edgecolors='black', linewidths=2)
    ax5.set_xlabel('Energy E')
    ax5.set_ylabel('Action S')
    ax5.set_title('Periodic orbits\n(regular)', fontweight='bold')
    ax5.grid(True, alpha=0.3)
    
    # Panel 6: Chaotic - Periodic orbits
    ax6 = fig.add_subplot(2, 4, 7)
    if orbits_chaos:
        E_chaos = [orb.energy for orb in orbits_chaos]
        S_chaos = [orb.action for orb in orbits_chaos]
        ax6.scatter(E_chaos, S_chaos, s=150, c='red', edgecolors='black', linewidths=2)
    ax6.set_xlabel('Energy E')
    ax6.set_ylabel('Action S')
    ax6.set_title('Periodic orbits\n(scarce, unstable)', fontweight='bold')
    ax6.grid(True, alpha=0.3)
    
    # Panel 7: Level spacing - Integrable
    ax7 = fig.add_subplot(2, 4, 4)
    if len(orbits_int) > 2:
        energies_int = sorted(set(orb.energy for orb in orbits_int))
        spacings_int = np.diff(energies_int)
        s_mean_int = np.mean(spacings_int)
        s_norm_int = spacings_int / s_mean_int
        
        ax7.hist(s_norm_int, bins=10, density=True, alpha=0.7, color='blue', edgecolor='black')
        s = np.linspace(0, 3, 100)
        ax7.plot(s, np.exp(-s), 'g--', linewidth=2, label='Poisson')
        ax7.legend()
    ax7.set_xlabel('Normalized spacing s')
    ax7.set_ylabel('P(s)')
    ax7.set_title('Level spacing\nPoisson distribution', fontweight='bold')
    ax7.grid(True, alpha=0.3)
    
    # Panel 8: Level spacing - Chaotic
    ax8 = fig.add_subplot(2, 4, 8)
    if len(orbits_chaos) > 2:
        energies_chaos = sorted(set(orb.energy for orb in orbits_chaos))
        spacings_chaos = np.diff(energies_chaos)
        s_mean_chaos = np.mean(spacings_chaos)
        s_norm_chaos = spacings_chaos / s_mean_chaos
        
        ax8.hist(s_norm_chaos, bins=10, density=True, alpha=0.7, color='red', edgecolor='black')
        s = np.linspace(0, 3, 100)
        wigner = (np.pi * s / 2) * np.exp(-np.pi * s**2 / 4)
        ax8.plot(s, wigner, 'b-', linewidth=2, label='Wigner')
        ax8.legend()
    ax8.set_xlabel('Normalized spacing s')
    ax8.set_ylabel('P(s)')
    ax8.set_title('Level spacing\nWigner distribution', fontweight='bold')
    ax8.grid(True, alpha=0.3)
    
    plt.suptitle('INTEGRABLE vs CHAOTIC: 2D Comparison', fontsize=18, fontweight='bold')
    plt.tight_layout()
    plt.savefig('2d_comparison_integrable_chaotic.png', dpi=150, bbox_inches='tight')
    plt.show()
    
    print("\n✓ Comparison complete")
    print(f"  Integrable: {len(orbits_int)} periodic orbits")
    print(f"  Chaotic: {len(orbits_chaos)} periodic orbits\n")


# ============================================================================
# MAIN
# ============================================================================

if __name__ == "__main__":
    print("""
╔══════════════════════════════════════════════════════════════════════╗
║                2D SYMBOL GEOMETRIC VISUALIZATION                      ║
║           H(x, y, ξ, η) - 4D Phase Space Analysis                    ║
╚══════════════════════════════════════════════════════════════════════╝

2D Extensions:
- 4D phase space (projections)
- Caustic curves (not just points)
- Poincaré sections
- Torus quantization (KAM theory)
- Angular momentum conservation
- 2D spectral density

Running examples...
    """)
    
    # Run all 2D examples
    print("\n" + "="*70)
    print("EXAMPLES GALLERY (2D)")
    print("="*70)

    print_theory_summary()
    
    print("\n[1/6] 2D Harmonic oscillator...")
    example_2d_harmonic_oscillator()
    
    print("\n[2/6] Anisotropic oscillator...")
    example_2d_anisotropic()
    
    print("\n[3/6] Mexican hat potential...")
    example_2d_mexican_hat()
    
    print("\n[4/6] Coupled oscillators...")
    example_2d_coupled()
    
    print("\n[5/6] 2D Laplacian...")
    example_2d_laplacian()
    
    print("\n[6/6] 2D Satium billiard...")
    example_2d_stadium_billiard()
    compare_integrable_chaotic_2d()

    print("\n" + "="*70)
    print("✓ All 2D visualizations complete!")
    print("="*70)

╔══════════════════════════════════════════════════════════════════════╗
║                2D SYMBOL GEOMETRIC VISUALIZATION                      ║
║           H(x, y, ξ, η) - 4D Phase Space Analysis                    ║
╚══════════════════════════════════════════════════════════════════════╝

2D Extensions:
- 4D phase space (projections)
- Caustic curves (not just points)
- Poincaré sections
- Torus quantization (KAM theory)
- Angular momentum conservation
- 2D spectral density

Running examples...
    

======================================================================
EXAMPLES GALLERY (2D)
======================================================================

╔══════════════════════════════════════════════════════════════════════╗
║        GEOMETRIC AND SEMI-CLASSICAL THEORY FOR 2D SYSTEMS            ║
╚══════════════════════════════════════════════════════════════════════╝
PHASE SPACE AND SYMPLECTIC GEOMETRY
────────────────────────────────────
- Phase space: T*ℝ² ≅ ℝ⁴ with coordinates (x,y,ξ,η)
- Symplectic form: ω = dx∧dξ + dy∧dη
- Hamiltonian flow: dx/dt = ∂H/∂ξ, dy/dt = ∂H/∂η, dξ/dt = -∂H/∂x, dη/dt = -∂H/∂y
- Lagrangian manifolds: maximal submanifolds where ω = 0
CAUSTICS AND CATASTROPHE THEORY
───────────────────────────────
- Caustics: envelopes of trajectories where Jacobian vanishes
- Arnold classification:
  * Fold: generic singularity, Maslov index μ=1
  * Cusp: higher-order singularity, Maslov index μ=2
- Full 4×4 Jacobian for rigorous detection
- Phase shift upon crossing: -μπ/2
KAM THEORY AND INTEGRABLE SYSTEMS
────────────────────────────────
- Integrable systems: two independent first integrals (H, L)
- Phase space foliated by 2D tori
- KAM theorem: persistence of tori under small perturbations
- Rotation numbers (ω₁, ω₂): frequencies on torus
- EBK quantization: S_i = 2πℏ(n_i + α_i) with Maslov corrections
SEMI-CLASSICAL QUANTUM ANALYSIS
──────────────────────────────
- Weyl law: N(E) ~ Vol({H≤E}) / (2πℏ)²
- Spectral expansion: N(E) = N_smooth(E) + N_osc(E)
- Oscillatory terms: contributions from periodic orbits and caustics
- Level spacing distribution:
  * Poisson e^(-s): integrable systems
  * Wigner (πs/2)e^(-πs²/4): chaotic systems
PRACTICAL APPLICATIONS
──────────────────────
1. Geometrical optics and light caustics
2. Semi-classical quantum mechanics
3. Quantum chaos theory
4. Spectral analysis of pseudo-differential operators
5. Visualization of classical-quantum transition
USAGE
─────
>>> from symbol_visualizer_2d import visualize_symbol_2d, print_theory_summary
>>> x, y = sp.symbols('x y', real=True)
>>> xi, eta = sp.symbols('xi eta', real=True)
>>> H = xi**2 + eta**2 + x**2*y**2 + 0.1*(x**2 - 1)**2  # Complex example
>>> fig, geos, orbits, caustics = visualize_symbol_2d(
...     symbol=H,
...     x_range=(-2, 2), y_range=(-2, 2),
...     xi_range=(-2, 2), eta_range=(-2, 2),
...     geodesics_params=[
...         (0.5, 0.5, 0.5, 0.5, 10, 'red'),
...         (-0.5, 0.5, -0.5, 0.5, 10, 'blue'),
...         (0.0, 1.0, 1.0, 0.0, 15, 'green')
...     ],
...     E_range=(0.1, 5),
...     hbar=0.05,
...     resolution=60
... )
>>> plt.show()


[1/6] 2D Harmonic oscillator...

=== 2D HARMONIC OSCILLATOR ===
H = ξ² + η² + x² + y²
Integrable system with circular orbits

Initializing 2D geometry engine for H = eta**2 + x**2 + xi**2 + y**2 with ℏ = 1.0

Computing phase space volume (Monte Carlo)...
  E=1.00, Volume=5.3395
  E=2.00, Volume=19.6733
  E=3.00, Volume=46.0080
  E=4.00, Volume=79.0301
  E=5.00, Volume=121.8758
  E=6.00, Volume=180.4032
  E=7.00, Volume=243.6998
  E=8.00, Volume=312.1805

✓ 4 geodesics computed
✓ 1 periodic orbits found
✓ 12 caustic curves detected


[2/6] Anisotropic oscillator...

=== 2D ANISOTROPIC OSCILLATOR ===
H = ξ² + η² + x² + 4y²
Different frequencies in x and y

Initializing 2D geometry engine for H = eta**2 + x**2 + xi**2 + 4*y**2 with ℏ = 1.0

Computing phase space volume (Monte Carlo)...
  E=1.00, Volume=2.6726
  E=2.29, Volume=13.7088
  E=3.57, Volume=31.5648
  E=4.86, Volume=57.6230
  E=6.14, Volume=93.4733
  E=7.43, Volume=134.6227
  E=8.71, Volume=186.0250
  E=10.00, Volume=246.6893

✓ Visualization complete


[3/6] Mexican hat potential...

=== 2D MEXICAN HAT POTENTIAL ===
H = ξ² + η² + (x² + y² - 1.5²)²
Circular symmetry with minimum at r=r₀

Initializing 2D geometry engine for H = eta**2 + xi**2 + (x**2 + y**2 - 2.25)**2 with ℏ = 0.5

Computing phase space volume (Monte Carlo)...
  E=0.10, Volume=0.3456
  E=0.51, Volume=4.9766
  E=0.93, Volume=11.4163
  E=1.34, Volume=20.3443
  E=1.76, Volume=31.6570
  E=2.17, Volume=42.1517
  E=2.59, Volume=54.7661
  E=3.00, Volume=68.4288

✓ Visualization complete


[4/6] Coupled oscillators...

=== 2D COUPLED OSCILLATORS ===
H = ξ² + η² + x² + y² + 0.5xy
Non-separable, mixed mode dynamics

Initializing 2D geometry engine for H = eta**2 + x**2 + 0.5*x*y + xi**2 + y**2 with ℏ = 1.0

Computing phase space volume (Monte Carlo)...
  E=1.00, Volume=5.4432
  E=2.00, Volume=19.9066
  E=3.00, Volume=47.6150
  E=4.00, Volume=82.8144
  E=5.00, Volume=128.6410
  E=6.00, Volume=185.0429
  E=7.00, Volume=248.6765
  E=8.00, Volume=326.6957

✓ Visualization complete


[5/6] 2D Laplacian...

=== 2D LAPLACIAN (FREE PARTICLE) ===
H = ξ² + η²
Straight line geodesics

Initializing 2D geometry engine for H = eta**2 + xi**2 with ℏ = 1.0

Computing phase space volume (Monte Carlo)...
  E=0.50, Volume=624.3840
  E=0.86, Volume=1084.6720
  E=1.21, Volume=1529.0880
  E=1.57, Volume=1978.7520
  E=1.93, Volume=2424.7040
  E=2.29, Volume=2891.1360
  E=2.64, Volume=3345.9200
  E=3.00, Volume=3777.7920