"""
active_learning.py — Loop di Active Learning per migliorare la NN del surrogato

Migliora iterativamente la rete neurale surrogata concentrandosi sulle zone dello
spazio decisionale che l'ottimizzatore effettivamente visita. 

Algoritmo (per ogni iterazione):

  1. Ottimizza il modello NN in Gurobi e trova (a*, y*) che massimizza l'obiettivo predetto.
  2. Valuta (a*, y*) con il ricorso esatto (build_recurse_only + solve_recurse).
  3. Calcola gap = |Q_predetto - Q_reale|.
  4. Se gap ≤ tolleranza (110 €) oppure Q_reale ≥ soglia_target: termina.
  5. Altrimenti: aggiunge nuovi punti al dataset:
       - 30 copie esatte di (y*, Q_reale)
       - 50 perturbazioni di y* (variazione del 5-15% su 3 prodotti casuali), ognuna valutata
         con il ricorso esatto, aggiunta 20 volte al dataset
  6. Riaddestra la NN su tutto il dataset aggiornato (warm_start).
  7. Salva i modelli aggiornati su disco.

Che file servono:
  - surrogate_mlp.pkl, scaler_input.pkl, scaler_target.pkl (da train_NN.py)
  - train_data_assortment.csv, test_data_assortment.csv (da create_NN_dataset.py)
  - dati_progetto_NN.pkl (da prepare_input_data)

Parametri principali (modificabili nel file, dopo la parte di acquisizione del dataset):
  - max_iter_loop  : numero massimo di iterazioni del loop (default: 30)
  - tolerance      : gap accettabile in € tra predetto e reale (default: 110)
  - perturbations  : numero di perturbazioni per iterazione (default: 50)
"""

import os
import sys
cartella_progetto = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
if cartella_progetto not in sys.path:
    sys.path.append(cartella_progetto)
import joblib
import pandas as pd
import numpy as np
from gurobipy import GRB
from prepare_data import *
from optimization import *


data=prepare_input_data(force_recalculate=False,SAA_flag=True,n_scenari_domanda_SAA=50,n_warm_scenarios=30,used_seed=1673,nome_file='dati_progetto_NN.pkl')

opt=AssortmentOptimizer(data["base"]["prodotti_tab"],data["base"]["sub_matrix"],data["base"]["prezzi"],data["base"]["cost"],data["base"]["holding"],data["base"]["capacita_scaffale_dict"],data["base"]["capacita_frigo_dict"],data["base"]["costo_fisso_scaffale"],data["base"]["penalty_lost_sales"],data["base"]["penale_SL"],data["base"]["SL_target"],flag_shelf=True)
scenari=data["saa_input"]["demand_weekly"]
prob=data["saa_input"]["scenario_prob"]
opt._mean_parameters(scenari,prob) #definisce i valori medi dei paramentri economici usando il vettore di probabilità fornito

dataset_train_file="train_data_assortment.csv"
dataset_test_file="test_data_assortment.csv"
if not os.path.exists(dataset_train_file):
    raise FileNotFoundError(f"Errore: Il file {dataset_train_file} deve esister per fare Active Learning.")

if not os.path.exists(dataset_test_file):
    raise FileNotFoundError(f"Il dataset di test non esiste")


#pulizia del dataset e formazione degli elemnti da usare 
df_train=pd.read_csv(dataset_train_file)
df_test=pd.read_csv(dataset_test_file)
num_p=int((df_train.shape[1]-1)/2)
X_quantity=df_train.iloc[:,num_p:2*num_p].values
y=df_train.iloc[:,-1].values

X_train_full=list(X_quantity)
y_train_full=list(y)

final_mlp=joblib.load('surrogate_mlp.pkl')
scaler_X=joblib.load('scaler_input.pkl')
scaler_y=joblib.load('scaler_target.pkl')
best_params=final_mlp.get_params()

final_mlp.warm_start=True #per warm-strat con pesi attuali
final_mlp.early_stopping=False
final_mlp.max_iter=150

max_iter_loop=30 #numero massimo di cicli di active learning 
tolerance=110.0  #tollearanza di delta tra soluzione vera e rete   
perturbations=50 #numero di campioni vicini a quello valutatao da aggiungere
 
opt.build_recurse_only(n_scenarios=30, write=False)
model_recurse=opt.models["recurse"]

for k in range(max_iter_loop):
    print(f"\n[Iterazione {k}/{max_iter_loop}] ottimizzazione primo stadio...")
    
    #ottimizzazione del problmea surrogato
    opt.build_model_NN(mlp_path='surrogate_mlp.pkl',scaler_path_input='scaler_input.pkl',scaler_path_target='scaler_target.pkl')
    model_nn=opt.models["nn"]
    model_nn.setParam('OutputFlag',0)
    model_nn.optimize()
    
    if model_nn.Status!=GRB.OPTIMAL:
        print("Non ci sono soluzioni ottime. Stop.")
        break

    #estrazione soluzioni ottenute
    a_opt=np.round(model_nn._a_nn.X).astype(int).flatten()
    y_opt=np.round(model_nn._y_nn.X).astype(int).flatten()
   
    #valutazione su ricorso vero
    predicted_recourse=model_nn.getVarByName("expected_recourse").X
    true_recourse=opt.solve_recurse(model=model_recurse,a_fixed=a_opt,y_fixed=y_opt)
    
    #valutazione del gap
    gap=abs(true_recourse-predicted_recourse)
    print(f"Predetto: {predicted_recourse:.2f} | Reale SAA: {true_recourse:.2f} | Gap: {gap:.2f}")
    
    if gap<=tolerance or true_recourse>=2160:
        print(f"Tolleranza raggiunta (Gap <= {tolerance}).")
        break

    #aggiunta per 30 volte la soluzione ottima con taget vero, che la rete la veda
    for _ in range(30):
        X_train_full.append(y_opt) 
        y_train_full.append(true_recourse)
    #aggiunta delle perturbazioni all'ottimo, con target annesso, per permettere di esplorale nelle vicinanze 
    prodotti_attivi=np.where(a_opt==1)[0]
    if len(prodotti_attivi)>0:
        for _ in range(perturbations):
            y_pert=y_opt.copy()
            da_alterare=np.random.choice(prodotti_attivi, size=min(3,len(prodotti_attivi)),replace=False) #cambiamo solo quantità di 3 prodotti
            for idx in da_alterare:
                variazione=max(1,int(y_pert[idx]*np.random.uniform(0.05, 0.15)))
                y_pert[idx]=max(1,y_pert[idx]+np.random.choice([-1, 1])*variazione)
            
            true_recourse_pert=opt.solve_recurse(model=model_recurse,a_fixed=a_opt,y_fixed=y_pert)
            for _ in range(20):
                X_train_full.append(y_pert) 
                y_train_full.append(true_recourse_pert)

    #aggiornamento scaler e riaddestramento rete
    X_scaled_updated=scaler_X.fit_transform(np.array(X_train_full))
    y_scaled_updated=scaler_y.fit_transform(np.array(y_train_full).reshape(-1,1)).ravel()
    final_mlp.best_loss_=np.inf
    final_mlp.fit(X_scaled_updated,y_scaled_updated)

    X_test=df_test.iloc[:,num_p:2*num_p].values
    y_test=df_test.iloc[:,-1].values

    X_test_scaled_loop=scaler_X.transform(X_test)
    y_test_scaled_loop=scaler_y.transform(y_test.reshape(-1,1)).ravel()
    
    score_attuale=final_mlp.score(X_test_scaled_loop,y_test_scaled_loop)
    print(f"R² su Test Set: {score_attuale:.4f}")

    joblib.dump(final_mlp,'surrogate_mlp.pkl')
    joblib.dump(scaler_X,'scaler_input.pkl')
    joblib.dump(scaler_y,'scaler_target.pkl')
