import pandas as pd
# Carica il JSON in un DataFrame
df = pd.read_json("all_T.json")

# Rimuovi o “parsa” le colonne stringa ridondanti
#    - "Scenario" → ridondante con "Orifice diameter (m)"
#    - "Weather"  → ridondante (hai già Stability class, Wind, Air T, Humidity, Atm stability)
df = df.drop(columns=["Scenario", "Weather", "Pressure vessel", "Stability class"])

# Rinomina le colonne per togliere unità e semplificare i nomi
df = df.rename(columns={
    "Composition H2 (%vol)"                                : "composition_h2",
    "Storage T (degC)"                                     : "storage_temp",
    "Storage P (bar)"                                      : "storage_pressure",
    "Elevation of release (m)"                             : "elevation",
    "Jet fire mass rate (kg/s)"                            : "mass_rate",
    "Flame length (m)"                                     : "flame_length",
    "Distance to intervention zone (5 kW/m2) (m)"           : "dist_intervention",
    "Distance to alert zone (3 kW/m2) (m)"                  : "dist_alert",
    "Distance to domino effect (8 kW/m2) (m)"               : "dist_domino",
    "Orifice diameter (m)"                                 : "orifice_diameter",
    "Ideal MW (g/mol)"                                     : "ideal_mw",
    "Wind (m/s)"                                           : "wind_speed",
    "Air T (degC)"                                         : "air_temp",
    "Humidity (%)"                                         : "humidity",
    "Atm stability"                                        : "atm_stability"
})

print("\n---------CHECK-----------\n")
# Assicurati che tutti i campi siano davvero numerici
df = df.apply(pd.to_numeric, errors="raise")

# Ora stampi normalmente
print(df.dtypes)
print(df.head(20))  # o print(d) per DataFrame

# Ottiene il numero di righe e colonne
n_rows, n_cols = df.shape
print(f"Numero di righe: {n_rows}, numero di colonne: {n_cols}")

# Calcola statistiche descrittive sulle colonne numeriche
stats = df.describe()
print(stats)

# Conta i valori NaN per colonna
missing_per_col = df.isnull().sum()
print("Valori mancanti per colonna:")
print(missing_per_col)
print("\n---------END CHECK-----------\n")


# ------------------------
# PREPROCESSING STANDARDSCALER
# ------------------------

from sklearn.model_selection import train_test_split      # per dividere train/test
from sklearn.preprocessing import StandardScaler          # per la standardizzazione

# 1) Seleziona features (X) e target (y)

feature_cols = [
    "composition_h2", "storage_temp", "storage_pressure", "elevation",
    "orifice_diameter", "ideal_mw", "wind_speed", "air_temp",
    "humidity", "atm_stability"
]
target_cols = [
    "mass_rate",
    "flame_length",
    "dist_intervention",
    "dist_alert",
    "dist_domino"
]

X = df[feature_cols].values    # array NumPy di shape (N, n_features)
y = df[target_cols].values     # array NumPy di shape (N, n_targets)

# 2) Primo split: TRAIN+VAL (90%) vs TEST (10%)
X_temp, X_test_raw, y_temp, y_test_raw = train_test_split(
    X, y,
    test_size=0.10,           # 10% per il test finale
    random_state=42,
    shuffle=True
)

# 3) Secondo split: da quel 90% ricavo TRAIN (80%) e VAL (10%)
#    Poiché X_temp è il 90% del totale, per ottenere il 10% del totale uso test_size = 10/90 = 1/9 ≃ 0.111...
X_train_raw, X_val_raw, y_train_raw, y_val_raw = train_test_split(
    X_temp, y_temp,
    test_size=1/9,            # ~0.111 per avere 10% del totale
    random_state=42,
    shuffle=True
)

# Controllo shape
print(f"X_train_raw: {X_train_raw.shape}")  # ~80% del totale
print(f"X_val_raw:   {X_val_raw.shape}")    # ~10% del totale
print(f"X_test_raw:  {X_test_raw.shape}")   # 10% del totale

# 3) Inizializza gli scaler per X e y
scaler_X = StandardScaler()
scaler_y = StandardScaler()

# 5) Fit e trasformazione dello scaler_X su TRAIN, poi applica a VAL e TEST
scaler_X.fit(X_train_raw)
X_train = scaler_X.transform(X_train_raw)
X_val   = scaler_X.transform(X_val_raw)
X_test  = scaler_X.transform(X_test_raw)

# 6) Fit e trasformazione dello scaler_y su TRAIN, poi applica a VAL e TEST
#    reshape(-1, n_targets) è già 2D, quindi possiamo trasformare direttamente
scaler_y.fit(y_train_raw)
y_train = scaler_y.transform(y_train_raw)
y_val   = scaler_y.transform(y_val_raw)
y_test  = scaler_y.transform(y_test_raw)

import torch
# 7) Conversione in tensori PyTorch float32
X_train = torch.from_numpy(X_train).float()
y_train = torch.from_numpy(y_train).float()
X_val   = torch.from_numpy(X_val).float()
y_val   = torch.from_numpy(y_val).float()
X_test  = torch.from_numpy(X_test).float()
y_test  = torch.from_numpy(y_test).float()

print("\n---------CHECK-----------\n")

# Shape dei dataset
print("\n--- Shapes dopo preprocess ---")
print("X_train:", X_train.shape, "y_train:", y_train.shape)
print("X_val:  ", X_val.shape,   "y_val:  ", y_val.shape)
print("X_test: ", X_test.shape,  "y_test: ", y_test.shape)

# Tipo dei dati
print("dtype X_train:", X_train.dtype)
print("dtype y_train:", y_train.dtype)
print("\n---------END CHECK-----------\n")

import torch.nn as nn

class JetFireModelReg(nn.Module):
    def __init__(self, D_in=10, H1=32, H2=16, D_out=5, p_dropout=0.3):
    
        # D_in  = numero di feature in input (qui 10)
        # H     = numero di neuroni in ciascun hidden layer
        # D_out = numero di target da prevedere (qui 5)
        # Hn    = numero di hidden layer

        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(D_in, H1),
            nn.ReLU(),
            nn.Dropout(p_dropout),        # regolarizzazione
            nn.Linear(H1, H2),
            nn.ReLU(),
            nn.Dropout(p_dropout),
            nn.Linear(H2, D_out)          # output lineare
        )

    def forward(self, x):
        return self.net(x)


### Setup Dataloader
from torch.utils.data import Dataset, DataLoader

class JetFireDataset(Dataset):
    def __init__(self, X, y):
        # X: tensor float32 (N, n_features)
        # y: tensor float32 (N, n_targets)
        self.features = X
        self.targets  = y

    def __len__(self):
        # Restituisce numero di campioni
        return len(self.targets)

    def __getitem__(self, idx):
        # Restituisce la coppia (feature, target) al sample idx
        return self.features[idx], self.targets[idx]

# Creazione dei DataLoader per train e validation
train_ds = JetFireDataset(X_train, y_train)  # X_train, y_train già tensori float32
val_ds   = JetFireDataset(X_val,   y_val)
test_ds     = JetFireDataset(X_test, y_test)

train_loader = DataLoader(train_ds, batch_size=32, shuffle=True)  # batch_size=32 consigliato :contentReference[oaicite:5]{index=5}
val_loader   = DataLoader(val_ds,   batch_size=64, shuffle=False) # validation non necessita shuffle
test_loader = DataLoader(test_ds, batch_size=64, shuffle=False)

# Setup di modello, loss, ottimizzatore e dispositivo
import torch.optim as optim

# Istanziamento, optimizer con L2 weight‐decay, criterio
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = JetFireModelReg(D_in=10, H1=32, H2=16, D_out=5, p_dropout=0.3).to(device)

criterion = nn.MSELoss()
optimizer = optim.Adam(
    model.parameters(),
    lr=1e-3,
    weight_decay=1e-4   # L2 regularization
)

# Early stopping setup
import copy
patience = 10
best_val_loss = float('inf')
best_model_wts = copy.deepcopy(model.state_dict())
epochs_no_improve = 0
num_epochs = 100

for epoch in range(1, num_epochs+1):
    # — train —
    model.train()
    running_train = 0.0
    for Xb, yb in train_loader:
        Xb, yb = Xb.to(device), yb.to(device)
        optimizer.zero_grad()
        y_pred = model(Xb)
        loss = criterion(y_pred, yb)
        loss.backward()
        optimizer.step()
        running_train += loss.item() * Xb.size(0)
    train_loss = running_train / len(train_loader.dataset)

    # — validation —
    model.eval()
    running_val = 0.0
    with torch.no_grad():
        for Xb, yb in val_loader:
            Xb, yb = Xb.to(device), yb.to(device)
            loss = criterion(model(Xb), yb)
            running_val += loss.item() * Xb.size(0)
    val_loss = running_val / len(val_loader.dataset)

    print(f"Epoch {epoch:03d}  Train Loss: {train_loss:.6f}  Val Loss: {val_loss:.6f}")

    # — controllo early stopping —
    if val_loss < best_val_loss - 1e-6:
        best_val_loss = val_loss
        best_model_wts = copy.deepcopy(model.state_dict())
        epochs_no_improve = 0
    else:
        epochs_no_improve += 1
        if epochs_no_improve >= patience:
            print(f"Early stopping alla epoca {epoch}")
            break


import matplotlib.pyplot as plt
import numpy as np
from sklearn.metrics import mean_squared_error, mean_absolute_error


# al termine del train & eval, sui dati di validazione:
y_val_pred_norm = model(X_val).detach().cpu().numpy()
# riporto sulla scala reale
y_val_pred = scaler_y.inverse_transform(y_val_pred_norm)
y_val_true = scaler_y.inverse_transform(y_val.cpu().numpy())

from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score

print("Metriche sui dati di validation (per target):")
for i, name in enumerate(target_cols):
    mae_i  = mean_absolute_error(y_val_true[:,i], y_val_pred[:,i])
    rmse_i = mean_squared_error(y_val_true[:,i], y_val_pred[:,i])
    r2_i   = r2_score(y_val_true[:,i], y_val_pred[:,i])
    print(f"{name}: MAE={mae_i:.3f}, RMSE={rmse_i:.3f}, R2={r2_i:.3f}")


# Carica i pesi migliori
model.load_state_dict(best_model_wts)
print(f"Best MSE: {best_val_loss:.4f}, RMSE: {np.sqrt(best_val_loss):.4f}")

# ——————————————————————————————————————————————
# 2) Predizioni sul validation set (ritorno alla scala originale)
# ——————————————————————————————————————————————

model.eval()
with torch.no_grad():
    # y_val_scaled: (N_val, 5) in scala standardizzata
    y_val_scaled = model(X_val.to(device)).cpu().numpy()

# Inverti la standardizzazione sui target
# scaler_y è lo StandardScaler che hai usato su y_train_raw
y_val_pred = scaler_y.inverse_transform(y_val_scaled)
y_val_true = scaler_y.inverse_transform(y_val.cpu().numpy())

# ——————————————————————————————————————————————
# 3) Calcolo dei metriche (MAE, RMSE) per ciascuna uscita
# ——————————————————————————————————————————————

mae_values  = mean_absolute_error(y_val_true, y_val_pred, multioutput='raw_values')
rmse_values = np.sqrt(mean_squared_error(y_val_true, y_val_pred, multioutput='raw_values'))


# ——————————————————————————————————————————————
# 4) Scatter True vs Predicted per ciascun target
# ——————————————————————————————————————————————

for i, col in enumerate(target_cols):
    plt.figure(figsize=(6,6))
    # Scatter dei valori veri contro quelli predetti
    plt.scatter(y_val_true[:,i], y_val_pred[:,i], alpha=0.6)
    # Diagonale ideale y=x
    lims = [
        np.min([y_val_true[:,i].min(), y_val_pred[:,i].min()]),
        np.max([y_val_true[:,i].max(), y_val_pred[:,i].max()]),
    ]
    plt.plot(lims, lims, '--', linewidth=1, label='y = x')
    plt.xlabel(f"True {col}")
    plt.ylabel(f"Predicted {col}")
    plt.title(f"{col}: True vs Predicted")
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    plt.show()

# ——————————————————————————————————————————————
# 5) Valutazione finale sul TEST set (10% mai toccato)
# ——————————————————————————————————————————————

# 5.1 Carica i pesi migliori
model.load_state_dict(best_model_wts)
model.eval()

# 5.2 Calcola la MSE media sul test_loader
test_mse = 0.0
with torch.no_grad():
    for Xb, yb in test_loader:
        Xb, yb = Xb.to(device), yb.to(device)
        test_mse += criterion(model(Xb), yb).item() * Xb.size(0)
test_mse /= len(test_loader.dataset)
print(f"\nTest Loss (MSE): {test_mse:.6f}, RMSE: {np.sqrt(test_mse):.6f}")

# 5.3 Predizioni su tutta la X_test e metriche per target
# (riportiamo tutto sulla scala reale)
with torch.no_grad():
    y_test_pred_norm = model(X_test.to(device)).cpu().numpy()

y_test_pred = scaler_y.inverse_transform(y_test_pred_norm)
y_test_true = scaler_y.inverse_transform(y_test.cpu().numpy())

from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score

print("\nMetriche sul TEST set (per target):")
for i, name in enumerate(target_cols):
    mae_i  = mean_absolute_error(y_test_true[:,i], y_test_pred[:,i])
    rmse_i = mean_squared_error(y_test_true[:,i], y_test_pred[:,i])
    r2_i   = r2_score(y_test_true[:,i], y_test_pred[:,i])
    print(f"{name}: MAE={mae_i:.3f}, RMSE={rmse_i:.3f}, R2={r2_i:.3f}")

# 5.4 (Opzionale) scatter true vs predetti sul test set
import matplotlib.pyplot as plt

for i, col in enumerate(target_cols):
    plt.figure(figsize=(5,5))
    plt.scatter(y_test_true[:,i], y_test_pred[:,i], alpha=0.6)
    m = max(y_test_true[:,i].max(), y_test_pred[:,i].max())
    plt.plot([0,m], [0,m], '--', color='grey')
    plt.xlabel(f"True {col}")
    plt.ylabel(f"Pred {col}")
    plt.title(f"Test set: {col}")
    plt.grid(True)
    plt.show()
