#!/usr/bin/env python3
"""
M1 Mirror Pure NSGA-II FEA Optimization V13
=============================================

Pure multi-objective optimization with NSGA-II sampler and FEA only.
No neural surrogate - every trial is a real FEA evaluation.

Key Features:
1. NSGA-II sampler for true multi-objective Pareto optimization
2. Seeds from V11 + V12 FEA trials (~110+ prior trials)
3. No surrogate bias - ground truth only
4. 3 objectives: rel_rms_40_vs_20, rel_rms_60_vs_20, mfg_90

Usage:
    python run_optimization.py --start
    python run_optimization.py --start --trials 50
    python run_optimization.py --start --trials 50 --resume

For 8-hour overnight run (~55 trials at 8-9 min/trial):
    python run_optimization.py --start --trials 55
"""

import sys
import os
import json
import time
import argparse
import logging
import sqlite3
import shutil
import re
from pathlib import Path
from typing import Dict, List, Tuple, Optional, Any
from dataclasses import dataclass, field
from datetime import datetime
import numpy as np

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

import optuna
from optuna.samplers import NSGAIISampler

# Atomizer imports
from optimization_engine.nx_solver import NXSolver
from optimization_engine.utils import ensure_nx_running
from optimization_engine.extractors import ZernikeExtractor

# ============================================================================
# Paths
# ============================================================================

STUDY_DIR = Path(__file__).parent
SETUP_DIR = STUDY_DIR / "1_setup"
ITERATIONS_DIR = STUDY_DIR / "2_iterations"
RESULTS_DIR = STUDY_DIR / "3_results"
CONFIG_PATH = SETUP_DIR / "optimization_config.json"

# Source studies for seeding
V11_DB = STUDY_DIR.parent / "m1_mirror_adaptive_V11" / "3_results" / "study.db"
V12_DB = STUDY_DIR.parent / "m1_mirror_adaptive_V12" / "3_results" / "study.db"

# Ensure directories exist
ITERATIONS_DIR.mkdir(exist_ok=True)
RESULTS_DIR.mkdir(exist_ok=True)

# Logging
LOG_FILE = RESULTS_DIR / "optimization.log"
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s | %(levelname)-8s | %(message)s',
    handlers=[
        logging.StreamHandler(sys.stdout),
        logging.FileHandler(LOG_FILE, mode='a')
    ]
)
logger = logging.getLogger(__name__)


# ============================================================================
# Objective names
# ============================================================================

OBJ_NAMES = [
    'rel_filtered_rms_40_vs_20',
    'rel_filtered_rms_60_vs_20',
    'mfg_90_optician_workload'
]

DESIGN_VAR_NAMES = [
    'lateral_inner_angle', 'lateral_outer_angle', 'lateral_outer_pivot',
    'lateral_inner_pivot', 'lateral_middle_pivot', 'lateral_closeness',
    'whiffle_min', 'whiffle_outer_to_vertical', 'whiffle_triangle_closeness',
    'blank_backface_angle', 'inner_circular_rib_dia'
]


# ============================================================================
# Prior Data Loader
# ============================================================================

def load_fea_trials_from_db(db_path: Path, label: str) -> List[Dict]:
    """Load FEA trials from an Optuna database."""
    if not db_path.exists():
        logger.warning(f"{label} database not found: {db_path}")
        return []

    fea_data = []
    conn = sqlite3.connect(str(db_path))

    try:
        cursor = conn.cursor()

        cursor.execute('''
            SELECT trial_id, number FROM trials
            WHERE state = 'COMPLETE'
        ''')
        trials = cursor.fetchall()

        for trial_id, trial_num in trials:
            # Get user attributes
            cursor.execute('''
                SELECT key, value_json FROM trial_user_attributes
                WHERE trial_id = ?
            ''', (trial_id,))
            attrs = {row[0]: json.loads(row[1]) for row in cursor.fetchall()}

            # Check if FEA trial (source contains 'FEA')
            source = attrs.get('source', 'FEA')
            if 'FEA' not in source:
                continue  # Skip NN trials

            # Get params
            cursor.execute('''
                SELECT param_name, param_value FROM trial_params
                WHERE trial_id = ?
            ''', (trial_id,))
            params = {name: float(value) for name, value in cursor.fetchall()}

            if not params:
                continue

            # Get objectives (stored as individual attributes or in 'objectives')
            objectives = {}
            if 'objectives' in attrs:
                objectives = attrs['objectives']
            else:
                # Try individual attributes
                for obj_name in OBJ_NAMES:
                    if obj_name in attrs:
                        objectives[obj_name] = attrs[obj_name]

            if all(k in objectives for k in OBJ_NAMES):
                fea_data.append({
                    'trial_num': trial_num,
                    'params': params,
                    'objectives': objectives,
                    'source': f'{label}_{source}'
                })

    except Exception as e:
        logger.error(f"Error loading {label} data: {e}")
    finally:
        conn.close()

    logger.info(f"Loaded {len(fea_data)} FEA trials from {label}")
    return fea_data


def load_all_prior_fea_data() -> List[Dict]:
    """Load FEA trials from V11 and V12."""
    all_data = []

    # V11 data
    v11_data = load_fea_trials_from_db(V11_DB, "V11")
    all_data.extend(v11_data)

    # V12 data
    v12_data = load_fea_trials_from_db(V12_DB, "V12")
    all_data.extend(v12_data)

    logger.info(f"Total prior FEA trials: {len(all_data)}")
    return all_data


# ============================================================================
# FEA Runner
# ============================================================================

class FEARunner:
    """Runs actual FEA simulations."""

    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.nx_solver = None
        self.nx_manager = None
        self.master_model_dir = SETUP_DIR / "model"

    def setup(self):
        """Setup NX and solver."""
        logger.info("Setting up NX session...")

        study_name = self.config.get('study_name', 'm1_mirror_adaptive_V13')

        try:
            self.nx_manager, nx_was_started = ensure_nx_running(
                session_id=study_name,
                auto_start=True,
                start_timeout=120
            )
            logger.info("NX session ready" + (" (started)" if nx_was_started else " (existing)"))
        except Exception as e:
            logger.error(f"Failed to setup NX: {e}")
            raise

        # Initialize solver
        nx_settings = self.config.get('nx_settings', {})
        nx_install_dir = nx_settings.get('nx_install_path', 'C:\\Program Files\\Siemens\\NX2506')
        version_match = re.search(r'NX(\d+)', nx_install_dir)
        nastran_version = version_match.group(1) if version_match else "2506"

        self.nx_solver = NXSolver(
            master_model_dir=str(self.master_model_dir),
            nx_install_dir=nx_install_dir,
            nastran_version=nastran_version,
            timeout=nx_settings.get('simulation_timeout_s', 600),
            use_iteration_folders=True,
            study_name="m1_mirror_adaptive_V13"
        )

    def run_fea(self, params: Dict[str, float], trial_num: int) -> Optional[Dict]:
        """Run FEA and extract objectives."""
        if self.nx_solver is None:
            self.setup()

        logger.info(f"  [FEA {trial_num}] Running simulation...")

        expressions = {var['expression_name']: params[var['name']]
                       for var in self.config['design_variables']}

        iter_folder = self.nx_solver.create_iteration_folder(
            iterations_base_dir=ITERATIONS_DIR,
            iteration_number=trial_num,
            expression_updates=expressions
        )

        try:
            nx_settings = self.config.get('nx_settings', {})
            sim_file = iter_folder / nx_settings.get('sim_file', 'ASSY_M1_assyfem1_sim1.sim')

            t_start = time.time()

            result = self.nx_solver.run_simulation(
                sim_file=sim_file,
                working_dir=iter_folder,
                expression_updates=expressions,
                solution_name=nx_settings.get('solution_name', 'Solution 1'),
                cleanup=False
            )

            solve_time = time.time() - t_start

            if not result['success']:
                logger.error(f"  [FEA {trial_num}] Solve failed: {result.get('error')}")
                return None

            logger.info(f"  [FEA {trial_num}] Solved in {solve_time:.1f}s")

            # Extract objectives
            op2_path = Path(result['op2_file'])
            objectives = self._extract_objectives(op2_path)

            if objectives is None:
                return None

            logger.info(f"  [FEA {trial_num}] 40-20: {objectives['rel_filtered_rms_40_vs_20']:.2f} nm")
            logger.info(f"  [FEA {trial_num}] 60-20: {objectives['rel_filtered_rms_60_vs_20']:.2f} nm")
            logger.info(f"  [FEA {trial_num}] Mfg:   {objectives['mfg_90_optician_workload']:.2f} nm")

            return {
                'trial_num': trial_num,
                'params': params,
                'objectives': objectives,
                'source': 'FEA',
                'solve_time': solve_time
            }

        except Exception as e:
            logger.error(f"  [FEA {trial_num}] Error: {e}")
            import traceback
            traceback.print_exc()
            return None

    def _extract_objectives(self, op2_path: Path) -> Optional[Dict[str, float]]:
        """Extract objectives using ZernikeExtractor."""
        try:
            zernike_settings = self.config.get('zernike_settings', {})

            extractor = ZernikeExtractor(
                op2_path,
                bdf_path=None,
                displacement_unit=zernike_settings.get('displacement_unit', 'mm'),
                n_modes=zernike_settings.get('n_modes', 50),
                filter_orders=zernike_settings.get('filter_low_orders', 4)
            )

            ref = zernike_settings.get('reference_subcase', '2')

            rel_40 = extractor.extract_relative("3", ref)
            rel_60 = extractor.extract_relative("4", ref)
            rel_90 = extractor.extract_relative("1", ref)

            return {
                'rel_filtered_rms_40_vs_20': rel_40['relative_filtered_rms_nm'],
                'rel_filtered_rms_60_vs_20': rel_60['relative_filtered_rms_nm'],
                'mfg_90_optician_workload': rel_90['relative_rms_filter_j1to3']
            }

        except Exception as e:
            logger.error(f"Zernike extraction failed: {e}")
            return None

    def cleanup(self):
        """Cleanup NX session."""
        if self.nx_manager:
            if self.nx_manager.can_close_nx():
                self.nx_manager.close_nx_if_allowed()
            self.nx_manager.cleanup()


# ============================================================================
# NSGA-II Optimizer
# ============================================================================

class NSGA2Optimizer:
    """Pure FEA multi-objective optimizer with NSGA-II."""

    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.fea_runner = FEARunner(config)

        # Load prior data for seeding
        self.prior_data = load_all_prior_fea_data()

        # Database
        self.db_path = RESULTS_DIR / "study.db"
        self.storage = optuna.storages.RDBStorage(f'sqlite:///{self.db_path}')

        # State
        self.trial_count = 0
        self.best_pareto = []

    def _get_next_trial_number(self) -> int:
        """Get the next trial number based on existing iterations."""
        existing = list(ITERATIONS_DIR.glob("iter*"))
        if not existing:
            return 1
        max_num = max(int(p.name.replace("iter", "")) for p in existing)
        return max_num + 1

    def seed_from_prior(self, study: optuna.Study):
        """Seed the study with prior FEA trials."""
        if not self.prior_data:
            logger.warning("No prior data to seed from")
            return

        logger.info(f"Seeding study with {len(self.prior_data)} prior FEA trials...")

        for i, d in enumerate(self.prior_data):
            try:
                # Create a trial with the prior data
                distributions = {}
                for var in self.config['design_variables']:
                    if var.get('enabled', False):
                        distributions[var['name']] = optuna.distributions.FloatDistribution(
                            var['min'], var['max']
                        )

                # Create frozen trial
                frozen_trial = optuna.trial.create_trial(
                    params=d['params'],
                    distributions=distributions,
                    values=[
                        d['objectives']['rel_filtered_rms_40_vs_20'],
                        d['objectives']['rel_filtered_rms_60_vs_20'],
                        d['objectives']['mfg_90_optician_workload']
                    ],
                    user_attrs={
                        'source': d.get('source', 'prior_FEA'),
                        'rel_filtered_rms_40_vs_20': d['objectives']['rel_filtered_rms_40_vs_20'],
                        'rel_filtered_rms_60_vs_20': d['objectives']['rel_filtered_rms_60_vs_20'],
                        'mfg_90_optician_workload': d['objectives']['mfg_90_optician_workload'],
                    }
                )

                study.add_trial(frozen_trial)

            except Exception as e:
                logger.warning(f"Failed to seed trial {i}: {e}")

        logger.info(f"Seeded {len(study.trials)} trials")

    def run(self, n_trials: int = 50, resume: bool = False):
        """Run NSGA-II optimization."""
        logger.info("\n" + "=" * 70)
        logger.info("M1 MIRROR NSGA-II PURE FEA OPTIMIZATION V13")
        logger.info("=" * 70)
        logger.info(f"Prior FEA trials: {len(self.prior_data)}")
        logger.info(f"New trials to run: {n_trials}")
        logger.info(f"Objectives: {OBJ_NAMES}")

        start_time = time.time()

        # Create or load study
        sampler = NSGAIISampler(
            population_size=self.config.get('nsga2_settings', {}).get('population_size', 20),
            crossover_prob=self.config.get('nsga2_settings', {}).get('crossover_prob', 0.9),
            mutation_prob=self.config.get('nsga2_settings', {}).get('mutation_prob', 0.1),
            seed=42
        )

        study = optuna.create_study(
            study_name="v13_nsga2",
            storage=self.storage,
            directions=['minimize', 'minimize', 'minimize'],  # 3 objectives
            sampler=sampler,
            load_if_exists=resume
        )

        # Seed with prior data if starting fresh
        if not resume or len(study.trials) == 0:
            self.seed_from_prior(study)

        self.trial_count = self._get_next_trial_number()
        logger.info(f"Starting from trial {self.trial_count}")

        # Run optimization
        def objective(trial: optuna.Trial) -> Tuple[float, float, float]:
            # Sample parameters
            params = {}
            for var in self.config['design_variables']:
                if var.get('enabled', False):
                    params[var['name']] = trial.suggest_float(var['name'], var['min'], var['max'])

            # Run FEA
            result = self.fea_runner.run_fea(params, self.trial_count)
            self.trial_count += 1

            if result is None:
                # Return worst-case values for failed trials
                return (1000.0, 1000.0, 1000.0)

            # Store objectives as user attributes
            trial.set_user_attr('source', 'FEA')
            trial.set_user_attr('rel_filtered_rms_40_vs_20', result['objectives']['rel_filtered_rms_40_vs_20'])
            trial.set_user_attr('rel_filtered_rms_60_vs_20', result['objectives']['rel_filtered_rms_60_vs_20'])
            trial.set_user_attr('mfg_90_optician_workload', result['objectives']['mfg_90_optician_workload'])
            trial.set_user_attr('solve_time', result.get('solve_time', 0))

            return (
                result['objectives']['rel_filtered_rms_40_vs_20'],
                result['objectives']['rel_filtered_rms_60_vs_20'],
                result['objectives']['mfg_90_optician_workload']
            )

        # Run
        try:
            study.optimize(
                objective,
                n_trials=n_trials,
                show_progress_bar=True,
                gc_after_trial=True
            )
        except KeyboardInterrupt:
            logger.info("\nOptimization interrupted by user")
        finally:
            self.fea_runner.cleanup()

        # Print results
        elapsed = time.time() - start_time
        self._print_results(study, elapsed)

    def _print_results(self, study: optuna.Study, elapsed: float):
        """Print optimization results."""
        logger.info("\n" + "=" * 70)
        logger.info("OPTIMIZATION COMPLETE")
        logger.info("=" * 70)
        logger.info(f"Time: {elapsed/60:.1f} min ({elapsed/3600:.2f} hours)")
        logger.info(f"Total trials: {len(study.trials)}")

        # Get Pareto front
        pareto_trials = study.best_trials
        logger.info(f"Pareto-optimal trials: {len(pareto_trials)}")

        # Print Pareto front
        logger.info("\nPareto Front:")
        logger.info("-" * 70)
        logger.info(f"{'Trial':>6}  {'40-20 (nm)':>12}  {'60-20 (nm)':>12}  {'Mfg (nm)':>12}")
        logger.info("-" * 70)

        pareto_data = []
        for trial in sorted(pareto_trials, key=lambda t: t.values[0]):
            logger.info(f"{trial.number:>6}  {trial.values[0]:>12.2f}  {trial.values[1]:>12.2f}  {trial.values[2]:>12.2f}")
            pareto_data.append({
                'trial': trial.number,
                'params': trial.params,
                'objectives': {
                    'rel_filtered_rms_40_vs_20': trial.values[0],
                    'rel_filtered_rms_60_vs_20': trial.values[1],
                    'mfg_90_optician_workload': trial.values[2]
                }
            })

        # Save results
        results = {
            'summary': {
                'total_trials': len(study.trials),
                'pareto_size': len(pareto_trials),
                'elapsed_hours': elapsed / 3600
            },
            'pareto_front': pareto_data
        }

        with open(RESULTS_DIR / 'final_results.json', 'w') as f:
            json.dump(results, f, indent=2)

        logger.info(f"\nResults saved to {RESULTS_DIR / 'final_results.json'}")


# ============================================================================
# Main
# ============================================================================

def main():
    parser = argparse.ArgumentParser(description='M1 Mirror NSGA-II V13')
    parser.add_argument('--start', action='store_true', help='Start optimization')
    parser.add_argument('--trials', type=int, default=50, help='Number of new FEA trials')
    parser.add_argument('--resume', action='store_true', help='Resume from existing study')
    args = parser.parse_args()

    if not args.start:
        print("M1 Mirror NSGA-II Pure FEA Optimization V13")
        print("=" * 50)
        print("\nUsage:")
        print("  python run_optimization.py --start")
        print("  python run_optimization.py --start --trials 55")
        print("  python run_optimization.py --start --trials 55 --resume")
        print("\nFor 8-hour overnight run (~55 trials at 8-9 min/trial):")
        print("  python run_optimization.py --start --trials 55")
        print("\nThis will:")
        print(f"  1. Load ~{107} FEA trials from V11 database")
        print(f"  2. Load additional FEA trials from V12 database")
        print("  3. Seed NSGA-II with all prior FEA data")
        print("  4. Run pure FEA multi-objective optimization")
        print("  5. No surrogate - every trial is real FEA")
        return

    with open(CONFIG_PATH, 'r') as f:
        config = json.load(f)

    optimizer = NSGA2Optimizer(config)
    optimizer.run(n_trials=args.trials, resume=args.resume)


if __name__ == "__main__":
    main()