"""
Bracket Stiffness Optimization - Intelligent Optimizer (Protocol 10)
====================================================================

Run multi-objective optimization using Protocol 10: Intelligent Multi-Strategy
Optimization (IMSO) to:
1. Maximize stiffness (k = F/δ)
2. Minimize mass (constraint: ≤ 0.2 kg)

Design Variables:
- support_angle: 30° - 90°
- tip_thickness: 15mm - 40mm

Usage:
    python run_optimization.py [--trials N] [--dashboard]
"""

import sys
import json
import argparse
from pathlib import Path
from datetime import datetime

# Add project root to path
project_root = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(project_root))

import optuna
from optimization_engine.core.intelligent_optimizer import IntelligentOptimizer
from optimization_engine.nx.solver import NXSolver
from bracket_stiffness_extractor import BracketStiffnessExtractor

# Import central configuration
import config as atomizer_config


def load_config(config_file: Path) -> dict:
    """Load optimization configuration from JSON."""
    with open(config_file, 'r') as f:
        return json.load(f)


def create_objective_function(config: dict, study_dir: Path):
    """
    Create the objective function for bracket optimization.

    Returns a function that takes an Optuna trial and returns objectives.
    """

    def objective(trial: optuna.Trial) -> tuple:
        """
        Optimization objective function.

        Args:
            trial: Optuna trial object

        Returns:
            (stiffness_neg, mass): Tuple of objectives
                - stiffness_neg: Negative stiffness (for minimization)
                - mass: Mass in kg
        """
        # Sample design variables
        design_vars = {}
        for dv in config['design_variables']:
            design_vars[dv['name']] = trial.suggest_float(
                dv['name'],
                dv['min'],
                dv['max']
            )

        print(f"\n{'='*60}")
        print(f"Trial #{trial.number}")
        print(f"{'='*60}")
        print(f"Design Variables:")
        for name, value in design_vars.items():
            print(f"  {name}: {value:.3f}")

        # Model paths
        model_dir = study_dir / "1_setup" / "model"
        model_file = model_dir / "Bracket.prt"
        sim_file = model_dir / "Bracket_sim1.sim"

        # Initialize NX solver using central config
        nx_solver = NXSolver(
            nastran_version=atomizer_config.NX_VERSION,
            timeout=atomizer_config.NASTRAN_TIMEOUT,
            use_journal=True,
            enable_session_management=True,
            study_name="bracket_stiffness_optimization_V2"
        )

        # Run simulation with design variable updates
        print(f"\nRunning simulation with updated design variables...")
        try:
            result = nx_solver.run_simulation(
                sim_file=sim_file,
                working_dir=model_dir,
                expression_updates=design_vars,
                solution_name=None  # Solve all solutions in the .sim file
            )

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

        except Exception as e:
            print(f"ERROR: Simulation failed: {e}")
            raise optuna.exceptions.TrialPruned()

        # Step 3: Extract results (stiffness and mass)
        print(f"Extracting results...")
        extractor = BracketStiffnessExtractor(model_dir=model_dir)

        try:
            results = extractor.extract_results()
        except Exception as e:
            print(f"ERROR: Extraction failed: {e}")
            raise optuna.exceptions.TrialPruned()

        stiffness = results['stiffness']  # N/mm
        mass = results['mass']  # kg
        mass_g = results['mass_g']  # grams

        # Check constraint: mass ≤ 0.2 kg
        mass_limit = config['constraints'][0]['value']
        constraint_satisfied = mass <= mass_limit

        # Store constraint status and all metrics
        trial.set_user_attr("constraint_satisfied", constraint_satisfied)
        trial.set_user_attr("mass_limit", mass_limit)
        trial.set_user_attr("mass_violation", max(0, mass - mass_limit))
        trial.set_user_attr("displacement", results['displacement'])
        trial.set_user_attr("force", results['force'])
        trial.set_user_attr("compliance", results['compliance'])
        trial.set_user_attr("mass_g", mass_g)

        if not constraint_satisfied:
            print(f"\n[!] CONSTRAINT VIOLATED: Mass {mass:.6f} kg > {mass_limit} kg")
            print(f"  Trial will be kept for surrogate modeling but not eligible for Pareto front")

        print(f"\n[OK] Trial Complete")
        print(f"  Stiffness: {stiffness:.2f} N/mm")
        print(f"  Mass: {mass:.6f} kg ({mass_g:.2f} g)")
        print(f"  Constraint: {'[OK] SATISFIED' if constraint_satisfied else '[X] VIOLATED'}")

        # Return tuple of objectives
        # Note: Optuna minimizes by default, so return negative for maximization
        return -stiffness, mass  # Maximize stiffness, minimize mass

    return objective


def run_optimization(
    config_file: Path,
    study_dir: Path,
    n_trials: int = None,
    dashboard: bool = False
):
    """
    Run bracket stiffness optimization using Protocol 10.

    Args:
        config_file: Path to optimization_config.json
        study_dir: Path to study directory
        n_trials: Number of trials (overrides config)
        dashboard: Enable real-time dashboard
    """
    # Load configuration
    config = load_config(config_file)
    study_name = config['study_name']

    if n_trials is None:
        n_trials = config['optimization_settings']['n_trials']

    # Setup results directory
    results_dir = study_dir / config['output_settings']['results_dir']
    results_dir.mkdir(exist_ok=True)

    # Extract design variable bounds
    design_variables = {
        dv['name']: (dv['min'], dv['max'])
        for dv in config['design_variables']
    }

    print(f"\n{'='*60}")
    print(f"BRACKET STIFFNESS OPTIMIZATION - PROTOCOL 10")
    print(f"{'='*60}")
    print(f"Study: {study_name}")
    print(f"Trials: {n_trials}")
    print(f"Objectives: Maximize stiffness, Minimize mass")
    print(f"Constraint: Mass <= 0.2 kg (kept for surrogate, filtered from Pareto)")
    print(f"Optimizer: Intelligent Multi-Strategy (Protocol 10)")
    print(f"Results: {results_dir}")
    print(f"{'='*60}\n")

    # Create intelligent optimizer
    intelligent_config = {
        'intelligent_optimization': {
            'enabled': True,
            'min_analysis_trials': 10,
            'stagnation_window': 10,
            'min_improvement_threshold': 0.001,
            'target_value': None,  # No specific target for multi-objective
            'enable_adaptive_surrogate': True,
            'enable_strategy_switching': True
        },
        'optimization_settings': config['optimization_settings'],
        'output_settings': config['output_settings']
    }

    optimizer = IntelligentOptimizer(
        study_name=study_name,
        study_dir=results_dir,
        config=intelligent_config,
        verbose=True
    )

    # Create objective function
    objective_fn = create_objective_function(config, study_dir)

    # Run optimization
    start_time = datetime.now()

    try:
        results = optimizer.optimize(
            objective_function=objective_fn,
            design_variables=design_variables,
            n_trials=n_trials,
            target_value=None,  # Multi-objective, no single target
            directions=["minimize", "minimize"]  # Minimize -stiffness (=maximize stiffness), minimize mass
        )
    except KeyboardInterrupt:
        print("\n\nOptimization interrupted by user.")
        results = None

    end_time = datetime.now()
    elapsed = (end_time - start_time).total_seconds()

    # Print results summary
    print(f"\n{'='*60}")
    print(f"OPTIMIZATION COMPLETE")
    print(f"{'='*60}")

    if results:
        print(f"\n[BEST] Best Solution Found:")
        print(f"  Stiffness: {-results['best_value'][0]:.2f} N/mm")  # Convert back from negative
        print(f"  Mass: {results['best_value'][1]:.6f} kg")
        print(f"  Parameters: {results['best_params']}")
        print(f"\n[STRATEGY] Strategy Performance:")
        print(f"  Final Strategy: {results.get('strategy_used', 'N/A')}")
        if 'landscape_analysis' in results and results['landscape_analysis'] is not None:
            print(f"  Landscape Type: {results['landscape_analysis'].get('landscape_type', 'N/A')}")

    # Access the Optuna study for detailed analysis
    study = optimizer.study

    if study:
        completed_trials = [t for t in study.trials if t.state == optuna.trial.TrialState.COMPLETE]
        pruned_trials = [t for t in study.trials if t.state == optuna.trial.TrialState.PRUNED]
        failed_trials = [t for t in study.trials if t.state == optuna.trial.TrialState.FAIL]

        # Count constraint violations
        constraint_violated = [t for t in completed_trials if not t.user_attrs.get("constraint_satisfied", True)]
        constraint_satisfied = [t for t in completed_trials if t.user_attrs.get("constraint_satisfied", True)]

        print(f"\n[STATS] Trial Statistics:")
        print(f"  Total trials: {len(study.trials)}")
        print(f"  Completed: {len(completed_trials)}")
        print(f"    [OK] Feasible (constraint satisfied): {len(constraint_satisfied)}")
        print(f"    [X] Infeasible (constraint violated): {len(constraint_violated)}")
        print(f"  Pruned: {len(pruned_trials)}")
        print(f"  Failed: {len(failed_trials)}")
        print(f"  Elapsed time: {elapsed:.1f} seconds ({elapsed/60:.1f} minutes)")

        # Get best trials (Pareto front) - filtered to only feasible solutions
        try:
            # Get all Pareto-optimal trials first
            all_best_trials = study.best_trials

            # Filter to only feasible solutions (constraint satisfied)
            feasible_best_trials = [
                t for t in all_best_trials
                if t.user_attrs.get("constraint_satisfied", True)
            ]

            print(f"\n[PARETO] Pareto Front Analysis:")
            print(f"  Pareto Front (all trials): {len(all_best_trials)} solutions")
            print(f"  Pareto Front (feasible only): {len(feasible_best_trials)} solutions")

            if len(feasible_best_trials) > 0:
                print(f"\n  Top 5 Feasible Solutions:")
                print(f"  {'Trial':<8} {'Stiffness':<15} {'Mass':<15} {'Angle':<12} {'Thickness':<12}")
                print(f"  {'-' * 70}")

                for i, trial in enumerate(feasible_best_trials[:5]):
                    stiff_neg, mass = trial.values
                    stiffness = -stiff_neg  # Convert back to positive
                    angle = trial.params.get('support_angle', 0)
                    thickness = trial.params.get('tip_thickness', 0)
                    print(f"  {trial.number:<8} {stiffness:<15.2f} {mass:<15.6f} {angle:<12.2f} {thickness:<12.2f}")
            else:
                print(f"\n  [!] Warning: No feasible solutions found in Pareto front!")

        except Exception as e:
            print(f"\n  Note: Could not retrieve Pareto front: {e}")

        # Save summary
        summary_file = results_dir / "optimization_summary.json"
        summary = {
            "study_name": study_name,
            "optimizer": "Protocol 10 - Intelligent Multi-Strategy",
            "n_trials": len(study.trials),
            "completed_trials": len(completed_trials),
            "feasible_trials": len(constraint_satisfied),
            "infeasible_trials": len(constraint_violated),
            "pruned_trials": len(pruned_trials),
            "failed_trials": len(failed_trials),
            "elapsed_seconds": elapsed,
            "pareto_front_all": len(all_best_trials) if 'all_best_trials' in locals() else 0,
            "pareto_front_feasible": len(feasible_best_trials) if 'feasible_best_trials' in locals() else 0,
            "best_solution": results if results else None,
            "timestamp": datetime.now().isoformat()
        }

        with open(summary_file, 'w') as f:
            json.dump(summary, f, indent=2)

        print(f"\n[SAVED] Summary saved to: {summary_file}")

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


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run bracket stiffness optimization with Protocol 10")
    parser.add_argument('--trials', type=int, default=None, help='Number of trials (default: from config)')
    parser.add_argument('--dashboard', action='store_true', help='Enable real-time dashboard')

    args = parser.parse_args()

    # Paths
    study_dir = Path(__file__).parent
    config_file = study_dir / "optimization_config.json"

    if not config_file.exists():
        print(f"ERROR: Configuration file not found: {config_file}")
        sys.exit(1)

    # Run optimization
    run_optimization(
        config_file=config_file,
        study_dir=study_dir,
        n_trials=args.trials,
        dashboard=args.dashboard
    )