"""
FlowNet Test Suite — Verify all components work.

Run with: python -m flownet.tests.test_core
"""

import torch
import sys
import os

# Add parent to path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

from flownet.core.particle import ParticleEncoder, ParticleState, ParticleUpdater, ParticleCombiner
from flownet.core.flow_field import FlowField, FlowFieldEfficient
from flownet.core.wave_processor import WaveProcessor, WaveAttention
from flownet.core.bond_system import BondSystem
from flownet.core.phase_engine import PhaseEngine
from flownet.core.consistency_field import ConsistencyField
from flownet.core.topological_memory import TopologicalMemory
from flownet.modules.flownet_model import FlowNetModel, FlowNetBlock
from flownet.utils.math_utils import (
    complex_wave, wave_interference, phase_coherence,
    pairwise_distances, compute_betti_numbers, topological_signature,
    kuramoto_step, lennard_jones_force
)


def test_particle_encoder():
    """Test particle encoding."""
    print("Testing ParticleEncoder...")
    encoder = ParticleEncoder(vocab_size=1000, d_semantic=128)
    token_ids = torch.randint(0, 1000, (2, 16))
    
    particles = encoder(token_ids)
    
    assert particles.semantic.shape == (2, 16, 128), f"Wrong shape: {particles.semantic.shape}"
    assert particles.phase.shape == (2, 16), f"Wrong phase shape: {particles.phase.shape}"
    assert particles.amplitude.shape == (2, 16), f"Wrong amplitude shape"
    assert (particles.amplitude >= 0).all(), "Amplitude should be non-negative"
    assert (particles.phase >= 0).all() and (particles.phase <= 2 * 3.14159).all(), "Phase out of range"
    print("  ✓ ParticleEncoder OK")


def test_flow_field():
    """Test flow field dynamics."""
    print("Testing FlowField...")
    d = 64
    flow = FlowField(d_semantic=d, n_propagation_steps=2)
    
    particles = ParticleState(
        semantic=torch.randn(2, 8, d),
        position=torch.randn(2, 8, 32),
        charge=torch.randn(2, 8, 16),
        mass=torch.ones(2, 8) * 0.5,
        spin=torch.randn(2, 8, 16),
        amplitude=torch.ones(2, 8) * 0.5,
        phase=torch.rand(2, 8) * 2 * 3.14159,
        memory_trace=torch.zeros(2, 8, 32),
    )
    
    updated, diag = flow(particles)
    assert updated.semantic.shape == particles.semantic.shape
    assert 'coherence' in diag
    print("  ✓ FlowField OK")


def test_flow_field_efficient():
    """Test efficient flow field."""
    print("Testing FlowFieldEfficient...")
    d = 64
    flow = FlowFieldEfficient(d_semantic=d, n_propagation_steps=2, local_radius=4)
    
    particles = ParticleState(
        semantic=torch.randn(2, 16, d),
        position=torch.randn(2, 16, 32),
        charge=torch.randn(2, 16, 16),
        mass=torch.ones(2, 16) * 0.5,
        spin=torch.randn(2, 16, 16),
        amplitude=torch.ones(2, 16) * 0.5,
        phase=torch.rand(2, 16) * 2 * 3.14159,
        memory_trace=torch.zeros(2, 16, 32),
    )
    
    updated, diag = flow(particles)
    assert updated.semantic.shape == particles.semantic.shape
    print("  ✓ FlowFieldEfficient OK")


def test_wave_processor():
    """Test wave processing."""
    print("Testing WaveProcessor...")
    d = 64
    wave = WaveProcessor(d_semantic=d, n_wave_heads=4, d_wave=16)
    
    particles = ParticleState(
        semantic=torch.randn(2, 8, d),
        position=torch.randn(2, 8, 32),
        charge=torch.randn(2, 8, 16),
        mass=torch.ones(2, 8) * 0.5,
        spin=torch.randn(2, 8, 16),
        amplitude=torch.ones(2, 8) * 0.5,
        phase=torch.rand(2, 8) * 2 * 3.14159,
        memory_trace=torch.zeros(2, 8, 32),
    )
    
    output, diag = wave(particles)
    assert output.shape == (2, 8, d), f"Wrong shape: {output.shape}"
    assert 'coherence' in diag
    print("  ✓ WaveProcessor OK")


def test_wave_attention():
    """Test wave-based attention."""
    print("Testing WaveAttention...")
    d = 64
    attn = WaveAttention(d_model=d, n_heads=4)
    
    x = torch.randn(2, 8, d)
    output = attn(x)
    assert output.shape == (2, 8, d)
    print("  ✓ WaveAttention OK")


def test_bond_system():
    """Test bond formation."""
    print("Testing BondSystem...")
    d = 64
    bonds = BondSystem(d_semantic=d, d_charge=16, d_spin=16)
    
    particles = ParticleState(
        semantic=torch.randn(2, 8, d),
        position=torch.randn(2, 8, 32),
        charge=torch.randn(2, 8, 16),
        mass=torch.ones(2, 8) * 0.5,
        spin=torch.randn(2, 8, 16),
        amplitude=torch.ones(2, 8) * 0.5,
        phase=torch.rand(2, 8) * 2 * 3.14159,
        memory_trace=torch.zeros(2, 8, 32),
    )
    
    bond_state, bond_output, diag = bonds(particles)
    assert bond_state.bond_strength.shape == (2, 8, 8)
    assert bond_output.shape == (2, 8, d)
    print("  ✓ BondSystem OK")


def test_phase_engine():
    """Test phase transition engine."""
    print("Testing PhaseEngine...")
    d = 64
    engine = PhaseEngine(d_semantic=d, min_steps=1, max_steps=4)
    
    particles = ParticleState(
        semantic=torch.randn(2, 8, d),
        position=torch.randn(2, 8, 32),
        charge=torch.randn(2, 8, 16),
        mass=torch.ones(2, 8) * 0.5,
        spin=torch.randn(2, 8, 16),
        amplitude=torch.ones(2, 8) * 0.5,
        phase=torch.rand(2, 8) * 2 * 3.14159,
        memory_trace=torch.zeros(2, 8, 32),
    )
    
    output, diag = engine(particles)
    assert output.semantic.shape == particles.semantic.shape
    assert 'total_steps' in diag
    print(f"  ✓ PhaseEngine OK (ran {diag['total_steps']} steps)")


def test_consistency_field():
    """Test consistency field."""
    print("Testing ConsistencyField...")
    d = 64
    field = ConsistencyField(d_semantic=d)
    
    semantic = torch.randn(2, 8, d)
    output, diag = field(semantic, n_minimize_steps=2)
    
    assert output.shape == semantic.shape
    assert 'energy_reduction' in diag
    print(f"  ✓ ConsistencyField OK (energy reduced by {diag['energy_reduction']:.4f})")


def test_topological_memory():
    """Test topological memory."""
    print("Testing TopologicalMemory...")
    d = 64
    memory = TopologicalMemory(d_semantic=d, max_entries=100)
    
    # Store some patterns
    for i in range(5):
        pattern = torch.randn(8, d)
        memory.store(pattern)
    
    # Query
    query = torch.randn(8, d)
    retrieved, indices, scores = memory.query(query, top_k=3)
    
    assert retrieved.shape == (3, 8, d)
    assert len(indices) == 3
    assert len(scores) == 3
    
    # Full forward pass
    context = torch.randn(2, 8, d)
    integrated, diag = memory(context, context, top_k=2)
    assert integrated.shape == context.shape
    
    print(f"  ✓ TopologicalMemory OK (stored {len(memory.entries)} entries)")


def test_math_utils():
    """Test mathematical utilities."""
    print("Testing math_utils...")
    
    # Wave functions
    amp = torch.ones(4)
    phase = torch.tensor([0.0, 3.14159, 1.5708, 4.7124])
    wave = complex_wave(amp, phase)
    assert wave.shape == (4,)
    
    # Interference
    w1 = complex_wave(torch.ones(4), torch.zeros(4))
    w2 = complex_wave(torch.ones(4), torch.zeros(4))
    interf = wave_interference(w1, w2)
    assert torch.allclose(interf, torch.ones(4) * 2, atol=0.1)  # constructive
    
    # Phase coherence
    in_phase = torch.zeros(8)
    assert phase_coherence(in_phase).item() > 0.99
    
    random_phases = torch.rand(8) * 2 * 3.14159
    assert phase_coherence(random_phases).item() < 0.99
    
    # Topology
    adj = torch.tensor([[0, 1, 0], [1, 0, 1], [0, 1, 0]]).float()
    betti = compute_betti_numbers(adj)
    assert betti[0] == 1  # one connected component
    assert betti[1] == 0  # no loops
    
    # Triangle has a loop
    adj_triangle = torch.tensor([[0, 1, 1], [1, 0, 1], [1, 1, 0]]).float()
    betti_triangle = compute_betti_numbers(adj_triangle)
    assert betti_triangle[1] == 1  # one loop
    
    sig = topological_signature(adj)
    assert sig.shape == (8,)
    
    print("  ✓ math_utils OK")


def test_flownet_block():
    """Test a single FlowNet block."""
    print("Testing FlowNetBlock...")
    d = 64
    block = FlowNetBlock(d_semantic=d, d_charge=16, d_spin=16, d_memory=32, n_wave_heads=4, d_wave=16)
    
    particles = ParticleState(
        semantic=torch.randn(2, 8, d),
        position=torch.randn(2, 8, 32),
        charge=torch.randn(2, 8, 16),
        mass=torch.ones(2, 8) * 0.5,
        spin=torch.randn(2, 8, 16),
        amplitude=torch.ones(2, 8) * 0.5,
        phase=torch.rand(2, 8) * 2 * 3.14159,
        memory_trace=torch.zeros(2, 8, 32),
    )
    
    updated, bond_state, diag = block(particles)
    assert updated.semantic.shape == particles.semantic.shape
    assert bond_state.bond_strength.shape == (2, 8, 8)
    print("  ✓ FlowNetBlock OK")


def test_flownet_model():
    """Test the complete FlowNet model."""
    print("Testing FlowNetModel...")
    
    model = FlowNetModel(
        vocab_size=1000,
        d_semantic=64,
        d_position=32,
        d_charge=16,
        d_spin=16,
        d_memory=32,
        n_blocks=2,
        n_wave_heads=4,
        d_wave=16,
        use_phase_engine=True,
        use_consistency=True,
        use_topological_memory=True,
        max_seq_len=128,
    )
    
    # Count parameters
    n_params = sum(p.numel() for p in model.parameters())
    print(f"  Parameters: {n_params:,}")
    
    # Forward pass
    token_ids = torch.randint(0, 1000, (2, 16))
    labels = torch.randint(0, 1000, (2, 16))
    
    output = model(token_ids, labels=labels)
    
    assert 'logits' in output
    assert 'loss' in output
    assert 'diagnostics' in output
    assert output['logits'].shape == (2, 16, 1000)
    assert output['loss'] is not None
    
    print(f"  Forward pass OK (loss: {output['loss'].item():.4f})")
    
    # Generation
    prompt = torch.randint(0, 1000, (1, 8))
    generated = model.generate(prompt, max_new_tokens=10)
    assert generated.shape == (1, 18)  # 8 prompt + 10 generated
    print(f"  Generation OK (output shape: {generated.shape})")
    
    print("  ✓ FlowNetModel OK")


def run_all_tests():
    """Run all tests."""
    print("=" * 60)
    print("FlowNet Test Suite")
    print("=" * 60)
    print()
    
    tests = [
        test_math_utils,
        test_particle_encoder,
        test_flow_field,
        test_flow_field_efficient,
        test_wave_processor,
        test_wave_attention,
        test_bond_system,
        test_phase_engine,
        test_consistency_field,
        test_topological_memory,
        test_flownet_block,
        test_flownet_model,
    ]
    
    passed = 0
    failed = 0
    
    for test_fn in tests:
        try:
            test_fn()
            passed += 1
        except Exception as e:
            print(f"  ✗ {test_fn.__name__} FAILED: {e}")
            import traceback
            traceback.print_exc()
            failed += 1
        print()
    
    print("=" * 60)
    print(f"Results: {passed} passed, {failed} failed out of {len(tests)}")
    print("=" * 60)
    
    return failed == 0


if __name__ == '__main__':
    success = run_all_tests()
    sys.exit(0 if success else 1)