#!/usr/bin/env python3
"""Simple test to verify the event_processor module works."""

import sys
sys.path.insert(0, '/home/samsc/me/rsc/transivent/src')

import numpy as np

# Test the event_processor module
try:
    from transivent import (
        extract_event_waveforms,
        calculate_msd_parallel,
        calculate_acf,
        fit_diffusion_linear,
        process_events_for_diffusion,
    )
    print("✓ Successfully imported event_processor functions")
except ImportError as e:
    print(f"✗ Import error: {e}")
    sys.exit(1)

def test_extract_event_waveforms():
    """Test event waveform extraction."""
    print("\nTesting extract_event_waveforms...")
    
    # Create test data
    t = np.linspace(0, 1, 1000)
    x = np.random.normal(0, 0.1, 1000)
    
    # Add some events
    x[200:300] += 1.0  # Event 1
    x[500:600] += 0.5  # Event 2
    
    # Define events
    events = np.array([[0.2, 0.3], [0.5, 0.6]])
    
    # Extract waveforms
    waveforms = extract_event_waveforms(t, x, events)
    
    assert len(waveforms) == 2, f"Expected 2 waveforms, got {len(waveforms)}"
    assert len(waveforms[0]) == 100, f"Expected 100 points in first event, got {len(waveforms[0])}"
    assert len(waveforms[1]) == 100, f"Expected 100 points in second event, got {len(waveforms[1])}"
    
    # Check that events have been extracted correctly
    assert np.mean(waveforms[0]) > 0.5, "First event should have positive mean"
    assert np.mean(waveforms[1]) > 0.2, "Second event should have positive mean"
    
    print("✓ extract_event_waveforms test passed")

def test_calculate_msd():
    """Test MSD calculation."""
    print("\nTesting calculate_msd_parallel...")
    
    # Create a simple Brownian motion trajectory
    n_points = 1000
    dt = 1e-6
    D_true = 1e-12  # True diffusion coefficient
    
    # Generate random walk
    np.random.seed(42)
    steps = np.random.normal(0, np.sqrt(2 * D_true * dt), n_points)
    x = np.cumsum(steps)
    
    # Calculate MSD
    taus, msds, counts = calculate_msd_parallel(x, dt=dt, max_lag=100, n_jobs=1)
    
    assert len(taus) == 100, f"Expected 100 lag times, got {len(taus)}"
    assert len(msds) == 100, f"Expected 100 MSD values, got {len(msds)}"
    assert len(counts) == 100, f"Expected 100 counts, got {len(counts)}"
    
    # Check that MSD increases with time (for diffusive motion)
    assert msds[50] > msds[10], "MSD should increase with lag time"
    
    print("✓ calculate_msd_parallel test passed")

def test_calculate_acf():
    """Test ACF calculation."""
    print("\nTesting calculate_acf...")
    
    # Create a simple signal
    n_points = 1000
    dt = 1e-6
    x = np.random.normal(0, 1, n_points)
    
    # Add some correlation
    for i in range(1, n_points):
        x[i] += 0.5 * x[i-1]
    
    # Calculate ACF
    lags, acf = calculate_acf(x, dt=dt, max_lag=100)
    
    assert len(lags) == 101, f"Expected 101 lag values (0 to 100), got {len(lags)}"
    assert len(acf) == 101, f"Expected 101 ACF values, got {len(acf)}"
    
    # Check that ACF at lag 0 is variance
    assert np.abs(acf[0] - np.var(x)) < 0.1, f"ACF(0) should be variance, got {acf[0]}"
    
    # Check that ACF decreases with lag (for correlated signal)
    assert acf[50] < acf[10], "ACF should decrease with lag time"
    
    print("✓ calculate_acf test passed")

def test_fit_diffusion():
    """Test diffusion coefficient fitting."""
    print("\nTesting fit_diffusion_linear...")
    
    # Create synthetic MSD data
    taus = np.linspace(0, 1e-5, 100)
    D_true = 1e-12
    msds = 2 * D_true * taus + 0.1 * np.random.normal(0, 1e-24, 100)  # Add noise
    
    # Fit diffusion
    D_fit = fit_diffusion_linear(taus, msds, time_limit=3e-5)
    
    assert not np.isnan(D_fit), "Diffusion coefficient should not be NaN"
    assert D_fit > 0, "Diffusion coefficient should be positive"
    
    # Check that fitted value is close to true value (within 50% due to noise)
    assert np.abs(D_fit - D_true) / D_true < 0.5, f"Fitted D ({D_fit}) too far from true D ({D_true})"
    
    print("✓ fit_diffusion_linear test passed")

def test_process_events_for_diffusion():
    """Test the high-level wrapper function."""
    print("\nTesting process_events_for_diffusion...")
    
    # Create test data
    t = np.linspace(0, 1, 1000)
    x = np.random.normal(0, 0.1, 1000)
    
    # Add some events
    x[200:300] += 1.0  # Event 1
    x[500:600] += 0.5  # Event 2
    
    # Define events
    events = np.array([[0.2, 0.3], [0.5, 0.6]])
    
    # Process events
    results = process_events_for_diffusion(
        name="test",
        sampling_interval=1e-3,  # 1 ms
        data_path="",
        t=t,
        x=x,
        events=events,
        max_lag=50,
        n_jobs=1,
    )
    
    assert "diffusion_coeffs" in results, "Results should contain diffusion_coeffs"
    assert "acf_values" in results, "Results should contain acf_values"
    assert "event_count" in results, "Results should contain event_count"
    assert "statistics" in results, "Results should contain statistics"
    
    assert len(results["diffusion_coeffs"]) == 2, f"Expected 2 diffusion coeffs, got {len(results['diffusion_coeffs'])}"
    assert len(results["acf_values"]) == 2, f"Expected 2 ACF values, got {len(results['acf_values'])}"
    assert results["event_count"] == 2, f"Expected event_count=2, got {results['event_count']}"
    
    print("✓ process_events_for_diffusion test passed")

def main():
    """Run all tests."""
    print("Running event_processor module tests...")
    
    try:
        test_extract_event_waveforms()
        test_calculate_msd()
        test_calculate_acf()
        test_fit_diffusion()
        test_process_events_for_diffusion()
        
        print("\n✅ All tests passed!")
        return 0
    except AssertionError as e:
        print(f"\n❌ Test failed: {e}")
        return 1
    except Exception as e:
        print(f"\n❌ Unexpected error: {e}")
        import traceback
        traceback.print_exc()
        return 1

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