"""
create_NN_dataset.py — Generazione del dataset per la rete neurale del surrogato

Genera il dataset di training e test per la rete neurale surrogata, usando una
strategia di exploration-exploitation guidata da Gurobi.

Logica di campionamento:

  - "Exploit": campiona assortimenti vicini alla soluzione SAA media (freq_a), che
    rappresenta la distribuzione di frequenza di a[i]=1 sui singoli scenari.
  - "Explore": campiona assortimenti completamente casuali.
  - "Extreme": casi limite (assortimento pieno, vuoto, quantità saturate) per coprire
    le zone di frontiera dello spazio decisionale.
  - Per ogni campione (a, y): risolve il modello di ricorso esatto per ottenere il
    valore di target.

Il meccanismo usa un modello Gurobi ausiliario ("Generator") che, dato un target (a, y)
desiderato, trova il punto feasible più vicino nel politopo dei vincoli fisici. Questo
garantisce che tutti i campioni siano effettivamente soluzioni ammissibili.

Struttura del dataset prodotto:

  Ogni riga: [a_0, ..., a_N, y_0, ..., y_N, target_profit]
  - a_i binarie   : variabili di assortimento
  - y_i interi    : quantità ordinate
  - target_profit : profitto atteso del ricorso calcolato esattamente

Output:
  - train_data_assortment.csv : dataset di training (90% dei dati)
  - test_data_assortment.csv  : dataset di test (10% dei dati)
"""

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 models import *
from optimization import *
from prepare_data import *
import pandas as pd
from sklearn.model_selection import train_test_split

def main():
    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:
        print("Impossibile trovare il dataset. Inizio generazione:")

        #usiamo la funzione di prepare_data.py per generare il dataset di addestramento
        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)
        df_tot=prepare_NN_training_data(opt,struct_scenari,scenari,prob,n_samples=200000,n_extreme=11000,seed=1849)

        #ottenuti i valori, andiamo a salvarli in modo organizzato in due file dividendo a priori test e train
        num_p=opt.num_opzioni
        cols=[f'a_{i}' for i in range(num_p)]+[f'y_{i}' for i in range(num_p)]+['target_profit']
        df_tot.columns=cols

        df_tot=df_tot.drop_duplicates(subset=cols[:-1]).reset_index(drop=True) 
        df_tot=df_tot.sample(frac=1,random_state=42).reset_index(drop=True)

        df_train,df_test=train_test_split(df_tot,test_size=0.1,random_state=42)
        column=df_train.columns[:-1]   
        df_train[column]=df_train[column].astype(int)
        df_train=df_train.sample(frac=1,random_state=42).reset_index(drop=True)
        
        df_train.to_csv(dataset_filename,index=False)
        df_test.to_csv(dataset_filename_test,index=False)
        print("Train set e Test set creati e splittati.")
        return df_train, df_test

if __name__ == "__main__":
    main()