"""
Monte Carlo System Test Suite
=============================

Comprehensive tests for the MC envelope mapping system.

Run with:
    python test_mc_system.py

Tests:
1. Sampler functionality
2. Validator constraint groups
3. Metrics computation
4. Store persistence
5. Executor (simulated mode)
6. Runner orchestration
7. ML training (optional)
"""

import sys
import os
import json
import tempfile
from pathlib import Path

# Add parent to path for imports
sys.path.insert(0, str(Path(__file__).parent.parent))
sys.path.insert(0, str(Path(__file__).parent))

import numpy as np


def test_sampler():
    """Test MCSampler."""
    print("\n" + "="*70)
    print("TEST 1: MCSampler")
    print("="*70)
    
    from nautilus_dolphin.mc import MCSampler, MCTrialConfig
    
    sampler = MCSampler(base_seed=42)
    
    # Test switch vector generation
    switches = sampler.generate_switch_vectors()
    print(f"  Switch vectors: {len(switches)}")
    assert 50 < len(switches) < 150, "Unexpected number of switch vectors"
    
    # Test trial generation (small)
    trials = sampler.generate_trials(n_samples_per_switch=5, max_trials=50)
    print(f"  Generated trials: {len(trials)}")
    assert len(trials) <= 50, "Too many trials generated"
    
    # Test champion generation
    champion = sampler.generate_champion_trial()
    assert champion.trial_id == -1, "Champion should have trial_id=-1"
    assert champion.vel_div_threshold == -0.02, "Champion threshold mismatch"
    
    # Test parameter ranges
    for trial in trials[:10]:
        assert -0.040 <= trial.vel_div_threshold <= -0.008, "vel_div_threshold out of range"
        assert 0.10 <= trial.min_leverage <= 1.50, "min_leverage out of range"
        assert 1.50 <= trial.max_leverage <= 12.00, "max_leverage out of range"
        assert 0.05 <= trial.fraction <= 0.40, "fraction out of range"
    
    print("  [PASS] Sampler tests passed")
    return True


def test_validator():
    """Test MCValidator."""
    print("\n" + "="*70)
    print("TEST 2: MCValidator")
    print("="*70)
    
    from nautilus_dolphin.mc import MCSampler, MCValidator, ValidationStatus
    
    sampler = MCSampler()
    validator = MCValidator(verbose=False)
    
    # Generate some trials
    trials = sampler.generate_trials(n_samples_per_switch=3, max_trials=30)
    
    # Validate
    results = validator.validate_batch(trials)
    
    # Stats
    stats = validator.get_validity_stats(results)
    print(f"  Total: {stats['total']}")
    print(f"  Valid: {stats['valid']} ({stats['validity_rate']*100:.1f}%)")
    
    # Most should be valid or rejected for specific reasons
    assert stats['total'] == len(trials), "Stats total mismatch"
    
    # Test specific constraints
    # Create an invalid config
    invalid_config = trials[0]._replace(
        vel_div_extreme=-0.01,  # Not extreme enough
        vel_div_threshold=-0.015  # Should fail: extreme >= threshold
    )
    result = validator.validate(invalid_config)
    assert not result.is_valid(), "Should reject invalid VD config"
    print(f"  Invalid config rejected: {result.reject_reason}")
    
    # Test leverage constraint
    invalid_config2 = trials[0]._replace(
        min_leverage=5.0,
        max_leverage=3.0  # Should fail: min >= max
    )
    result2 = validator.validate(invalid_config2)
    assert not result2.is_valid(), "Should reject invalid leverage config"
    
    print("  [PASS] Validator tests passed")
    return True


def test_metrics():
    """Test MCMetrics."""
    print("\n" + "="*70)
    print("TEST 3: MCMetrics")
    print("="*70)
    
    from nautilus_dolphin.mc import MCSampler, MCMetrics, MCTrialResult
    
    sampler = MCSampler()
    config = sampler.generate_champion_trial()
    
    metrics = MCMetrics(initial_capital=25000.0)
    
    # Create dummy data
    trades = [
        {'pnl': 100, 'pnl_pct': 0.004, 'exit_type': 'tp', 'bars_held': 50},
        {'pnl': -50, 'pnl_pct': -0.002, 'exit_type': 'stop', 'bars_held': 20},
        {'pnl': 150, 'pnl_pct': 0.006, 'exit_type': 'tp', 'bars_held': 80},
    ] * 20
    
    daily_pnls = [50, -20, 80, -10, 100, -30, 60, 40, -15, 90] * 5
    
    date_stats = [{'date': f'2026-01-{i%31+1:02d}', 'pnl': daily_pnls[i]} 
                  for i in range(len(daily_pnls))]
    
    signal_stats = {
        'dc_skip_rate': 0.1,
        'ob_skip_rate': 0.05,
        'dc_confirm_rate': 0.7,
        'irp_match_rate': 0.6,
        'entry_attempt_rate': 0.3,
        'signal_to_trade_rate': 0.15,
    }
    
    # Compute metrics
    result = metrics.compute(config, trades, daily_pnls, date_stats, signal_stats)
    
    print(f"  ROI: {result.roi_pct:.2f}%")
    print(f"  Profit Factor: {result.profit_factor:.2f}")
    print(f"  Win Rate: {result.win_rate:.2%}")
    print(f"  Sharpe: {result.sharpe_ratio:.2f}")
    print(f"  Max DD: {result.max_drawdown_pct:.2f}%")
    print(f"  N Trades: {result.n_trades}")
    
    # Verify labels
    print(f"  Profitable: {result.profitable}")
    print(f"  Champion Region: {result.champion_region}")
    print(f"  Catastrophic: {result.catastrophic}")
    
    # Test serialization
    result_dict = result.to_dict()
    assert 'P_vel_div_threshold' in result_dict, "Config params not in dict"
    assert 'M_roi_pct' in result_dict, "Metrics not in dict"
    assert 'L_profitable' in result_dict, "Labels not in dict"
    
    print("  [PASS] Metrics tests passed")
    return True


def test_store():
    """Test MCStore."""
    print("\n" + "="*70)
    print("TEST 4: MCStore")
    print("="*70)
    
    from nautilus_dolphin.mc import MCStore, MCSampler, MCMetrics, MCTrialResult
    from nautilus_dolphin.mc.mc_validator import ValidationResult, ValidationStatus
    
    # Use temp directory
    with tempfile.TemporaryDirectory() as tmpdir:
        store = MCStore(output_dir=tmpdir, batch_size=10)
        
        # Test validation result storage
        val_results = [
            ValidationResult(ValidationStatus.VALID, i)
            for i in range(5)
        ]
        val_results.append(ValidationResult(
            ValidationStatus.REJECTED_V2, 5, "Test rejection"
        ))
        
        store.save_validation_results(val_results, batch_id=0)
        
        # Test trial result storage
        sampler = MCSampler()
        config = sampler.generate_champion_trial()
        
        trial_results = []
        for i in range(5):
            result = MCTrialResult(
                trial_id=i,
                config=config._replace(trial_id=i),
                roi_pct=np.random.uniform(-10, 50),
                n_trades=np.random.randint(20, 200),
                status='completed'
            )
            result.compute_labels()
            trial_results.append(result)
        
        store.save_trial_results(trial_results, batch_id=1)
        
        # Test query
        entries = store.query_index(limit=10)
        print(f"  Index entries: {len(entries)}")
        assert len(entries) == 5, "Index query returned wrong count"
        
        # Test stats
        stats = store.get_corpus_stats()
        print(f"  Corpus stats: {stats}")
        assert stats['total_trials'] == 5, "Corpus stats wrong"
        
        print("  [PASS] Store tests passed")
    
    return True


def test_executor():
    """Test MCExecutor (simulated mode)."""
    print("\n" + "="*70)
    print("TEST 5: MCExecutor (Simulated)")
    print("="*70)
    
    from nautilus_dolphin.mc import MCSampler, MCExecutor
    
    sampler = MCSampler()
    executor = MCExecutor(initial_capital=25000.0, verbose=False)
    
    # Generate a few trials
    trials = sampler.generate_trials(n_samples_per_switch=2, max_trials=5)
    
    # Execute
    results = executor.execute_batch(trials, progress_interval=2)
    
    print(f"  Executed {len(results)} trials")
    
    for r in results:
        print(f"    Trial {r.trial_id}: ROI={r.roi_pct:.2f}%, Trades={r.n_trades}, Status={r.status}")
    
    # All should complete
    completed = sum(1 for r in results if r.status == 'completed')
    print(f"  Completed: {completed}/{len(results)}")
    
    assert completed > 0, "No trials completed"
    
    print("  [PASS] Executor tests passed")
    return True


def test_runner():
    """Test MCRunner (small scale)."""
    print("\n" + "="*70)
    print("TEST 6: MCRunner (Small Scale)")
    print("="*70)
    
    from nautilus_dolphin.mc import MCRunner
    
    with tempfile.TemporaryDirectory() as tmpdir:
        runner = MCRunner(
            output_dir=tmpdir,
            n_workers=1,
            batch_size=5,
            base_seed=42,
            verbose=False
        )
        
        # Run small envelope mapping
        stats = runner.run_envelope_mapping(
            n_samples_per_switch=2,
            max_trials=10,
            resume=False
        )
        
        print(f"  Total trials: {stats['total_trials']}")
        print(f"  Champion region: {stats['champion_count']}")
        print(f"  Catastrophic: {stats['catastrophic_count']}")
        
        assert stats['total_trials'] > 0, "No trials completed"
        
        # Test report generation
        report = runner.generate_report()
        assert 'Corpus Statistics' in report, "Report missing content"
        
        print("  [PASS] Runner tests passed")
    
    return True


def test_cli():
    """Test CLI entry point."""
    print("\n" + "="*70)
    print("TEST 7: CLI Entry Point")
    print("="*70)
    
    from nautilus_dolphin.run_mc_envelope import create_parser, cmd_sample, cmd_report
    
    # Test parser
    parser = create_parser()
    
    # Test sample command args
    args = parser.parse_args(['--mode', 'sample', '--n-samples', '5', '--max-trials', '10'])
    assert args.mode == 'sample'
    assert args.n_samples == 5
    
    print("  Parser works")
    
    # Test actual sample command (small)
    with tempfile.TemporaryDirectory() as tmpdir:
        args.output_dir = tmpdir
        args.seed = 42
        args.quiet = True
        
        result = cmd_sample(args)
        assert result == 0, "Sample command failed"
        
        # Check output
        config_path = Path(tmpdir) / "manifests" / "all_configs.json"
        assert config_path.exists(), "Config file not created"
        
        with open(config_path) as f:
            data = json.load(f)
        
        print(f"  Generated {len(data)} configs")
        assert len(data) <= 10, "Too many configs generated"
    
    print("  [PASS] CLI tests passed")
    return True


def test_integration():
    """Integration test: full pipeline."""
    print("\n" + "="*70)
    print("TEST 8: Full Integration")
    print("="*70)
    
    with tempfile.TemporaryDirectory() as tmpdir:
        # Step 1: Generate
        from nautilus_dolphin.mc import MCSampler
        sampler = MCSampler(base_seed=42)
        trials = sampler.generate_trials(n_samples_per_switch=2, max_trials=10)
        
        # Step 2: Validate
        from nautilus_dolphin.mc import MCValidator
        validator = MCValidator()
        valid_trials = [t for t in trials if validator.validate(t).is_valid()]
        print(f"  Valid trials: {len(valid_trials)}/{len(trials)}")
        
        # Step 3: Execute
        from nautilus_dolphin.mc import MCExecutor
        executor = MCExecutor(verbose=False)
        results = executor.execute_batch(valid_trials[:5])
        
        # Step 4: Store
        from nautilus_dolphin.mc import MCStore
        store = MCStore(output_dir=tmpdir)
        store.save_trial_results(results, batch_id=1)
        
        # Step 5: Query
        stats = store.get_corpus_stats()
        print(f"  Stored trials: {stats['total_trials']}")
        
        assert stats['total_trials'] > 0, "No trials stored"
        
        print("  [PASS] Integration test passed")
    
    return True


def run_all_tests():
    """Run all tests."""
    print("\n" + "="*70)
    print("MONTE CARLO SYSTEM TEST SUITE")
    print("="*70)
    
    tests = [
        ("Sampler", test_sampler),
        ("Validator", test_validator),
        ("Metrics", test_metrics),
        ("Store", test_store),
        ("Executor", test_executor),
        ("Runner", test_runner),
        ("CLI", test_cli),
        ("Integration", test_integration),
    ]
    
    passed = 0
    failed = 0
    
    for name, test_func in tests:
        try:
            if test_func():
                passed += 1
            else:
                failed += 1
                print(f"  [FAIL] {name} test failed")
        except Exception as e:
            failed += 1
            print(f"  [FAIL] {name} test failed with exception: {e}")
            import traceback
            traceback.print_exc()
    
    print("\n" + "="*70)
    print("TEST SUMMARY")
    print("="*70)
    print(f"Passed: {passed}/{len(tests)}")
    print(f"Failed: {failed}/{len(tests)}")
    
    if failed == 0:
        print("\n[OK] ALL TESTS PASSED!")
        return 0
    else:
        print("\n[FAIL] SOME TESTS FAILED")
        return 1


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