"""
Create a new circular plate frequency tuning study using Protocol 10.

This script creates a complete study configured for intelligent multi-strategy
optimization (IMSO) to test the self-tuning framework.
"""

import json
import shutil
from pathlib import Path

# Study configuration
STUDY_NAME = "circular_plate_frequency_tuning_intelligent_optimizer"
BASE_DIR = Path(__file__).parent
STUDIES_DIR = BASE_DIR / "studies"
STUDY_DIR = STUDIES_DIR / STUDY_NAME

# Source model files (copy from examples)
SOURCE_MODEL_DIR = BASE_DIR / "examples" / "Models" / "Circular Plate"

def create_study_structure():
    """Create complete study directory structure."""
    print(f"Creating study: {STUDY_NAME}")

    # Create directories
    setup_dir = STUDY_DIR / "1_setup"
    model_dir = setup_dir / "model"
    results_dir = STUDY_DIR / "2_results"
    reports_dir = STUDY_DIR / "3_reports"

    for directory in [setup_dir, model_dir, results_dir, reports_dir]:
        directory.mkdir(parents=True, exist_ok=True)
        print(f"  Created: {directory.relative_to(BASE_DIR)}")

    return setup_dir, model_dir, results_dir, reports_dir


def copy_model_files(model_dir):
    """Copy model files from examples."""
    print("\nCopying model files...")

    model_files = [
        "Circular_Plate.prt",
        "Circular_Plate_sim1.sim",
        "Circular_Plate_fem1.fem",
        "Circular_Plate_fem1_i.prt"
    ]

    for filename in model_files:
        source = SOURCE_MODEL_DIR / filename
        dest = model_dir / filename

        if source.exists():
            shutil.copy2(source, dest)
            print(f"  Copied: {filename}")
        else:
            print(f"  WARNING: Source not found: {filename}")

    return list(model_dir.glob("*"))


def create_workflow_config(setup_dir):
    """Create workflow configuration."""
    print("\nCreating workflow configuration...")

    workflow = {
        "study_name": STUDY_NAME,
        "optimization_request": "Tune the first natural frequency mode to exactly 115 Hz using intelligent multi-strategy optimization",
        "design_variables": [
            {
                "parameter": "inner_diameter",
                "bounds": [50, 150],
                "units": "mm",
                "description": "Inner diameter of circular plate"
            },
            {
                "parameter": "plate_thickness",
                "bounds": [2, 10],
                "units": "mm",
                "description": "Thickness of circular plate"
            }
        ],
        "objectives": [
            {
                "name": "frequency_error",
                "goal": "minimize",
                "extraction": {
                    "action": "extract_first_natural_frequency",
                    "params": {
                        "mode_number": 1,
                        "target_frequency": 115.0
                    }
                }
            }
        ],
        "constraints": [
            {
                "name": "frequency_tolerance",
                "type": "less_than",
                "threshold": 0.1,
                "description": "Error from target must be less than 0.1 Hz"
            }
        ]
    }

    workflow_file = setup_dir / "workflow_config.json"
    with open(workflow_file, 'w', encoding='utf-8') as f:
        json.dump(workflow, f, indent=2)

    print(f"  Saved: {workflow_file.relative_to(BASE_DIR)}")
    return workflow


def create_optimization_config(setup_dir):
    """Create Protocol 10 optimization configuration."""
    print("\nCreating Protocol 10 optimization configuration...")

    config = {
        "_description": "Protocol 10: Intelligent Multi-Strategy Optimization - Circular Plate Test",
        "_version": "1.0",

        "study_name": STUDY_NAME,
        "direction": "minimize",

        "intelligent_optimization": {
            "_description": "Protocol 10 - Automatic landscape analysis and strategy selection",
            "enabled": True,

            "characterization_trials": 15,
            "stagnation_window": 10,
            "min_improvement_threshold": 0.001,
            "min_analysis_trials": 10,
            "reanalysis_interval": 15,

            "strategy_preferences": {
                "prefer_cmaes_for_smooth": True,
                "prefer_tpe_for_multimodal": True,
                "enable_hybrid_strategies": False
            }
        },

        "sampler": {
            "_description": "Fallback sampler if Protocol 10 disabled",
            "type": "TPESampler",
            "params": {
                "n_startup_trials": 10,
                "n_ei_candidates": 24,
                "multivariate": True,
                "warn_independent_sampling": True
            }
        },

        "pruner": {
            "type": "MedianPruner",
            "params": {
                "n_startup_trials": 5,
                "n_warmup_steps": 0
            }
        },

        "adaptive_strategy": {
            "_description": "Protocol 8 - Adaptive exploitation based on surrogate confidence",
            "enabled": True,
            "min_confidence_for_exploitation": 0.65,
            "min_trials_for_confidence": 15,
            "target_confidence_metrics": {
                "convergence_weight": 0.4,
                "coverage_weight": 0.3,
                "stability_weight": 0.3
            }
        },

        "trials": {
            "n_trials": 100,
            "timeout": None,
            "catch": []
        },

        "reporting": {
            "auto_generate_plots": True,
            "include_optuna_visualizations": True,
            "include_confidence_report": True,
            "include_strategy_performance": True,
            "save_intelligence_report": True
        },

        "verbosity": {
            "print_landscape_report": True,
            "print_strategy_recommendation": True,
            "print_phase_transitions": True,
            "print_confidence_updates": True,
            "log_to_file": True
        },

        "optimization_notes": "Protocol 10 Test: Atomizer will automatically characterize the circular plate problem, select the best optimization algorithm (TPE, CMA-ES, or GP-BO), and adapt strategy if stagnation is detected. Expected: smooth_unimodal landscape → CMA-ES recommendation."
    }

    config_file = setup_dir / "optimization_config.json"
    with open(config_file, 'w', encoding='utf-8') as f:
        json.dump(config, f, indent=2)

    print(f"  Saved: {config_file.relative_to(BASE_DIR)}")
    return config


def create_runner_script(study_dir, workflow, config):
    """Create optimization runner using Protocol 10."""
    print("\nCreating Protocol 10 optimization runner...")

    runner_code = '''"""
Intelligent Multi-Strategy Optimization Runner
Study: circular_plate_frequency_tuning_intelligent_optimizer

This runner uses Protocol 10 (IMSO) to automatically:
1. Characterize the optimization landscape
2. Select the best optimization algorithm
3. Adapt strategy dynamically if needed

Generated: 2025-11-19
Protocol: 10 (Intelligent Multi-Strategy Optimization)
"""

import sys
import json
import optuna
from pathlib import Path

# Add optimization engine to path
sys.path.insert(0, str(Path(__file__).parent.parent.parent))

from optimization_engine.core.intelligent_optimizer import IntelligentOptimizer
from optimization_engine.nx.updater import NXParameterUpdater
from optimization_engine.nx.solver import NXSolver
from optimization_engine.extractors.frequency_extractor import extract_first_frequency
from optimization_engine.reporting.markdown_report import generate_markdown_report


def main():
    """Run Protocol 10 intelligent optimization."""

    # Setup paths
    study_dir = Path(__file__).parent
    setup_dir = study_dir / "1_setup"
    model_dir = setup_dir / "model"
    results_dir = study_dir / "2_results"
    reports_dir = study_dir / "3_reports"

    # Create directories
    results_dir.mkdir(exist_ok=True)
    reports_dir.mkdir(exist_ok=True)

    # Load configuration
    print("\\nLoading configuration...")
    with open(setup_dir / "workflow_config.json") as f:
        workflow = json.load(f)

    with open(setup_dir / "optimization_config.json") as f:
        opt_config = json.load(f)

    print(f"Study: {workflow['study_name']}")
    print(f"Protocol 10: {opt_config['intelligent_optimization']['enabled']}")

    # Model files
    prt_file = model_dir / "Circular_Plate.prt"
    sim_file = model_dir / "Circular_Plate_sim1.sim"

    # Initialize NX components
    updater = NXParameterUpdater(str(prt_file))
    solver = NXSolver()

    # Incremental history tracking
    history_file = results_dir / "optimization_history_incremental.json"
    history = []

    def objective(trial):
        """Objective function for optimization."""

        # Sample design variables
        inner_diameter = trial.suggest_float('inner_diameter', 50, 150)
        plate_thickness = trial.suggest_float('plate_thickness', 2, 10)

        params = {
            'inner_diameter': inner_diameter,
            'plate_thickness': plate_thickness
        }

        print(f"\\n  Trial #{trial.number}")
        print(f"    Inner Diameter: {inner_diameter:.4f} mm")
        print(f"    Plate Thickness: {plate_thickness:.4f} mm")

        # Update CAD model
        updater.update_expressions(params)

        # Run simulation (use discovered solution name from benchmarking)
        result = solver.run_simulation(str(sim_file), solution_name="Solution_Normal_Modes")

        if not result['success']:
            print(f"    Simulation FAILED: {result.get('error', 'Unknown error')}")
            raise optuna.TrialPruned()

        # Extract frequency
        op2_file = result['op2_file']
        frequency = extract_first_frequency(op2_file, mode_number=1)

        # Calculate objective (error from target)
        target_frequency = 115.0
        objective_value = abs(frequency - target_frequency)

        print(f"    Frequency: {frequency:.4f} Hz")
        print(f"    Target: {target_frequency:.4f} Hz")
        print(f"    Error: {objective_value:.4f} Hz")

        # Save to incremental history
        trial_data = {
            "trial_number": trial.number,
            "design_variables": params,
            "results": {"first_frequency": frequency},
            "objective": objective_value
        }

        history.append(trial_data)

        with open(history_file, 'w', encoding='utf-8') as f:
            json.dump(history, f, indent=2)

        return objective_value

    # Create intelligent optimizer
    print("\\n" + "="*70)
    print("  PROTOCOL 10: INTELLIGENT MULTI-STRATEGY OPTIMIZATION")
    print("="*70)

    optimizer = IntelligentOptimizer(
        study_name=workflow['study_name'],
        study_dir=results_dir,
        config=opt_config,
        verbose=True
    )

    # Extract design variable bounds
    design_vars = {
        var['parameter']: tuple(var['bounds'])
        for var in workflow['design_variables']
    }

    # Run optimization
    results = optimizer.optimize(
        objective_function=objective,
        design_variables=design_vars,
        n_trials=opt_config['trials']['n_trials'],
        target_value=115.0,
        tolerance=0.1
    )

    # Save intelligence report
    optimizer.save_intelligence_report()

    # Generate markdown report
    print("\\nGenerating optimization report...")

    # Load study for Optuna visualizations
    storage = f"sqlite:///{results_dir / 'study.db'}"
    study = optuna.load_study(
        study_name=workflow['study_name'],
        storage=storage
    )

    report = generate_markdown_report(
        history_file=history_file,
        target_value=115.0,
        tolerance=0.1,
        reports_dir=reports_dir,
        study=study
    )

    report_file = reports_dir / "OPTIMIZATION_REPORT.md"
    with open(report_file, 'w', encoding='utf-8') as f:
        f.write(report)

    print(f"\\nReport saved: {report_file}")

    # Print final summary
    print("\\n" + "="*70)
    print("  PROTOCOL 10 TEST COMPLETE")
    print("="*70)
    print(f"Best Frequency: {results['best_value'] + 115.0:.4f} Hz")
    print(f"Best Error: {results['best_value']:.4f} Hz")
    print(f"Best Parameters:")
    for param, value in results['best_params'].items():
        print(f"  {param}: {value:.4f}")

    if 'landscape_analysis' in results and results['landscape_analysis'].get('ready'):
        landscape = results['landscape_analysis']
        print(f"\\nLandscape Type: {landscape['landscape_type'].upper()}")
        print(f"Recommended Strategy: {results.get('final_strategy', 'N/A').upper()}")

    print("="*70)


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

    runner_file = study_dir / "run_optimization.py"
    with open(runner_file, 'w', encoding='utf-8') as f:
        f.write(runner_code)

    print(f"  Saved: {runner_file.relative_to(BASE_DIR)}")
    return runner_file


def create_readme(study_dir):
    """Create README for the study."""
    print("\nCreating README...")

    readme = f"""# {STUDY_NAME}

**Protocol 10 Test Study** - Intelligent Multi-Strategy Optimization

## Overview

This study tests Atomizer's Protocol 10 (IMSO) framework on a circular plate frequency tuning problem.

**Goal**: Tune first natural frequency to exactly 115 Hz

**Protocol 10 Features Tested**:
- Automatic landscape characterization (smoothness, multimodality, correlation)
- Intelligent strategy selection (TPE, CMA-ES, or GP-BO)
- Dynamic strategy switching based on stagnation detection
- Comprehensive decision logging for transparency

## Expected Behavior

### Stage 1: Landscape Characterization (Trials 1-15)
- Random exploration to gather data
- Analyze problem characteristics
- Expected classification: `smooth_unimodal` with strong parameter correlation

### Stage 2: Strategy Selection (Trial 15)
- Expected recommendation: **CMA-ES** (92% confidence)
- Reasoning: "Smooth unimodal with strong correlation - CMA-ES converges quickly"

### Stage 3: Adaptive Optimization (Trials 16-100)
- Run with CMA-ES sampler
- Monitor for stagnation
- Switch strategies if needed (unlikely for this problem)

## Study Structure

```
{STUDY_NAME}/
├── 1_setup/
│   ├── model/                    # CAD and simulation files
│   ├── workflow_config.json      # Optimization goals
│   └── optimization_config.json  # Protocol 10 configuration
├── 2_results/
│   ├── study.db                  # Optuna database
│   ├── optimization_history_incremental.json
│   └── intelligent_optimizer/    # Protocol 10 tracking
│       ├── strategy_transitions.json
│       ├── strategy_performance.json
│       └── intelligence_report.json
└── 3_reports/
    └── OPTIMIZATION_REPORT.md    # Final report with visualizations
```

## Running the Optimization

```bash
# Activate environment
conda activate test_env

# Run optimization
python run_optimization.py
```

## What to Look For

### Console Output

**Landscape Analysis Report**:
```
======================================================================
  LANDSCAPE ANALYSIS REPORT
======================================================================
  Type: SMOOTH_UNIMODAL
  Smoothness: 0.7X (smooth)
  Multimodal: NO (1 modes)
  Parameter Correlation: 0.6X (strong)
```

**Strategy Recommendation**:
```
======================================================================
  STRATEGY RECOMMENDATION
======================================================================
  Recommended: CMAES
  Confidence: 92.0%
  Reasoning: Smooth unimodal with strong correlation - CMA-ES converges quickly
```

**Phase Transitions** (if any):
```
======================================================================
  STRATEGY TRANSITION
======================================================================
  Trial #45
  TPE → CMAES
  Reason: Stagnation detected
```

### Intelligence Report

Check `2_results/intelligent_optimizer/intelligence_report.json` for:
- Complete landscape analysis
- Strategy recommendation reasoning
- All transition events
- Performance breakdown by strategy

### Optimization Report

Check `3_reports/OPTIMIZATION_REPORT.md` for:
- Best result (should be < 0.1 Hz error)
- Convergence plots
- Optuna visualizations
- (Future: Protocol 10 analysis section)

## Expected Results

**Baseline (TPE only)**: ~160 trials to achieve 0.18 Hz error

**Protocol 10 (Intelligent)**: ~60-80 trials to achieve < 0.1 Hz error

**Improvement**: 40-50% faster convergence by selecting optimal algorithm

## Configuration

See [`optimization_config.json`](1_setup/optimization_config.json) for full Protocol 10 settings.

Key parameters:
- `characterization_trials`: 15 (initial exploration)
- `stagnation_window`: 10 (trials to check for stagnation)
- `min_improvement_threshold`: 0.001 (0.1% minimum improvement)

## References

- [PROTOCOL.md](../../PROTOCOL.md) - Complete Protocol 10 documentation
- [Protocol 10 Implementation Summary](../../docs/PROTOCOL_10_IMPLEMENTATION_SUMMARY.md)
- [Example Configuration](../../examples/optimization_config_protocol10.json)

---

*Study created: 2025-11-19*
*Protocol: 10 (Intelligent Multi-Strategy Optimization)*
"""

    readme_file = study_dir / "README.md"
    with open(readme_file, 'w', encoding='utf-8') as f:
        f.write(readme)

    print(f"  Saved: {readme_file.relative_to(BASE_DIR)}")


def main():
    """Create complete Protocol 10 test study."""
    print("\n" + "="*70)
    print("  CREATING PROTOCOL 10 TEST STUDY")
    print("="*70)

    # Create structure
    setup_dir, model_dir, results_dir, reports_dir = create_study_structure()

    # Copy model files
    model_files = copy_model_files(model_dir)

    # Create configurations
    workflow = create_workflow_config(setup_dir)
    config = create_optimization_config(setup_dir)

    # Create runner
    runner_file = create_runner_script(STUDY_DIR, workflow, config)

    # Create README
    create_readme(STUDY_DIR)

    print("\n" + "="*70)
    print("  STUDY CREATION COMPLETE")
    print("="*70)
    print(f"\nStudy directory: {STUDY_DIR.relative_to(BASE_DIR)}")
    print(f"\nTo run optimization:")
    print(f"  cd {STUDY_DIR.relative_to(BASE_DIR)}")
    print(f"  python run_optimization.py")
    print("\n" + "="*70)


if __name__ == "__main__":
    main()