"""Classical N-body gravitational system demonstration using TNFR framework.

⚠️ **IMPORTANT**: This example ASSUMES Newtonian gravitational potential:
   U = -G*m₁*m₂/r

This is an **external assumption**, NOT derived from TNFR physics!

For a PURE TNFR formulation (no gravitational assumption), see:
   examples/domain_applications/nbody_tnfr_pure.py

That implementation derives dynamics from coherence potential and phase
synchronization, with NO assumptions about classical force laws.

What This Example Shows
------------------------

This example demonstrates how TNFR can **reproduce** classical mechanics
when we explicitly assume the classical potential. It shows:

1. **Mass as inverse frequency**: m = 1/νf
   Particles with high mass have low reorganization rate (inertia)
   
2. **Force as coherence gradient**: F = -∇U (ASSUMED U from gravity)
   Classical gravitational force mapped to TNFR framework
   
3. **Nodal equation**: ∂EPI/∂t = νf · ΔNFR
   Particle trajectories emerge from structural evolution
   
4. **Conservation laws emerge naturally**:
   Energy, momentum, angular momentum preserved

Comparison:
-----------

**This script** (nbody_gravitational.py):
- Assumes: U = -G*m₁*m₂/r (Newtonian gravity)
- Computes: F = -∇U, then ΔNFR = F/m
- Shows: TNFR can reproduce classical mechanics

**Pure TNFR** (nbody_tnfr_pure.py):
- Assumes: NOTHING (pure TNFR)
- Computes: ΔNFR from Hamiltonian commutator
- Shows: Attraction emerges from coherence/phase sync

Examples shown:
- Two-body circular orbit (Earth-Moon analogy)
- Three-body figure-8 orbit (choreographic solution)
- Solar system approximation (Sun + planets)

Run this script to see how TNFR reproduces classical mechanics
when classical potentials are explicitly assumed!
"""

import numpy as np
from tnfr.dynamics.nbody import NBodySystem


def print_section(title: str) -> None:
    """Print a formatted section header."""
    print()
    print("=" * 70)
    print(f" {title}")
    print("=" * 70)
    print()


def two_body_circular_orbit():
    """Demonstrate two-body circular orbit (Earth-Moon system)."""
    print_section("Example 1: Two-Body Circular Orbit (Earth-Moon)")
    
    print("Setting up Earth-Moon system in TNFR framework...")
    print("  - Earth mass: 1.0 (dimensionless)")
    print("  - Moon mass: 0.3 (adjusted for TNFR νf limits)")
    print("  - Initial separation: 1.0 (dimensionless units)")
    print()
    
    # Create system
    M_earth = 1.0
    M_moon = 0.3  # Adjusted to keep νf within TNFR default limits (νf = 1/0.3 = 3.33 < 10)
    system = NBodySystem(
        n_bodies=2,
        masses=[M_earth, M_moon],
        G=1.0
    )
    
    # Initial conditions for circular orbit
    r = 1.0  # Separation
    v_orbit = np.sqrt(system.G * (M_earth + M_moon) / r)
    
    positions = np.array([
        [0.0, 0.0, 0.0],  # Earth at origin
        [r, 0.0, 0.0]     # Moon at distance r
    ])
    velocities = np.array([
        [0.0, 0.0, 0.0],          # Earth at rest (CM frame)
        [0.0, v_orbit, 0.0]       # Moon with tangential velocity
    ])
    
    system.set_state(positions, velocities)
    
    print("Initial state:")
    print(f"  Earth position: {positions[0]}")
    print(f"  Moon position:  {positions[1]}")
    print(f"  Orbital velocity: {v_orbit:.4f}")
    print()
    
    # Initial energy and momentum
    K0, U0, E0 = system.compute_energy()
    P0 = system.compute_momentum()
    L0 = system.compute_angular_momentum()
    
    print(f"Initial energy:")
    print(f"  Kinetic:   {K0:+.6f}")
    print(f"  Potential: {U0:+.6f}")
    print(f"  Total:     {E0:+.6f}")
    print()
    print(f"Initial momentum: {P0}")
    print(f"Initial angular momentum: {L0}")
    print()
    
    # TNFR perspective
    print("TNFR Structural View:")
    print(f"  Earth νf (structural frequency): {1.0/M_earth:.4f} Hz_str")
    print(f"  Moon νf (structural frequency):  {1.0/M_moon:.4f} Hz_str")
    print("  → Higher νf = faster structural reorganization")
    print()
    
    # Evolve system
    print("Evolving system through 2 orbital periods...")
    T_orbit = 2 * np.pi * np.sqrt(r**3 / (system.G * (M_earth + M_moon)))
    print(f"Orbital period T = {T_orbit:.4f}")
    
    history = system.evolve(t_final=2*T_orbit, dt=0.01, store_interval=10)
    
    # Final state
    K_f, U_f, E_f = system.compute_energy()
    P_f = system.compute_momentum()
    L_f = system.compute_angular_momentum()
    
    print()
    print("Final state (after 2 periods):")
    print(f"Final energy:")
    print(f"  Kinetic:   {K_f:+.6f}")
    print(f"  Potential: {U_f:+.6f}")
    print(f"  Total:     {E_f:+.6f}")
    print()
    
    # Conservation checks
    print("Conservation Laws (TNFR Structural Invariants):")
    dE = abs(E_f - E0)
    rel_dE = dE / abs(E0) * 100
    print(f"  Energy drift:     ΔE/E₀ = {rel_dE:.4f}%")
    
    dP = np.linalg.norm(P_f - P0)
    print(f"  Momentum drift:   |ΔP|   = {dP:.2e}")
    
    dL = np.linalg.norm(L_f - L0)
    rel_dL = dL / np.linalg.norm(L0) * 100
    print(f"  Angular momentum: ΔL/L₀  = {rel_dL:.4f}%")
    print()
    
    if rel_dE < 1.0:
        print("✓ Excellent energy conservation! TNFR nodal equation working correctly.")
    else:
        print("⚠ Energy drift detected. Consider smaller time step.")
    
    print()
    print("Interpretation in TNFR paradigm:")
    print("  • Orbit is a stable resonance pattern between two nodes")
    print("  • Gravitational potential = coherence potential")
    print("  • System evolved to maximize coherence (minimize U)")
    print("  • Conservation laws emerge from Hamiltonian structure")
    
    return system, history


def three_body_symmetric():
    """Demonstrate three-body system with symmetry."""
    print_section("Example 2: Three-Body Symmetric Configuration")
    
    print("Setting up three equal masses in equilateral triangle...")
    print("  - All masses: 1.0")
    print("  - Triangle side length: 1.0")
    print()
    
    # Create system
    system = NBodySystem(
        n_bodies=3,
        masses=[1.0, 1.0, 1.0],
        G=1.0
    )
    
    # Equilateral triangle
    a = 1.0
    h = a * np.sqrt(3) / 2
    positions = np.array([
        [0.0, 0.0, 0.0],
        [a, 0.0, 0.0],
        [a/2, h, 0.0]
    ])
    
    # All bodies rotating about center of mass
    # Center is at (a/2, h/3, 0)
    center = positions.mean(axis=0)
    
    # Tangential velocities (perpendicular to radius from center)
    v_mag = 0.4  # Tuned for stable orbit
    velocities = np.zeros((3, 3))
    for i in range(3):
        r_vec = positions[i] - center
        # Perpendicular in xy-plane: (x, y) → (-y, x)
        tangent = np.array([-r_vec[1], r_vec[0], 0.0])
        tangent = tangent / np.linalg.norm(tangent)
        velocities[i] = v_mag * tangent
    
    system.set_state(positions, velocities)
    
    print("Initial configuration:")
    for i in range(3):
        print(f"  Body {i+1}: position {positions[i]}, "
              f"velocity {velocities[i]}")
    print()
    
    # Energy
    K0, U0, E0 = system.compute_energy()
    print(f"Initial energy: K={K0:.4f}, U={U0:.4f}, Total={E0:.4f}")
    print()
    
    print("TNFR Structural Interpretation:")
    print("  • Three resonant nodes with equal νf = 1.0 Hz_str")
    print("  • Symmetric coupling (all edges have same weight)")
    print("  • System seeks coherent rotating pattern")
    print()
    
    # Evolve
    print("Evolving system...")
    history = system.evolve(t_final=5.0, dt=0.01, store_interval=20)
    
    # Conservation check
    E_final = history['energy'][-1]
    rel_dE = abs(E_final - E0) / abs(E0) * 100
    
    print(f"Final energy drift: {rel_dE:.4f}%")
    
    if rel_dE < 5.0:
        print("✓ Good energy conservation for chaotic 3-body problem!")
    
    print()
    print("Note: Three-body problem is chaotic (no closed-form solution)")
    print("      But TNFR framework still preserves structural invariants!")
    
    return system, history


def solar_system_miniature():
    """Demonstrate miniature solar system (Sun + 3 planets)."""
    print_section("Example 3: Miniature Solar System (Sun + 3 Planets)")
    
    print("Setting up Sun with 3 planets at different orbits...")
    print()
    
    # Masses (Sun >> planets, but keeping νf within bounds)
    M_sun = 10.0  # Large but not too large
    M_planets = [1.0, 0.5, 0.2]  # Adjusted to keep νf < 10
    
    all_masses = [M_sun] + M_planets
    
    system = NBodySystem(
        n_bodies=4,
        masses=all_masses,
        G=1.0
    )
    
    # Initial positions (Sun at origin, planets in circular orbits)
    r_orbits = [1.0, 2.0, 3.5]  # Orbital radii
    positions = [np.array([0.0, 0.0, 0.0])]  # Sun
    velocities = [np.array([0.0, 0.0, 0.0])]  # Sun at rest
    
    print("Planet configurations:")
    for i, (r, m) in enumerate(zip(r_orbits, M_planets)):
        # Circular orbit velocity: v = sqrt(G*M_sun/r)
        v = np.sqrt(system.G * M_sun / r)
        
        # Position: equally spaced angles
        angle = 2 * np.pi * i / 3
        pos = np.array([r * np.cos(angle), r * np.sin(angle), 0.0])
        
        # Velocity: tangent to orbit
        vel = v * np.array([-np.sin(angle), np.cos(angle), 0.0])
        
        positions.append(pos)
        velocities.append(vel)
        
        print(f"  Planet {i+1}: mass={m:.2f}, radius={r:.1f}, "
              f"velocity={v:.4f}")
        print(f"           νf={1.0/m:.2f} Hz_str (structural frequency)")
    
    print()
    
    positions = np.array(positions)
    velocities = np.array(velocities)
    
    system.set_state(positions, velocities)
    
    # Energy
    K0, U0, E0 = system.compute_energy()
    print(f"Initial total energy: {E0:.6f}")
    print()
    
    print("TNFR Network Structure:")
    print(f"  Sun νf:        {1.0/M_sun:.4f} Hz_str (very low, high inertia)")
    print(f"  Planets νf:    {[f'{1.0/m:.2f}' for m in M_planets]} Hz_str")
    print("  → Sun barely reorganizes, planets adapt quickly")
    print("  → Gravitational coherence potential dominates dynamics")
    print()
    
    # Evolve
    print("Evolving for several orbital periods...")
    T_inner = 2 * np.pi * np.sqrt(r_orbits[0]**3 / (system.G * M_sun))
    print(f"Inner planet period: T ≈ {T_inner:.4f}")
    
    history = system.evolve(t_final=2*T_inner, dt=0.01, store_interval=20)
    
    # Conservation
    E_final = history['energy'][-1]
    rel_dE = abs(E_final - E0) / abs(E0) * 100
    
    print()
    print(f"Energy conservation: ΔE/E₀ = {rel_dE:.4f}%")
    
    if rel_dE < 1.0:
        print("✓ Excellent! Multi-body system remains coherent!")
    
    print()
    print("Physical Insight:")
    print("  Classical solar system = low-dissonance TNFR network")
    print("  Planetary motion = emergent coherence pattern")
    print("  Stability = structural resonance lock")
    
    return system, history


def demonstrate_coherence_potential():
    """Demonstrate coherence potential concept."""
    print_section("Bonus: Coherence Potential Demonstration")
    
    print("Showing how gravitational potential = coherence potential...")
    print()
    
    # Two bodies starting far apart, falling together
    system = NBodySystem(n_bodies=2, masses=[1.0, 1.0], G=1.0)
    
    positions = np.array([
        [0.0, 0.0, 0.0],
        [5.0, 0.0, 0.0]  # Far apart
    ])
    velocities = np.array([
        [0.0, 0.0, 0.0],
        [0.0, 0.0, 0.0]  # Start at rest
    ])
    
    system.set_state(positions, velocities)
    
    print("Initial state: Two bodies at rest, separated by r = 5.0")
    print()
    
    _, U0, _ = system.compute_energy()
    print(f"Initial potential (coherence): U = {U0:.6f}")
    print()
    
    # Evolve
    history = system.evolve(t_final=3.0, dt=0.01, store_interval=10)
    
    # Track potential over time
    print("Potential evolution (coherence increasing):")
    U_history = history['potential']
    times = history['time']
    
    for i in range(0, len(times), len(times)//5):
        print(f"  t={times[i]:5.2f}: U={U_history[i]:+.6f}")
    
    print()
    print("Observation:")
    print(f"  U decreased from {U0:.6f} to {U_history[-1]:.6f}")
    print("  → System moved toward HIGHER coherence (more negative U)")
    print("  → Gravitational attraction = coherence maximization!")
    print()
    print("TNFR Interpretation:")
    print("  Bodies aren't 'attracted' by a force")
    print("  They naturally evolve toward configurations of higher coherence")
    print("  Potential U encodes the structural stability landscape")


def main():
    """Run all demonstrations."""
    print()
    print("*" * 70)
    print("*" + " " * 68 + "*")
    print("*  N-Body Gravitational System in TNFR Framework" + " " * 20 + "*")
    print("*" + " " * 68 + "*")
    print("*  Classical Mechanics as Low-Dissonance Coherence Regime" + " " * 12 + "*")
    print("*" + " " * 68 + "*")
    print("*" * 70)
    
    try:
        # Example 1: Two-body orbit
        sys1, hist1 = two_body_circular_orbit()
        
        # Example 2: Three-body
        sys2, hist2 = three_body_symmetric()
        
        # Example 3: Solar system
        sys3, hist3 = solar_system_miniature()
        
        # Bonus: Coherence demonstration
        demonstrate_coherence_potential()
        
        # Summary
        print_section("Summary")
        print("Successfully demonstrated classical N-body problem in TNFR!")
        print()
        print("Key Results:")
        print("  ✓ Mass = 1/νf (structural frequency)")
        print("  ✓ Force = coherence gradient")
        print("  ✓ Trajectories from nodal equation ∂EPI/∂t = νf·ΔNFR")
        print("  ✓ Conservation laws preserved")
        print("  ✓ Stable orbits = resonance patterns")
        print()
        print("TNFR successfully reproduces classical mechanics!")
        print("Classical physics emerges as low-dissonance limit of TNFR.")
        print()
        
        # Offer to plot if matplotlib available
        try:
            import matplotlib
            print("Matplotlib detected! Plotting trajectories...")
            
            # Plot two-body orbit
            fig1 = sys1.plot_trajectories(hist1, show_energy=True)
            fig1.suptitle("Two-Body Orbit (Earth-Moon)", fontsize=14, y=0.98)
            
            # Plot three-body
            fig2 = sys2.plot_trajectories(hist2, show_energy=True)
            fig2.suptitle("Three-Body Symmetric System", fontsize=14, y=0.98)
            
            # Plot solar system
            fig3 = sys3.plot_trajectories(hist3, show_energy=True)
            fig3.suptitle("Miniature Solar System", fontsize=14, y=0.98)
            
            import matplotlib.pyplot as plt
            plt.show()
            
        except ImportError:
            print("To see plots, install matplotlib:")
            print("  pip install 'tnfr[viz-basic]'")
        
    except Exception as e:
        print(f"Error: {e}")
        import traceback
        traceback.print_exc()
        return 1
    
    return 0


if __name__ == "__main__":
    exit(main())
