import os
import time
import numpy as np
import pandas as pd
import warnings

# Attempt to import config, but provide fallbacks to make the file standalone
try:
    import config
except ImportError:
    class MockConfig:
        BASE_DIR = "./"
        RANDOM_STATE = 42
        # Default physical bounds for [thickness, twist, solidity, tsr]
        BOUNDS_LOWER = [0.10, 0.0, 0.05, 4.0]
        BOUNDS_UPPER = [0.30, 20.0, 0.15, 10.0]
        POPULATION_SIZE = 100
        N_GENERATIONS = 50
    config = MockConfig()

# Suppress sklearn/pymoo warnings for cleaner CLI output
warnings.filterwarnings("ignore")

try:
    from pymoo.core.problem import Problem
    from pymoo.algorithms.moo.nsga2 import NSGA2
    from pymoo.optimize import minimize
    from pymoo.operators.crossover.sbx import SBX
    from pymoo.operators.mutation.pm import PM
    from pymoo.operators.sampling.rnd import FloatRandomSampling
    from pymoo.termination import get_termination
except ImportError:
    raise ImportError("pymoo is required for optimization. Install via: pip install pymoo")

class WindTurbineProblem(Problem):
    """
    Formulates the surrogate-driven Multi-Objective Optimization Problem (MOOP).
    Inherits from pymoo's base Problem class for vectorized evaluation.
    
    Objectives:
      1. Maximize Aerodynamic Efficiency (Cp) -> minimize -Cp
      2. Minimize Torque Ripple Factor (TRF)
    """
    def __init__(self, cp_builder, trf_builder, **kwargs):
        self.cp_builder = cp_builder
        self.trf_builder = trf_builder
        
        # 4 Variables: [thickness, twist, solidity, tsr]
        n_var = 4 
        # 2 Objectives: [Minimize -Cp, Minimize TRF]
        n_obj = 2 
        
        # Extrapolate physical bounds from config
        xl = np.array(config.BOUNDS_LOWER)
        xu = np.array(config.BOUNDS_UPPER)
        
        super().__init__(n_var=n_var, 
                         n_obj=n_obj, 
                         n_ieq_constr=0, # No explicit inequality constraints in this formulation
                         xl=xl, 
                         xu=xu, 
                         **kwargs)

    def _evaluate(self, x, out, *args, **kwargs):
        """
        Vectorized evaluation function called by NSGA-II.
        Evaluates an entire generation of designs at once for maximum speed.
        
        Args:
            x (np.ndarray): 2D array of shape (pop_size, n_vars) containing the physical inputs.
            out (dict): Output dictionary that pymoo expects us to populate with objective values ("F").
        """
        # Aerodynamic surrogate (Cp) uses all 4 variables: ['thickness', 'twist', 'solidity', 'tsr']
        x_cp = x
        
        # Structural surrogate (TRF) uses only the first 3 variables: ['thickness', 'twist', 'solidity']
        # We slice the array dynamically to prevent feature mismatch
        x_trf = x[:, :3]
        
        # Query the Gaussian Processes
        # predict() returns (mean, sigma). We only need the mean (index 0) for optimization objectives.
        # If we wanted risk-averse optimization, we could penalize based on sigma.
        cp_pred, _ = self.cp_builder.predict(x_cp)
        trf_pred, _ = self.trf_builder.predict(x_trf)
        
        # Objective 1: Pymoo minimizes by default. To MAXIMIZE Cp, we minimize -Cp.
        f1 = -1.0 * cp_pred
        
        # Objective 2: MINIMIZE TRF (already in the correct direction).
        f2 = trf_pred
        
        # Bind the objectives into a single 2D array of shape (pop_size, n_obj)
        out["F"] = np.column_stack([f1, f2])


def run_optimization(cp_builder, trf_builder, dynamic_seed=True, save_csv=True):
    """
    Executes the NSGA-II genetic algorithm to find the optimal Pareto front
    trading off Aerodynamic Efficiency vs. Structural Load.
    
    Args:
        cp_builder: The trained SurrogateBuilder for Cp.
        trf_builder: The trained SurrogateBuilder for TRF.
        dynamic_seed: Whether to generate a fresh random seed for exploration.
    """
    print("\n" + "="*60)
    print(" 🧬 STARTING OPTIMIZATION (NSGA-II)")
    print("="*60)
    
    # 1. Seed Management (The core of the dynamic exploration architecture)
    if dynamic_seed:
        # Generate a fresh seed for every run to explore new genetic combinations
        seed = np.random.randint(1, 1_000_000)
        print(f" -> 🎲 Using Dynamic Seed for Exploration: {seed}")
    else:
        # Fallback to the frozen seed if we need absolute reproducibility
        seed = config.RANDOM_STATE
        print(f" -> 🔒 Using Static Deterministic Seed: {seed}")

    # 2. Problem Definition
    problem = WindTurbineProblem(cp_builder, trf_builder)

    # 3. Algorithm Configuration
    algorithm = NSGA2(
        pop_size=getattr(config, 'POPULATION_SIZE', 100),
        sampling=FloatRandomSampling(),
        # SBX (Simulated Binary Crossover) - Creates children near parents
        crossover=SBX(prob=0.9, eta=15),
        # Polynomial Mutation - Perturbs parameters to escape local minima
        mutation=PM(prob=0.25, eta=20),
        eliminate_duplicates=True
    )
    
    # 4. Termination Criteria
    n_gens = getattr(config, 'N_GENERATIONS', 100)
    termination = get_termination("n_gen", n_gens)
    
    print(f" -> Population Size : {algorithm.pop_size}")
    print(f" -> Generations     : {n_gens}")
    print(f" -> Total Evals     : {algorithm.pop_size * n_gens}")
    print("\n -> Evolving population... Please wait.")

    # 5. Run the Optimizer
    start_time = time.time()
    res = minimize(
        problem,
        algorithm,
        termination,
        seed=seed,
        save_history=False,
        verbose=False
    )
    opt_time = time.time() - start_time
    
    print(f" ✅ Optimization completed in {opt_time:.2f} seconds.")
    
    # 6. Extract the Pareto Front Data
    # res.X contains the optimal physical variables
    # res.F contains the objective values (Remember we minimized -Cp, so we flip it back)
    optimal_X = res.X
    optimal_Cp = -1.0 * res.F[:, 0]
    optimal_TRF = res.F[:, 1]
    
    # Get uncertainties (sigmas) for the pareto points to quantify trust
    _, cp_sigma = cp_builder.predict(optimal_X)
    _, trf_sigma = trf_builder.predict(optimal_X[:, :3])

    # Compile into a clean Pandas DataFrame
    pareto_df = pd.DataFrame({
        'thickness': optimal_X[:, 0],
        'twist': optimal_X[:, 1],
        'solidity': optimal_X[:, 2],
        'tsr': optimal_X[:, 3],
        'Cp_Optimal': optimal_Cp,
        'Cp_Uncertainty': cp_sigma,
        'TRF_Optimal': optimal_TRF,
        'TRF_Uncertainty': trf_sigma
    })
    
    # Sort the Pareto front by Aerodynamic efficiency (highest to lowest)
    pareto_df = pareto_df.sort_values(by='Cp_Optimal', ascending=False).reset_index(drop=True)
    
    if save_csv:
        output_file = os.path.join(getattr(config, 'BASE_DIR', './'), 'pareto_front_optimal.csv')
        pareto_df.to_csv(output_file, index=False)
        print(f" -> 💾 Saved {len(pareto_df)} optimal Pareto designs to: {output_file}")
        
    print("\nTop 3 High-Efficiency Designs:")
    print(pareto_df.head(3)[['thickness', 'twist', 'solidity', 'tsr', 'Cp_Optimal', 'TRF_Optimal']])
    
    return pareto_df

if __name__ == "__main__":
    # If run as a standalone script, attempt to load the pre-trained surrogates and test
    print("Running optimizer in standalone diagnostic mode...")
    
    try:
        from surrogates import SurrogateBuilder
        
        # Initialize mock builders
        cp_builder = SurrogateBuilder('Cp', ['thickness', 'twist', 'solidity', 'tsr'])
        trf_builder = SurrogateBuilder('TRF', ['thickness', 'twist', 'solidity'])
        
        # Load pre-trained models from the previous step
        cp_builder.load_model(os.path.join(config.BASE_DIR, "CP_gp_model.pkl"))
        trf_builder.load_model(os.path.join(config.BASE_DIR, "TRF_gp_model.pkl"))
        
        # Execute optimization
        pareto_results = run_optimization(cp_builder, trf_builder, dynamic_seed=True)
        
    except FileNotFoundError:
        print("❌ Error: Could not find pre-trained surrogate models. Please run surrogates.py first.")
    except ImportError as e:
        print(f"❌ Error importing dependencies: {e}")