#!/usr/bin/env python
"""
L-BFGS Polish for M1 Mirror V14

This script uses gradient-based optimization to polish the best designs
found by the surrogate-assisted optimization.

Key advantage: L-BFGS exploits the differentiability of the trained MLP
surrogate to find precise local optima 100-1000x faster than TPE/CMA-ES.

Usage:
    python run_lbfgs_polish.py                    # Use top FEA as starts
    python run_lbfgs_polish.py --n-starts 50      # More starting points
    python run_lbfgs_polish.py --grid-then-grad   # Hybrid grid + gradient

After running:
    - Results saved to 3_results/lbfgs_results.json
    - Top candidates ready for FEA validation
"""

import sys
from pathlib import Path

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

import json
import argparse
from datetime import datetime

from optimization_engine.gradient_optimizer import (
    GradientOptimizer,
    run_lbfgs_polish,
    OptimizationResult
)
from optimization_engine.generic_surrogate import GenericSurrogate


def main():
    parser = argparse.ArgumentParser(
        description="L-BFGS Polish - Gradient-based refinement using trained surrogate"
    )
    parser.add_argument(
        "--n-starts", type=int, default=20,
        help="Number of starting points (default: 20)"
    )
    parser.add_argument(
        "--n-iterations", type=int, default=100,
        help="L-BFGS iterations per starting point (default: 100)"
    )
    parser.add_argument(
        "--grid-then-grad", action="store_true",
        help="Use hybrid grid search + gradient refinement"
    )
    parser.add_argument(
        "--n-grid", type=int, default=500,
        help="Grid samples for hybrid mode (default: 500)"
    )
    parser.add_argument(
        "--random-only", action="store_true",
        help="Use only random starting points (ignore FEA database)"
    )
    parser.add_argument(
        "--validate-with-fea", action="store_true",
        help="Validate top results with actual FEA (requires NX)"
    )
    parser.add_argument(
        "--top-n", type=int, default=5,
        help="Number of top results to validate with FEA (default: 5)"
    )

    args = parser.parse_args()

    # Paths
    config_path = STUDY_DIR / "1_setup" / "optimization_config.json"
    results_dir = STUDY_DIR / "3_results"
    surrogate_path = results_dir / "surrogate_best.pt"

    # Check prerequisites
    if not surrogate_path.exists():
        # Try other locations
        alt_paths = [
            results_dir / "surrogate_iter3.pt",
            results_dir / "surrogate_iter2.pt",
            results_dir / "surrogate_iter1.pt",
            results_dir / "surrogate_initial.pt",
        ]
        for alt in alt_paths:
            if alt.exists():
                surrogate_path = alt
                break
        else:
            print(f"ERROR: No trained surrogate found in {results_dir}")
            print("Run training first: python run_nn_optimization.py --train")
            return 1

    print(f"\n{'#'*70}")
    print(f"# L-BFGS GRADIENT POLISH - M1 Mirror V14")
    print(f"{'#'*70}")
    print(f"Surrogate: {surrogate_path.name}")
    print(f"Config: {config_path}")

    # Load config
    with open(config_path) as f:
        config = json.load(f)

    # Normalize config for surrogate
    normalized_config = normalize_config(config)

    # Load surrogate
    print(f"\nLoading surrogate model...")
    surrogate = GenericSurrogate(normalized_config)
    surrogate.load(surrogate_path)

    # Extract weights
    weights = [obj.get('weight', 1.0) for obj in config.get('objectives', [])]
    directions = [obj.get('direction', 'minimize') for obj in config.get('objectives', [])]

    print(f"Design variables: {len(surrogate.design_var_names)}")
    print(f"Objectives: {surrogate.objective_names}")
    print(f"Weights: {weights}")

    # Create optimizer
    optimizer = GradientOptimizer(
        surrogate=surrogate,
        objective_weights=weights,
        objective_directions=directions
    )

    # Run optimization
    if args.grid_then_grad:
        # Hybrid mode: grid search then gradient refinement
        print(f"\nUsing HYBRID mode: {args.n_grid} grid samples -> top {args.n_starts} gradient polish")
        results = optimizer.grid_search_then_gradient(
            n_grid_samples=args.n_grid,
            n_top_for_gradient=args.n_starts,
            n_iterations=args.n_iterations,
            verbose=True
        )
    else:
        # Standard multi-start L-BFGS
        results = run_lbfgs_polish(
            study_dir=STUDY_DIR,
            n_starts=args.n_starts,
            use_top_fea=not args.random_only,
            n_iterations=args.n_iterations,
            verbose=True
        )

    # Ensure results is a list
    if isinstance(results, OptimizationResult):
        results = [results]

    # Save detailed results
    output = {
        'timestamp': datetime.now().isoformat(),
        'study': 'm1_mirror_adaptive_V14',
        'method': 'L-BFGS',
        'n_starts': args.n_starts,
        'n_iterations': args.n_iterations,
        'mode': 'grid_then_grad' if args.grid_then_grad else 'multi_start',
        'surrogate_path': str(surrogate_path),
        'results': [r.to_dict() for r in results]
    }

    output_path = results_dir / "lbfgs_results.json"
    with open(output_path, 'w') as f:
        json.dump(output, f, indent=2)

    print(f"\n{'='*70}")
    print(f"RESULTS SAVED: {output_path}")
    print(f"{'='*70}")

    # Summary table
    print(f"\nTOP {min(10, len(results))} RESULTS:")
    print("-" * 70)
    print(f"{'Rank':<5} {'WS':>10} {'Improve':>10} {'Converged':>10}")
    print("-" * 70)

    for i, r in enumerate(results[:10]):
        converged = "Yes" if r.converged else "No"
        improve = f"{r.improvement:.2f}" if r.improvement > 0 else "-"
        print(f"{i+1:<5} {r.weighted_sum:>10.4f} {improve:>10} {converged:>10}")

    # Show best parameters
    best = results[0]
    print(f"\nBEST DESIGN (WS={best.weighted_sum:.4f}):")
    print("-" * 70)
    for name, val in best.params.items():
        bounds = surrogate.design_var_bounds[name]
        pct = (val - bounds[0]) / (bounds[1] - bounds[0]) * 100
        at_bound = " [AT BOUND]" if pct < 1 or pct > 99 else ""
        print(f"  {name}: {val:.4f} ({pct:.0f}% of range){at_bound}")

    print(f"\nOBJECTIVES:")
    for name, val in best.objectives.items():
        print(f"  {name}: {val:.4f}")

    # Optional FEA validation
    if args.validate_with_fea:
        print(f"\n{'='*70}")
        print(f"FEA VALIDATION (Top {args.top_n})")
        print(f"{'='*70}")

        try:
            from optimization_engine.nx_solver import NXSolver

            model_dir = STUDY_DIR / "1_setup" / "model"
            sim_file = model_dir / config['nx_settings']['sim_file']

            solver = NXSolver(nastran_version="2506")

            for i, result in enumerate(results[:args.top_n]):
                print(f"\n[{i+1}/{args.top_n}] Validating design WS={result.weighted_sum:.4f}...")

                fea_result = solver.run_simulation(
                    sim_file=sim_file,
                    working_dir=model_dir,
                    expression_updates=result.params,
                    solution_name=config['nx_settings'].get('solution_name'),
                    cleanup=True
                )

                if fea_result['success']:
                    print(f"  FEA SUCCESS - OP2: {fea_result['op2_file']}")
                    # TODO: Extract and compare objectives
                else:
                    print(f"  FEA FAILED")

        except ImportError:
            print("NX Solver not available for validation")
        except Exception as e:
            print(f"FEA validation error: {e}")

    print(f"\n{'#'*70}")
    print(f"# L-BFGS POLISH COMPLETE")
    print(f"{'#'*70}")
    print(f"\nNext steps:")
    print(f"  1. Review top designs in {output_path}")
    print(f"  2. Validate promising candidates with FEA:")
    print(f"     python run_lbfgs_polish.py --validate-with-fea --top-n 3")
    print(f"  3. Or add to turbo loop for continued optimization")

    return 0


def normalize_config(config):
    """Normalize config format for GenericSurrogate."""
    normalized = {
        'study_name': config.get('study_name', 'unnamed_study'),
        'description': config.get('description', ''),
        'design_variables': [],
        'objectives': [],
        'constraints': [],
        'simulation': {},
        'neural_acceleration': config.get('neural_acceleration', {}),
    }

    # Normalize design variables
    for var in config.get('design_variables', []):
        normalized['design_variables'].append({
            'name': var.get('parameter') or var.get('name'),
            'type': var.get('type', 'continuous'),
            'min': var.get('bounds', [var.get('min', 0), var.get('max', 1)])[0] if 'bounds' in var else var.get('min', 0),
            'max': var.get('bounds', [var.get('min', 0), var.get('max', 1)])[1] if 'bounds' in var else var.get('max', 1),
        })

    # Normalize objectives
    for obj in config.get('objectives', []):
        normalized['objectives'].append({
            'name': obj.get('name'),
            'direction': obj.get('goal') or obj.get('direction', 'minimize'),
            'weight': obj.get('weight', 1.0),
        })

    # Normalize simulation
    sim = config.get('simulation', config.get('nx_settings', {}))
    normalized['simulation'] = {
        'sim_file': sim.get('sim_file', ''),
        'dat_file': sim.get('dat_file', ''),
        'solution_name': sim.get('solution_name', 'Solution 1'),
    }

    return normalized


if __name__ == "__main__":
    sys.exit(main())
