"""
train_NN.py — Training della rete neurale per il surrogato

Addestra un MLPRegressor (scikit-learn) per stimare il valore atteso del ricorso
in funzione delle sole quantità ordinate y.

Come funziona:

  1. Carica il dataset da train_data_assortment.csv (generato da create_NN_dataset.py)
     o lo genera se non esiste.
  2. Usa come input X le colonne y_i (quantità), come target la colonna target_profit.
  3. Split: 85% temp / 15% test; poi temp lo divide 80% train / 20% validation.
  4. Normalizzazione: MinMaxScaler su X, StandardScaler su y.
  5. Hyperparameter tuning con OPTUNA (30 trial):
       - n_layers: 1 o 2
       - n_units per layer: da 16 a 64 (step 8)
       - alpha (L2 regularization): 1e-5 ... 1e-2
       - learning_rate_init: 1e-4 ... 1e-2
  6. Riaddestra la rete migliore su train+val.
  7. Salva i modelli.

Output:
  - surrogate_mlp.pkl    : rete neurale surrogata finale
  - scaler_input.pkl     : MinMaxScaler per le quantità y
  - scaler_target.pkl    : StandardScaler per il target (profitto)

Nota importante: i file .pkl devono trovarsi nella stessa directory da cui si lancia main_NN.py.
"""

import sys
import os
cartella_progetto = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
if cartella_progetto not in sys.path:
    sys.path.append(cartella_progetto)
from prepare_data import *
from optimization import *
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, MinMaxScaler
from sklearn.neural_network import MLPRegressor
import os
import joblib
import optuna
import optuna.visualization as vis
import create_NN_dataset

dataset_filename="train_data_assortment.csv"
dataset_filename_test="test_data_assortment.csv"

data=prepare_input_data(force_recalculate=True, SAA_flag=True, n_scenari_domanda_SAA=30, n_warm_scenarios=30, used_seed=1673, nome_file='dati_progetto_NN.pkl')
scenari=data["saa_input"]["demand_weekly"]
prob=data["saa_input"]["scenario_prob"]

struct_scenari=data["saa_input"]["warm_demand"]

if os.path.exists(dataset_filename) and os.path.exists(dataset_filename_test):
    print(f"Dataset find: upload '{dataset_filename}' and '{dataset_filename_test}' ")
    df_train=pd.read_csv(dataset_filename)
    df_test=pd.read_csv(dataset_filename_test)
else:
    df_train, df_test=create_NN_dataset.main()

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

#Split in Train,test e Validation set(85(80/20)/15)
X_temp,X_test,y_temp,y_test=train_test_split(X_quantity,y,test_size=0.15,random_state=42)

X_train,X_val,y_train,y_val=train_test_split(X_temp,y_temp,test_size=0.2,random_state=42)

scaler_y=StandardScaler()
y_scaled=scaler_y.fit_transform(y_train.reshape(-1,1)).ravel()

scaler_X=MinMaxScaler()
X_scaled=scaler_X.fit_transform(X_train)

X_temp_scaled=scaler_X.transform(X_temp)
y_temp_scaled=scaler_y.transform(y_temp.reshape(-1,1)).ravel()

X_val_scaled=scaler_X.transform(X_val)
y_val_scaled=scaler_y.transform(y_val.reshape(-1,1)).ravel()

X_test_scaled=scaler_X.transform(X_test)
y_test_scaled=scaler_y.transform(y_test.reshape(-1,1)).ravel()


def objective(trial):
    n_layers=trial.suggest_int('n_layers',1,2) 
    layers=[]
    for i in range(n_layers):
        n_units=trial.suggest_int(f'n_units_l{i}',16,64,step=8)
        layers.append(n_units)

    hidden_layer_sizes=tuple(layers)


    alpha=trial.suggest_float('alpha', 1e-5, 1e-2, log=True)
    learning_rate_init=trial.suggest_float('learning_rate_init', 1e-4, 1e-2, log=True)

    mlp=MLPRegressor(hidden_layer_sizes=hidden_layer_sizes,activation='relu',solver='adam',alpha=alpha,batch_size='auto',learning_rate_init=learning_rate_init,max_iter=1000,early_stopping=True,validation_fraction=0.1,random_state=42,verbose=False)
    mlp.fit(X_scaled,y_scaled)

    val_score=mlp.score(X_val_scaled,y_val_scaled)

    return val_score

study=optuna.create_study(direction='maximize',study_name='Gurobi_nn_parameter')

study.optimize(objective,n_trials=30)
print(f"Best R² Score over the Validation Set: {study.best_value:.4f}")
print("Find the best parameters")
for k, v in study.best_params.items():
    print(f"  {k}: {v}")

best_params=study.best_params
best_layers=[best_params[f'n_units_l{i}'] for i in range(best_params['n_layers'])]

print("Starting the network lerning phase")

best_mlp=MLPRegressor(hidden_layer_sizes=tuple(best_layers),activation='relu',solver='adam',learning_rate_init=best_params['learning_rate_init'],alpha=best_params['alpha'],max_iter=1000,batch_size='auto',early_stopping=True,validation_fraction=0.1,random_state=42,verbose=True)

best_mlp.fit(X_temp_scaled,y_temp_scaled)

train_score=best_mlp.score(X_temp_scaled,y_temp_scaled)
val_score=best_mlp.score(X_test_scaled,y_test_scaled)
print(f"R² Score over Train Set: {train_score:.4f}")
print(f"R² Score over Test Set: {val_score:.4f}")


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

print("Models saved succesfully:'scaler_input.pkl', 'surrogate_mlp.pkl' and 'scaler_target.pkl'")


print("\nOptuna graph generation")

vis.plot_optimization_history(study).show()
vis.plot_param_importances(study).show()
vis.plot_slice(study).show()



