"""
Test script for Protocol 10 v2.0 Adaptive Characterization.

This script demonstrates the new adaptive characterization feature that
intelligently determines when enough landscape exploration has been done.

Expected behavior:
- Simple problems: Stop at ~10-15 trials
- Complex problems: Continue to ~20-30 trials
"""

import numpy as np
import optuna
from pathlib import Path
from optimization_engine.processors.adaptive_characterization import CharacterizationStoppingCriterion
from optimization_engine.reporting.landscape_analyzer import LandscapeAnalyzer


def simple_smooth_function(trial):
    """Simple smooth quadratic function (should stop early ~10-15 trials)."""
    x = trial.suggest_float('x', -10, 10)
    y = trial.suggest_float('y', -10, 10)

    # Simple quadratic bowl
    return (x - 3)**2 + (y + 2)**2


def complex_multimodal_function(trial):
    """Complex multimodal function (should need more trials ~20-30)."""
    x = trial.suggest_float('x', -5, 5)
    y = trial.suggest_float('y', -5, 5)

    # Rastrigin function (multimodal, many local minima)
    A = 10
    n = 2
    return A * n + ((x**2 - A * np.cos(2 * np.pi * x)) +
                    (y**2 - A * np.cos(2 * np.pi * y)))


def test_adaptive_characterization(
    objective_function,
    function_name: str,
    expected_trials_range: tuple
):
    """Test adaptive characterization on a given function."""

    print(f"\n{'='*70}")
    print(f"  TESTING: {function_name}")
    print(f"  Expected trials: {expected_trials_range[0]}-{expected_trials_range[1]}")
    print(f"{'='*70}\n")

    # Setup tracking directory
    tracking_dir = Path(f"test_results/adaptive_char_{function_name.lower().replace(' ', '_')}")
    tracking_dir.mkdir(parents=True, exist_ok=True)

    # Create components
    analyzer = LandscapeAnalyzer(min_trials_for_analysis=10)
    stopping_criterion = CharacterizationStoppingCriterion(
        min_trials=10,
        max_trials=30,
        confidence_threshold=0.85,
        check_interval=5,
        verbose=True,
        tracking_dir=tracking_dir
    )

    # Create study
    study = optuna.create_study(
        study_name=f"test_{function_name.lower().replace(' ', '_')}",
        direction='minimize',
        sampler=optuna.samplers.RandomSampler()
    )

    # Run adaptive characterization
    check_interval = 5
    while not stopping_criterion.should_stop(study):
        # Run batch of trials
        study.optimize(objective_function, n_trials=check_interval)

        # Analyze landscape
        landscape = analyzer.analyze(study)

        # Update stopping criterion
        if landscape.get('ready', False):
            completed_trials = [t for t in study.trials if t.state == optuna.trial.TrialState.COMPLETE]
            stopping_criterion.update(landscape, len(completed_trials))

    # Print results
    completed_trials = [t for t in study.trials if t.state == optuna.trial.TrialState.COMPLETE]
    actual_trials = len(completed_trials)

    print(stopping_criterion.get_summary_report())

    # Verify expectation
    in_range = expected_trials_range[0] <= actual_trials <= expected_trials_range[1]
    status = "PASS" if in_range else "FAIL"

    print(f"\n{'='*70}")
    print(f"  RESULT: {status}")
    print(f"  Actual trials: {actual_trials}")
    print(f"  Expected range: {expected_trials_range[0]}-{expected_trials_range[1]}")
    print(f"  In range: {'YES' if in_range else 'NO'}")
    print(f"  Stop reason: {stopping_criterion.stop_reason}")
    print(f"  Final confidence: {stopping_criterion.final_confidence:.1%}")
    print(f"{'='*70}\n")

    return {
        'function': function_name,
        'expected_range': expected_trials_range,
        'actual_trials': actual_trials,
        'in_range': in_range,
        'stop_reason': stopping_criterion.stop_reason,
        'confidence': stopping_criterion.final_confidence
    }


def main():
    """Run all adaptive characterization tests."""

    print("\n" + "="*70)
    print("  PROTOCOL 10 v2.0: ADAPTIVE CHARACTERIZATION TESTS")
    print("="*70)

    results = []

    # Test 1: Simple smooth function (should stop early)
    result1 = test_adaptive_characterization(
        objective_function=simple_smooth_function,
        function_name="Simple Smooth Quadratic",
        expected_trials_range=(10, 20)
    )
    results.append(result1)

    # Test 2: Complex multimodal function (should need more trials)
    result2 = test_adaptive_characterization(
        objective_function=complex_multimodal_function,
        function_name="Complex Multimodal (Rastrigin)",
        expected_trials_range=(15, 30)
    )
    results.append(result2)

    # Summary
    print("\n" + "="*70)
    print("  TEST SUMMARY")
    print("="*70)

    for result in results:
        status = "PASS" if result['in_range'] else "FAIL"
        print(f"\n  [{status}] {result['function']}")
        print(f"        Expected: {result['expected_range'][0]}-{result['expected_range'][1]} trials")
        print(f"        Actual: {result['actual_trials']} trials")
        print(f"        Confidence: {result['confidence']:.1%}")

    # Overall pass/fail
    all_passed = all(r['in_range'] for r in results)
    overall_status = "ALL TESTS PASSED" if all_passed else "SOME TESTS FAILED"

    print(f"\n{'='*70}")
    print(f"  {overall_status}")
    print(f"{'='*70}\n")


if __name__ == "__main__":
    main()