import time
import numpy as np
import matplotlib.pyplot as plt
import copy
from mopso import MOPSO 
from surrogate_handler import SurrogateHandler
import pandas as pd


import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D # Nécessaire pour la 3D

def plot_pareto_3d(archive, model_type:str):
    fig = plt.figure(figsize=(12, 8))
    ax = fig.add_subplot(111, projection='3d')

    # Extraction des scores depuis l'archive
    # f_best[0] = Coût, f_best[1] = Insatisfaction, f_best[2] = Stress Réseau
    f1 = [p.f_best[0] for p in archive] 
    f2 = [p.f_best[1] for p in archive] 
    f3 = [p.f_best[2] for p in archive] 

    # Création du nuage de points
    img = ax.scatter(f1, f2, f3, c=f3, cmap='viridis', s=60, edgecolors='black')
    
    ax.set_xlabel('Coût (€)')
    ax.set_ylabel('Insatisfaction (SoC manquant)')
    ax.set_zlabel('Pic Réseau (kW)')
    ax.set_title(f'Front de Pareto des Solutions Non-Dominées ({model_type})')
    
    # Barre de couleur
    cbar = fig.colorbar(img, ax=ax, pad=0.1)
    cbar.set_label('Intensité du Pic Réseau (kW)')
    
    # Sauvegarde et affichage
    filename = f"{model_type}_pareto_3d.png"
    plt.savefig(filename)
    print(f"Graphique sauvegardé sous : {filename}")
    plt.show()



class SmartMOPSO(MOPSO):
    def __init__(self, model_type=None, **kwargs):
        super().__init__(**kwargs)
        
        # Initialize Surrogate Handler if model_type is provided
        self.use_surrogate = (model_type is not None)
        if self.use_surrogate:
            self.surrogate_handler = SurrogateHandler(model_type)

            # Pre-fill with initial particle data
            for p in self.particles:
                self.surrogate_handler.add_data(p.x, p.f_current[1])

    def iterate(self):
        train_freq = 10 # Retrain every 10 iterations
        
        # Check if retraining is needed
        if self.use_surrogate and (self.t % train_freq == 0):
            self.surrogate_handler.train()

        # Determine if AI prediction should be used
        use_ai = (self.use_surrogate and 
                  self.surrogate_handler.is_trained and 
                  self.t % train_freq != 0)

        # Main loop (overriding original logic to manage control flow)
        for t in range(self.t): 
            self.select_leader()
            
            for i in range(self.n):
                # Movement 
                self.particles[i].update_velocity(self.leader.x, self.c1, self.c2, self.w)
                self.particles[i].update_position()
                self.particles[i].keep_boudaries(self.A_max)

                if use_ai:
                    # Fast exact calculation (f1, f3)
                    f1 = self.particles[i].f1(self.prices)
                    f3 = self.particles[i].f3()
                    
                    # Slow prediction (f2) by using Surrogate
                    f2_pred = self.surrogate_handler.predict(self.particles[i].x)
                    
                    # Inject scores without running the expensive 'updating_socs'
                    self.particles[i].f_current = [f1, f2_pred, f3]
                    
                else:
                    # Standard Calculation (Slow and Exact)
                    self.particles[i].updating_socs(self.socs, self.capacities)
                    self.particles[i].evaluate(self.prices, self.socs, self.socs_req, self.times)
                    
                    # Capture data for AI training
                    if self.use_surrogate:
                        self.surrogate_handler.add_data(self.particles[i].x, self.particles[i].f_current[1])

                self.particles[i].update_best()
            
            self.update_archive()


def calculate_elec_prices(csv_file:str, sep:str=';'):
    elec_df = pd.read_csv(filepath_or_buffer=csv_file, sep=sep, skipinitialspace=True)

    # Mean of Winter and Summer of 2025 electric prices (Euros/MWh)
    elec_mean = (elec_df['Winter 2025'].mean() + elec_df['Summer 2025'].mean())/2

    # Standard variation of Winter and Summer of 2025 electric prices (Euros/MWh)
    elec_std = (elec_df['Winter 2025'].std() + elec_df['Summer 2025'].std())/2

    elec_mean = elec_mean / 1000
    elec_std = elec_std / 1000


    print(f'Electricity prices:\n - Mean: ${elec_mean}€/Mwh\n - Std:  ${elec_std}€/Mwh')
    return elec_mean, elec_std

def generate_capacities(csv_file:str, nb_vehicles:int, seed:int=42, sep:str=';'):
    cap_df = pd.read_csv(filepath_or_buffer=csv_file, sep=sep)

    # Getting back all kind of battery capacities with unique values
    all_capacities = cap_df['Battery Capacity kwh'].dropna().unique()

    # Extracting random values for generating the array of capacities
    capacities = pd.Series(all_capacities).sample(n=nb_vehicles, random_state=seed)

    print(f'Capacities of vehicles (kwh): ${capacities}')
    return capacities.tolist()

def get_power_constants(nb_vehicles:int, nb_consumers:int=67000000):
    mean_consumption = (87028 + 46847 + 52374 + 29819)/4 # Mean of consumption in France in 2025 (estimate according to data/grid_capacity.txt)
    sim_ratio = nb_vehicles / nb_consumers # Ratio to reduce A_max of simulation to realistic restrictions

    a_max = sim_ratio * mean_consumption
    x_max = a_max / nb_vehicles # For init, uniform charging/discharging for every vehicle
    x_min = -x_max
    return a_max, x_max, x_min



def run_scenario(scenario_name, capacities:list, price_mean:float, price_std:float, model_type=None, n:int=20, t:int=30, w:float=0.4, c1:float=0.3, c2:float=0.2, archive_size:int=10, nb_vehicles:int=10, delta_t:int=60, nb_of_ticks:int=48):
    A_MAX, X_MAX, X_MIN = get_power_constants(nb_vehicles=nb_vehicles)


    print(f"\n--- Launching Scenario: {scenario_name} ---")
    # Simulation parameters
    params = {
        'A_max': A_MAX, 'price_mean': price_mean, 'price_std': price_std,
        'capacities': capacities, 'n': n, 't': t, 
        'w': w, 'c1': c1, 'c2': c2,
        'nb_vehicles': nb_vehicles, 'delta_t': delta_t, 'nb_of_ticks': nb_of_ticks,
        'x_min':X_MIN, 'x_max':X_MAX
    }
    
    # Instantiate extended class
    optimizer = SmartMOPSO(model_type=model_type, **params)
    
    start_time = time.time()
    
    # Run simulation
    optimizer.iterate()
        
    end_time = time.time()
    duration = end_time - start_time
    
    # Retrieve best f2 (e.g. from archive)
    best_f2 = min([p.f_best[1] for p in optimizer.archive]) if optimizer.archive else 0
    
    print(f"Finished in {duration:.2f} seconds.")
    print(f"Best f2 found: {best_f2:.4f}")
    
    return duration, best_f2, optimizer.archive



import matplotlib.pyplot as plt
import numpy as np

def plot_time_benchmark(nb_particles_list, results_dict):
    
    t_mopso = [item[0] for item in results_dict['MOPSO']]
    t_mlp   = [item[0] for item in results_dict['MLP']]
    t_rf    = [item[0] for item in results_dict['RF']]

    plt.figure(figsize=(10, 6))
    
    plt.plot(nb_particles_list, t_mopso, 'o-', label='Sans IA (MOPSO)', color='#1f77b4', linewidth=2)
    plt.plot(nb_particles_list, t_mlp,   's--', label='Avec MLP', color='#ff7f0e', linewidth=2)
    plt.plot(nb_particles_list, t_rf,    '^-.', label='Avec Random Forest', color='#2ca02c', linewidth=2)
    
    plt.title("Temps d'exécution selon le nombre de particules", fontsize=14, fontweight='bold')
    plt.xlabel("Nombre de Particules", fontsize=12)
    plt.ylabel("Temps (s)", fontsize=12)
    plt.grid(True, linestyle=':', alpha=0.7)
    plt.legend(fontsize=11)
    

    
    plt.tight_layout()
    plt.show()


import matplotlib.pyplot as plt

def plot_f2_benchmark(nb_particles_list, results_dict):
    s_mopso = [item[1] for item in results_dict['MOPSO']]
    s_mlp   = [item[1] for item in results_dict['MLP']]
    s_rf    = [item[1] for item in results_dict['RF']]


    plt.figure(figsize=(10, 6))
    
    plt.plot(nb_particles_list, s_mopso, 'o-', label='Sans IA (MOPSO)', color='#1f77b4', linewidth=2)
    plt.plot(nb_particles_list, s_mlp,   's--', label='Avec MLP', color='#ff7f0e', linewidth=2)
    plt.plot(nb_particles_list, s_rf,    '^-.', label='Avec Random Forest', color='#2ca02c', linewidth=2)
    
    plt.title("Meilleur Score F2 (Convergence) selon le nombre de particules", fontsize=14, fontweight='bold')
    plt.xlabel("Nombre de Particules (log scale)", fontsize=12)
    plt.ylabel("Meilleur F2 Score", fontsize=12)
    plt.grid(True, linestyle=':', alpha=0.7)
    plt.legend(fontsize=11)
    
    plt.xscale('log') 
    
    plt.tight_layout()
    plt.show()






def main():

    # CSV files
    elec_price_csv = 'data/elec_prices.csv'
    capacity_csv = 'data/vehicle_capacity.csv'

    # Global Simulation parameters
    T = 30 # Number of iterations (for the particles)
    W = 0.4 # Inertia (for exploration)
    C1 = 0.3 # Individual trust
    C2 = 0.2 # Social trust
    ARC_SIZE = 10 # Archive size
    nb_vehicle = 20

    P_MEAN, P_STD = calculate_elec_prices(elec_price_csv)
    CAPACITIES = generate_capacities(capacity_csv, nb_vehicles=nb_vehicle)

    NB_TICKS = 48
    DELTA = 60

    results = {
        'MOPSO':[],
        'MLP': [],
        'RF': []
    }
    nb_particles = [20,50,500,1000,10000]

    for k in range(len(nb_particles)):
        # 1. Without Surrogate (Baseline)
        d1, f1_score, _ = run_scenario(
            "Only MOPSO", 
            capacities=CAPACITIES, 
            price_mean=P_MEAN, 
            price_std=P_STD, 
            nb_vehicles=nb_vehicle,  # Important pour la cohérence
            model_type=None,
            n=nb_particles[k]
        )
        results['MOPSO'].append((d1, f1_score))

        # 2. With MLP
        d2, f2_score, _ = run_scenario(
            "With MLP", 
            capacities=CAPACITIES, 
            price_mean=P_MEAN, 
            price_std=P_STD, 
            nb_vehicles=nb_vehicle,
            model_type='mlp',
            n=nb_particles[k]
        )
        results['MLP'].append((d2, f2_score))

        # 3. With Random Forest
        d3, f3_score, _ = run_scenario(
            "With Random Forest", 
            capacities=CAPACITIES, 
            price_mean=P_MEAN, 
            price_std=P_STD, 
            nb_vehicles=nb_vehicle,
            model_type='rf',
            n=nb_particles[k]
        )
        results['RF'].append((d3, f3_score))
        # --- DISPLAY RESULTS ---
    
    
    print("\n=== SUMMARY ===")
    print(f"{'Mode':<15} | {'Time (s)':<10} | {'Best f2':<10}")
    print("-" * 45)
    for k, v in results.items():
        for i in range(len(nb_particles)):
            print(f"{k:<15}_{nb_particles[i]:<15} | {v[i][0]:<10.2f} | {v[i][1]:<10.4f}")
    
    plot_time_benchmark(nb_particles, results)
    plot_f2_benchmark(nb_particles, results)



main()