#!/usr/bin/env python3
"""
M1 Mirror Cost Reduction - Flat Back V9 (Self-Aware Turbo v3)
==============================================================

SAT v3 improvements over V8:
1. Uses ALL campaign data (601 samples from V5-V8)
2. Adaptive exploration schedule (15% -> 8% -> 3%)
3. Mass threshold at 118 kg (sweet spot, not 115 kg)
4. 70% exploitation near best design
5. Seeded with best known design (WS=218.26)
6. L-BFGS polish phase (last 10 trials)

Target: Beat WS=218.26 (current best from V7)

Usage:
    python run_sat_optimization.py --trials 200
    python run_sat_optimization.py --trials 50 --resume

Author: Atomizer
Created: 2025-12-30
"""

import sys
import os

LICENSE_SERVER = "28000@dalidou;28000@100.80.199.40"
os.environ['SPLM_LICENSE_SERVER'] = LICENSE_SERVER
print(f"[LICENSE] SPLM_LICENSE_SERVER set to: {LICENSE_SERVER}")

STUDY_DIR = os.path.dirname(os.path.abspath(__file__))
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(STUDY_DIR)))
sys.path.insert(0, PROJECT_ROOT)

import json
import re
import time
import logging
import argparse
import sqlite3
from pathlib import Path
from datetime import datetime
from typing import Dict, Any, Optional, List, Tuple, Set

import numpy as np
from scipy.spatial.distance import cdist
from scipy.optimize import minimize

from optimization_engine.nx.solver import NXSolver
from optimization_engine.extractors import ZernikeOPDExtractor
from optimization_engine.surrogates import EnsembleSurrogate, OODDetector, create_and_train_ensemble

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

STUDY_DIR = Path(__file__).parent
SETUP_DIR = STUDY_DIR / "1_setup"
MODEL_DIR = SETUP_DIR / "model"
ITERATIONS_DIR = STUDY_DIR / "2_iterations"
RESULTS_DIR = STUDY_DIR / "3_results"
DB_PATH = RESULTS_DIR / "study.db"
SURROGATE_DIR = RESULTS_DIR / "surrogate"
CONFIG_PATH = SETUP_DIR / "optimization_config.json"

for d in [ITERATIONS_DIR, RESULTS_DIR, SURROGATE_DIR]:
    d.mkdir(parents=True, exist_ok=True)

# ============================================================================
# Logging
# ============================================================================

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

# ============================================================================
# Configuration
# ============================================================================

with open(CONFIG_PATH) as f:
    CONFIG = json.load(f)

STUDY_NAME = CONFIG["study_name"]
DESIGN_VARS = [dv for dv in CONFIG["design_variables"] if dv.get("enabled", True)]
PARAM_NAMES = [dv['name'] for dv in DESIGN_VARS]
N_PARAMS = len(DESIGN_VARS)
PARAM_BOUNDS = [(dv['min'], dv['max']) for dv in DESIGN_VARS]

# Compute parameter ranges for normalization
PARAM_RANGES = np.array([dv['max'] - dv['min'] for dv in DESIGN_VARS])
PARAM_MINS = np.array([dv['min'] for dv in DESIGN_VARS])

OBJECTIVE_WEIGHTS = {'wfe_40_20': 6.0, 'wfe_60_20': 5.0, 'mfg_90': 3.0}
MAX_MASS_KG = 120.0
MASS_PENALTY = 1000.0

# SAT v3 settings
SAT = CONFIG.get('sat_settings', {})
N_ENSEMBLE = SAT.get('n_ensemble_models', 5)
CONFIDENCE_THRESHOLD = SAT.get('confidence_threshold', 0.7)
CANDIDATES_PER_ROUND = SAT.get('candidates_per_round', 1000)
FEA_PER_ROUND = SAT.get('fea_per_round', 1)
RETRAIN_EVERY = SAT.get('retrain_every', 25)
MIN_TRAINING_SAMPLES = SAT.get('min_training_samples', 50)
MIN_DISTANCE_THRESHOLD = SAT.get('min_distance_threshold', 0.03)
JITTER_SCALE = SAT.get('jitter_scale', 0.01)

# Adaptive exploration schedule
EXPLORE_SCHEDULE = SAT.get('exploration_schedule', {})
PHASE1_TRIALS = EXPLORE_SCHEDULE.get('phase1_trials', 30)
PHASE1_WEIGHT = EXPLORE_SCHEDULE.get('phase1_weight', 0.15)
PHASE2_TRIALS = EXPLORE_SCHEDULE.get('phase2_trials', 80)
PHASE2_WEIGHT = EXPLORE_SCHEDULE.get('phase2_weight', 0.08)
PHASE3_WEIGHT = EXPLORE_SCHEDULE.get('phase3_weight', 0.03)

# Mass targeting
MASS_SOFT_THRESHOLD = SAT.get('mass_soft_threshold', 118.0)

# Exploitation settings
EXPLOIT_NEAR_BEST_RATIO = SAT.get('exploit_near_best_ratio', 0.7)
EXPLOIT_SCALE = SAT.get('exploit_scale', 0.05)

# L-BFGS polish
LBFGS_POLISH_TRIALS = SAT.get('lbfgs_polish_trials', 10)
LBFGS_TRUST_RADIUS = SAT.get('lbfgs_trust_radius', 0.1)

# Seed design
SEED_DESIGN = CONFIG.get('seed_design', {}).get('params', None)


def compute_weighted_sum(objectives: Dict[str, float], mass_kg: float) -> float:
    ws = (OBJECTIVE_WEIGHTS['wfe_40_20'] * objectives['wfe_40_20'] +
          OBJECTIVE_WEIGHTS['wfe_60_20'] * objectives['wfe_60_20'] +
          OBJECTIVE_WEIGHTS['mfg_90'] * objectives['mfg_90'])
    if mass_kg > MAX_MASS_KG:
        ws += MASS_PENALTY * (mass_kg - MAX_MASS_KG)
    return ws


def get_exploration_weight(trial_num: int, total_trials: int) -> float:
    """Adaptive exploration weight that decreases over time."""
    if trial_num <= PHASE1_TRIALS:
        return PHASE1_WEIGHT
    elif trial_num <= PHASE2_TRIALS:
        return PHASE2_WEIGHT
    else:
        return PHASE3_WEIGHT


# ============================================================================
# Evaluated Points Tracker
# ============================================================================

class EvaluatedPointsTracker:
    """Tracks all evaluated parameter sets to prevent duplicates."""

    def __init__(self, param_ranges: np.ndarray, min_distance: float = 0.03):
        self.param_ranges = param_ranges
        self.min_distance = min_distance
        self.evaluated_points: List[np.ndarray] = []
        self._evaluated_set: Set[tuple] = set()

    def add_point(self, x: np.ndarray):
        x = np.asarray(x).flatten()
        x_rounded = tuple(np.round(x, 2))
        if x_rounded not in self._evaluated_set:
            self._evaluated_set.add(x_rounded)
            self.evaluated_points.append(x)

    def is_duplicate(self, x: np.ndarray) -> bool:
        x = np.asarray(x).flatten()
        x_rounded = tuple(np.round(x, 2))
        if x_rounded in self._evaluated_set:
            return True
        if len(self.evaluated_points) == 0:
            return False
        x_norm = x / self.param_ranges
        evaluated_norm = np.array(self.evaluated_points) / self.param_ranges
        distances = cdist([x_norm], evaluated_norm, metric='euclidean')[0]
        return distances.min() < self.min_distance

    def min_distance_to_evaluated(self, x: np.ndarray) -> float:
        if len(self.evaluated_points) == 0:
            return float('inf')
        x_norm = x / self.param_ranges
        evaluated_norm = np.array(self.evaluated_points) / self.param_ranges
        distances = cdist([x_norm], evaluated_norm, metric='euclidean')[0]
        return distances.min()

    def filter_candidates(self, candidates: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        if len(self.evaluated_points) == 0:
            return candidates, np.full(len(candidates), float('inf'))
        candidates_norm = candidates / self.param_ranges
        evaluated_norm = np.array(self.evaluated_points) / self.param_ranges
        distances = cdist(candidates_norm, evaluated_norm, metric='euclidean')
        min_distances = distances.min(axis=1)
        mask = min_distances >= self.min_distance
        return candidates[mask], min_distances[mask]

    def __len__(self):
        return len(self.evaluated_points)


# ============================================================================
# Database
# ============================================================================

def init_database():
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()

    cursor.execute('CREATE TABLE IF NOT EXISTS studies (study_id INTEGER PRIMARY KEY, study_name TEXT UNIQUE)')
    cursor.execute('INSERT OR IGNORE INTO studies (study_id, study_name) VALUES (1, ?)', (STUDY_NAME,))

    cursor.execute('''CREATE TABLE IF NOT EXISTS trials (
        trial_id INTEGER PRIMARY KEY, study_id INTEGER DEFAULT 1, number INTEGER,
        state TEXT DEFAULT 'COMPLETE', datetime_start TEXT, datetime_complete TEXT)''')

    cursor.execute('''CREATE TABLE IF NOT EXISTS trial_values (
        trial_value_id INTEGER PRIMARY KEY, trial_id INTEGER, objective INTEGER DEFAULT 0, value REAL)''')

    cursor.execute('''CREATE TABLE IF NOT EXISTS trial_params (
        trial_param_id INTEGER PRIMARY KEY, trial_id INTEGER, param_name TEXT,
        param_value REAL, distribution_json TEXT DEFAULT '{}')''')

    cursor.execute('''CREATE TABLE IF NOT EXISTS trial_user_attributes (
        trial_user_attribute_id INTEGER PRIMARY KEY, trial_id INTEGER, key TEXT, value_json TEXT)''')

    cursor.execute('''CREATE TABLE IF NOT EXISTS study_directions (
        study_direction_id INTEGER PRIMARY KEY, study_id INTEGER, objective INTEGER DEFAULT 0,
        direction TEXT DEFAULT 'MINIMIZE')''')
    cursor.execute('INSERT OR IGNORE INTO study_directions VALUES (1, 1, 0, "MINIMIZE")')

    cursor.execute('''CREATE TABLE IF NOT EXISTS version_info (
        version_info_id INTEGER PRIMARY KEY, schema_version TEXT, library_version TEXT)''')
    cursor.execute('INSERT OR IGNORE INTO version_info VALUES (1, "0.9.0", "3.0.0")')

    cursor.execute('CREATE TABLE IF NOT EXISTS study_info (key TEXT PRIMARY KEY, value TEXT)')
    cursor.execute('INSERT OR REPLACE INTO study_info VALUES (?, ?)', ('study_name', STUDY_NAME))
    cursor.execute('INSERT OR REPLACE INTO study_info VALUES (?, ?)', ('algorithm', 'SAT_v3'))
    cursor.execute('INSERT OR REPLACE INTO study_info VALUES (?, ?)', ('start_time', datetime.now().isoformat()))

    conn.commit()
    conn.close()
    logger.info(f"[DB] Initialized: {DB_PATH}")


def log_trial_to_db(trial_num: int, params: Dict, objectives: Dict, ws: float,
                    mass_kg: float, is_feasible: bool, source: str = 'fea'):
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()

    cursor.execute('INSERT OR REPLACE INTO trials VALUES (?, 1, ?, "COMPLETE", ?, ?)',
                   (trial_num, trial_num, datetime.now().isoformat(), datetime.now().isoformat()))

    cursor.execute('DELETE FROM trial_values WHERE trial_id = ?', (trial_num,))
    cursor.execute('INSERT INTO trial_values (trial_id, objective, value) VALUES (?, 0, ?)', (trial_num, ws))

    cursor.execute('DELETE FROM trial_params WHERE trial_id = ?', (trial_num,))
    for name, value in params.items():
        cursor.execute('INSERT INTO trial_params (trial_id, param_name, param_value) VALUES (?, ?, ?)',
                       (trial_num, name, value))

    cursor.execute('DELETE FROM trial_user_attributes WHERE trial_id = ?', (trial_num,))
    for key, value in objectives.items():
        cursor.execute('INSERT INTO trial_user_attributes (trial_id, key, value_json) VALUES (?, ?, ?)',
                       (trial_num, f"obj_{key}", json.dumps(value)))
    cursor.execute('INSERT INTO trial_user_attributes VALUES (NULL, ?, "mass_kg", ?)', (trial_num, json.dumps(mass_kg)))
    cursor.execute('INSERT INTO trial_user_attributes VALUES (NULL, ?, "is_feasible", ?)', (trial_num, json.dumps(is_feasible)))
    cursor.execute('INSERT INTO trial_user_attributes VALUES (NULL, ?, "source", ?)', (trial_num, json.dumps(source)))

    conn.commit()
    conn.close()


def get_trial_count() -> int:
    if not DB_PATH.exists():
        return 0
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    cursor.execute('SELECT COUNT(*) FROM trials')
    count = cursor.fetchone()[0]
    conn.close()
    return count


# ============================================================================
# Training Data Loading
# ============================================================================

def load_training_data_from_db(db_path: Path, include_infeasible: bool = True) -> Tuple[np.ndarray, np.ndarray]:
    """Load training data from a study database."""
    if not db_path.exists():
        return np.array([]), np.array([])

    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    cursor.execute('SELECT trial_id FROM trials WHERE state = "COMPLETE"')
    trial_ids = [row[0] for row in cursor.fetchall()]

    X_list, Y_list = [], []

    for trial_id in trial_ids:
        cursor.execute('SELECT param_name, param_value FROM trial_params WHERE trial_id = ?', (trial_id,))
        params_raw = {row[0]: row[1] for row in cursor.fetchall()}
        params = {(k.split(']', 1)[1] if ']' in k else k): v for k, v in params_raw.items()}

        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()}

        mass_kg = attrs.get('mass_kg', 999.0)

        if mass_kg > 200.0:
            continue

        if not include_infeasible and mass_kg > MAX_MASS_KG:
            continue

        x = []
        skip = False
        for dv in DESIGN_VARS:
            if dv['name'] not in params:
                skip = True
                break
            x.append(params[dv['name']])
        if skip:
            continue

        wfe_40 = attrs.get('obj_wfe_40_20') or attrs.get('wfe_40_20') or attrs.get('rel_filtered_rms_40_vs_20')
        wfe_60 = attrs.get('obj_wfe_60_20') or attrs.get('wfe_60_20') or attrs.get('rel_filtered_rms_60_vs_20')
        mfg_90 = attrs.get('obj_mfg_90') or attrs.get('mfg_90') or attrs.get('rel_filtered_j1to3_90_vs_20')

        if wfe_40 is None or wfe_60 is None or mfg_90 is None:
            continue

        ws = 6.0 * wfe_40 + 5.0 * wfe_60 + 3.0 * mfg_90
        if mass_kg > MAX_MASS_KG:
            ws += MASS_PENALTY * (mass_kg - MAX_MASS_KG)

        X_list.append(x)
        Y_list.append([wfe_40, wfe_60, mfg_90, mass_kg, ws])

    conn.close()
    return (np.array(X_list), np.array(Y_list)) if X_list else (np.array([]), np.array([]))


def load_all_training_data() -> Tuple[np.ndarray, np.ndarray]:
    X_all, Y_all = [], []

    for source in CONFIG.get('training_data_sources', []):
        db_path = STUDY_DIR / source['path']
        X, Y = load_training_data_from_db(db_path)
        if len(X) > 0:
            X_all.append(X)
            Y_all.append(Y)
            logger.info(f"[DATA] Loaded {len(X)} from {source['study']}")

    X, Y = load_training_data_from_db(DB_PATH)
    if len(X) > 0:
        X_all.append(X)
        Y_all.append(Y)
        logger.info(f"[DATA] Loaded {len(X)} from current study")

    if not X_all:
        return np.array([]), np.array([])

    X_combined = np.vstack(X_all)
    Y_combined = np.vstack(Y_all)
    logger.info(f"[DATA] Total: {len(X_combined)} samples")
    return X_combined, Y_combined


# ============================================================================
# FEA Execution
# ============================================================================

def setup_nx_solver() -> NXSolver:
    nx_settings = CONFIG.get('nx_settings', {})
    nx_install_dir = nx_settings.get('nx_install_path', 'C:\\Program Files\\Siemens\\DesigncenterNX2512')
    version_match = re.search(r'NX(\d+)|DesigncenterNX(\d+)', nx_install_dir)
    nastran_version = (version_match.group(1) or version_match.group(2)) if version_match else "2512"

    solver = NXSolver(
        master_model_dir=str(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=STUDY_NAME
    )
    logger.info(f"[NX] Solver ready (Nastran {nastran_version})")
    return solver


def extract_objectives(op2_path: Path, working_dir: Path) -> Optional[Dict[str, float]]:
    try:
        zernike_settings = CONFIG.get("zernike_settings", {})
        extractor = ZernikeOPDExtractor(
            op2_path, figure_path=None, 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")
        wfe_40 = extractor.extract_relative("3", ref)
        wfe_60 = extractor.extract_relative("4", ref)
        mfg_90 = extractor.extract_relative("1", ref)

        mass = None
        for f in working_dir.glob("*_props.json"):
            with open(f) as fp:
                mass = json.load(fp).get("mass_kg")
                break
        if mass is None:
            mass_file = working_dir / "_temp_mass.txt"
            if mass_file.exists():
                mass = float(mass_file.read_text().strip())
        if mass is None:
            mass = 999.0

        return {
            "wfe_40_20": wfe_40["relative_filtered_rms_nm"],
            "wfe_60_20": wfe_60["relative_filtered_rms_nm"],
            "mfg_90": mfg_90["relative_rms_filter_j1to3"],
            "mass_kg": mass
        }
    except Exception as e:
        logger.error(f"Extraction failed: {e}")
        return None


def run_fea_trial(nx_solver: NXSolver, trial_num: int, params: Dict) -> Optional[Dict]:
    logger.info(f"[TRIAL {trial_num:04d}] Starting FEA...")

    expressions = {dv['expression_name']: params[dv['name']] for dv in DESIGN_VARS}
    for fixed in CONFIG.get('fixed_parameters', []):
        expressions[fixed['name']] = fixed['value']

    iter_folder = nx_solver.create_iteration_folder(ITERATIONS_DIR, trial_num, expressions)

    try:
        nx_settings = CONFIG.get('nx_settings', {})
        sim_file = iter_folder / nx_settings.get('sim_file', 'ASSY_M1_assyfem1_sim1.sim')
        t_start = time.time()

        result = 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"[TRIAL {trial_num:04d}] FEA failed")
            return None

        logger.info(f"[TRIAL {trial_num:04d}] Solved in {solve_time:.1f}s")
        objectives = extract_objectives(Path(result['op2_file']), iter_folder)

        if objectives is None:
            return None

        mass_kg = objectives.pop('mass_kg')
        is_feasible = mass_kg <= MAX_MASS_KG
        ws = compute_weighted_sum(objectives, mass_kg)

        logger.info(f"[TRIAL {trial_num:04d}] WFE 40-20={objectives['wfe_40_20']:.2f}nm, "
                    f"60-20={objectives['wfe_60_20']:.2f}nm, Mfg90={objectives['mfg_90']:.2f}nm, "
                    f"Mass={mass_kg:.2f}kg, WS={ws:.2f}")

        return {'trial_num': trial_num, 'params': params, 'objectives': objectives,
                'mass_kg': mass_kg, 'weighted_sum': ws, 'is_feasible': is_feasible, 'solve_time': solve_time}
    except Exception as e:
        logger.error(f"[TRIAL {trial_num:04d}] Error: {e}")
        return None


# ============================================================================
# SAT v3 Algorithm
# ============================================================================

def params_to_array(params: Dict[str, float]) -> np.ndarray:
    return np.array([params[dv['name']] for dv in DESIGN_VARS])


def array_to_params(x: np.ndarray) -> Dict[str, float]:
    return {dv['name']: round(float(x[i]), 2) for i, dv in enumerate(DESIGN_VARS)}


def sample_random_candidate() -> np.ndarray:
    return np.array([np.random.uniform(dv['min'], dv['max']) for dv in DESIGN_VARS])


def sample_near_point(center: np.ndarray, scale: float = 0.05) -> np.ndarray:
    """Sample near a point with given scale (fraction of range)."""
    x = center + np.random.normal(0, scale, size=len(center)) * PARAM_RANGES
    for i, dv in enumerate(DESIGN_VARS):
        x[i] = np.clip(x[i], dv['min'], dv['max'])
    return x


def add_jitter(x: np.ndarray, scale: float = JITTER_SCALE) -> np.ndarray:
    jitter = np.random.uniform(-scale, scale, size=len(x)) * PARAM_RANGES
    x_jittered = x + jitter
    for i, dv in enumerate(DESIGN_VARS):
        x_jittered[i] = np.clip(x_jittered[i], dv['min'], dv['max'])
    return x_jittered


def run_sat_optimization(n_trials: int = 200, resume: bool = False):
    print("\n" + "="*70)
    print("M1 MIRROR FLAT BACK V9 - SELF-AWARE TURBO v3")
    print("="*70)
    print(f"Algorithm: SAT v3 with adaptive exploration + L-BFGS polish")
    print(f"Trials: {n_trials}")
    print(f"Training data: V5 + V6 + V7 + V8 (601 samples)")
    print(f"Mass threshold: {MASS_SOFT_THRESHOLD} kg (sweet spot)")
    print(f"Exploration schedule: {PHASE1_WEIGHT} -> {PHASE2_WEIGHT} -> {PHASE3_WEIGHT}")
    print(f"Target: Beat WS=218.26")
    print("="*70 + "\n")

    init_database()
    nx_solver = setup_nx_solver()

    # Load ALL training data
    X_train, Y_train = load_all_training_data()
    logger.info(f"[SAT] Loaded {len(X_train)} training samples")

    # Initialize tracker with ALL evaluated points
    tracker = EvaluatedPointsTracker(PARAM_RANGES, min_distance=MIN_DISTANCE_THRESHOLD)
    for x in X_train:
        tracker.add_point(x)
    logger.info(f"[SAT] Tracking {len(tracker)} evaluated points")

    # Initialize surrogate
    surrogate = None
    if len(X_train) >= MIN_TRAINING_SAMPLES:
        logger.info(f"[SAT] Training ensemble surrogate on {len(X_train)} samples...")
        surrogate = create_and_train_ensemble(
            X_train, Y_train,
            n_models=N_ENSEMBLE,
            epochs=SAT.get('training_epochs', 800)
        )
        surrogate.save(SURROGATE_DIR)

    # Track best - initialize with seed design if provided
    best_ws = float('inf')
    best_trial = None
    best_x = None

    # Use seed design as initial best
    if SEED_DESIGN:
        best_x = params_to_array(SEED_DESIGN)
        logger.info(f"[SAT] Seeded with best known design")

    # Find actual best from training data
    if len(Y_train) > 0:
        feasible_mask = Y_train[:, 3] <= MAX_MASS_KG
        if feasible_mask.any():
            feasible_Y = Y_train[feasible_mask]
            feasible_X = X_train[feasible_mask]
            best_idx = feasible_Y[:, 4].argmin()
            best_ws = feasible_Y[best_idx, 4]
            best_x = feasible_X[best_idx]
            logger.info(f"[SAT] Best feasible from training: WS={best_ws:.2f}")

    trial_counter = get_trial_count()
    trials_run = 0
    unique_designs = 0
    total_trials = n_trials

    while trials_run < n_trials:
        remaining = n_trials - trials_run
        current_trial = trials_run + 1

        # Determine phase
        if remaining <= LBFGS_POLISH_TRIALS and surrogate is not None and best_x is not None:
            phase = "LBFGS_POLISH"
        elif current_trial <= PHASE1_TRIALS:
            phase = "PHASE1_EXPLORE"
        elif current_trial <= PHASE2_TRIALS:
            phase = "PHASE2_BALANCED"
        else:
            phase = "PHASE3_EXPLOIT"

        exploration_weight = get_exploration_weight(current_trial, total_trials)

        # Generate candidates based on phase
        candidates = []

        if phase == "LBFGS_POLISH":
            # L-BFGS polish: small perturbations around best
            logger.info(f"[{phase}] Polishing near best (remaining: {remaining})")
            for _ in range(CANDIDATES_PER_ROUND):
                candidates.append(sample_near_point(best_x, scale=LBFGS_TRUST_RADIUS))
        else:
            # SAT-guided with adaptive exploration
            exploit_ratio = EXPLOIT_NEAR_BEST_RATIO if best_x is not None else 0.0

            for _ in range(CANDIDATES_PER_ROUND):
                r = np.random.random()
                if r < exploit_ratio and best_x is not None:
                    # Exploit: sample near best with small scale
                    scale = np.random.uniform(EXPLOIT_SCALE, EXPLOIT_SCALE * 3)
                    candidates.append(sample_near_point(best_x, scale=scale))
                else:
                    # Explore: random sampling
                    candidates.append(sample_random_candidate())

        candidates = np.array(candidates)

        # Filter candidates
        valid_candidates, distances = tracker.filter_candidates(candidates)

        if len(valid_candidates) == 0:
            logger.warning(f"[{phase}] All candidates rejected! Forcing random exploration.")
            for attempt in range(100):
                x = sample_random_candidate()
                x = add_jitter(x, scale=0.05)
                if not tracker.is_duplicate(x):
                    valid_candidates = np.array([x])
                    distances = np.array([tracker.min_distance_to_evaluated(x)])
                    break
            else:
                logger.error("[SAT] Could not find unique candidate after 100 attempts!")
                break

        # Select best candidate using surrogate
        if surrogate is not None and len(valid_candidates) > 0:
            predictions = surrogate.predict_with_confidence(valid_candidates)

            pred_mass = predictions['mean'][:, 3]
            pred_ws = predictions['mean'][:, 4]

            # Mass penalty with sweet spot threshold (118 kg, not 115)
            mass_penalty = np.maximum(0, pred_mass - MASS_SOFT_THRESHOLD)
            feasibility_score = mass_penalty * 5.0

            # Acquisition with adaptive exploration weight
            norm_ws = (pred_ws - pred_ws.min()) / (pred_ws.max() - pred_ws.min() + 1e-8)
            norm_dist = distances / (distances.max() + 1e-8)
            norm_mass_penalty = feasibility_score / (feasibility_score.max() + 1e-8)

            # Adaptive acquisition
            acquisition = norm_ws - exploration_weight * norm_dist + norm_mass_penalty

            best_idx = acquisition.argmin()
            x_selected = valid_candidates[best_idx]
            source = f'{phase.lower()}_{predictions["recommendation"][best_idx]}'

            logger.info(f"[{phase}] Trial {current_trial}: pred_WS={pred_ws[best_idx]:.2f}, "
                        f"pred_mass={pred_mass[best_idx]:.1f}kg, "
                        f"dist={distances[best_idx]:.3f}, explore_w={exploration_weight:.2f}")
        else:
            best_idx = np.random.randint(len(valid_candidates))
            x_selected = valid_candidates[best_idx]
            source = 'random'

        # Add jitter
        x_selected = add_jitter(x_selected, scale=JITTER_SCALE)

        if tracker.is_duplicate(x_selected):
            x_selected = add_jitter(x_selected, scale=0.03)

        params = array_to_params(x_selected)

        # Run FEA
        trial_counter += 1
        result = run_fea_trial(nx_solver, trial_counter, params)
        trials_run += 1

        if result is None:
            logger.warning(f"[{phase}] Trial {trial_counter} failed")
            continue

        tracker.add_point(params_to_array(params))
        unique_designs += 1

        log_trial_to_db(trial_counter, params, result['objectives'],
                        result['weighted_sum'], result['mass_kg'], result['is_feasible'], source)

        # Track best
        if result['is_feasible'] and result['weighted_sum'] < best_ws:
            improvement = best_ws - result['weighted_sum']
            best_ws = result['weighted_sum']
            best_trial = trial_counter
            best_x = params_to_array(params)
            logger.info(f"\n{'*'*60}")
            logger.info(f"*** NEW BEST! Trial {trial_counter}: WS={best_ws:.2f} (improved by {improvement:.2f}) ***")
            logger.info(f"{'*'*60}\n")

        # Retrain surrogate periodically
        if trials_run % RETRAIN_EVERY == 0 and trials_run > 0:
            logger.info(f"[SAT] Retraining surrogate with {len(tracker)} samples...")
            X_train, Y_train = load_all_training_data()
            surrogate = create_and_train_ensemble(
                X_train, Y_train,
                n_models=N_ENSEMBLE,
                epochs=SAT.get('training_epochs', 800)
            )
            surrogate.save(SURROGATE_DIR)

    # Final summary
    print("\n" + "="*70)
    print("OPTIMIZATION COMPLETE")
    print("="*70)
    print(f"Trials completed: {trials_run}")
    print(f"Unique designs: {unique_designs} ({100*unique_designs/max(1,trials_run):.1f}%)")
    print(f"Total tracked points: {len(tracker)}")
    print(f"Best WS: {best_ws:.2f} (trial {best_trial})")
    print(f"Target: 218.26")
    if best_ws < 218.26:
        print(f"*** SUCCESS! Beat target by {218.26 - best_ws:.2f} ***")
    else:
        print(f"Gap to target: {best_ws - 218.26:.2f}")
    print("="*70)


def main():
    parser = argparse.ArgumentParser(description='M1 Mirror Flat Back V9 - SAT v3 Optimization')
    parser.add_argument('--trials', type=int, default=200, help='Number of trials')
    parser.add_argument('--resume', action='store_true', help='Resume from existing study')
    args = parser.parse_args()

    run_sat_optimization(n_trials=args.trials, resume=args.resume)


if __name__ == "__main__":
    main()
