# # File: src/samudra_ai/trainer.py
# file lama, udah OK


from sklearn.preprocessing import MinMaxScaler
from sklearn.model_selection import train_test_split
import numpy as np
import tensorflow as tf
import random
import matplotlib.pyplot as plt
import os

def prepare_training_data(x_data, y_data, time_seq, seed):
    np.random.seed(seed)
    random.seed(seed)
    tf.random.set_seed(seed)

    min_samples = min(x_data.shape[0], y_data.shape[0])

    x_np = np.nan_to_num(
        x_data.isel(time=slice(0, min_samples)).values,
        nan=0.0
    ).astype(np.float32)

    y_np = np.nan_to_num(
        y_data.isel(time=slice(0, min_samples)).values,
        nan=0.0
    ).astype(np.float32)

    scaler_X = MinMaxScaler()
    X_scaled = scaler_X.fit_transform(
        x_np.reshape(x_np.shape[0], -1)
    ).reshape(x_np.shape).astype(np.float32)

    scaler_y = MinMaxScaler(feature_range=(-1, 1))
    y_scaled = scaler_y.fit_transform(
        y_np.reshape(y_np.shape[0], -1)
    ).reshape(y_np.shape).astype(np.float32)

    n_samples = len(X_scaled) - time_seq + 1

    X_seq = np.empty((n_samples, time_seq, *X_scaled.shape[1:]), dtype=np.float32)
    y_seq = np.empty((n_samples, *y_scaled.shape[1:]), dtype=np.float32)

    for i in range(n_samples):
        X_seq[i] = X_scaled[i:i+time_seq]
        y_seq[i] = y_scaled[i+time_seq-1]

    # channel dim
    X_seq = X_seq[..., np.newaxis]
    y_seq = y_seq[..., np.newaxis]

    X_train, X_test, y_train, y_test = train_test_split(
        X_seq, y_seq, test_size=0.2, random_state=seed, shuffle=False
    )

    return X_train, X_test, y_train, y_test, scaler_X, scaler_y

def plot_training_history(history, output_dir=None):
    if output_dir: os.makedirs(output_dir, exist_ok=True)
    plt.figure(figsize=(10, 6))
    plt.plot(history.history['loss'], label='Train Loss')
    plt.plot(history.history['val_loss'], label='Val Loss')
    plt.title("Kurva Loss Training")
    plt.xlabel("Epoch")
    plt.ylabel("Loss")
    plt.legend()
    if output_dir:
        plt.savefig(os.path.join(output_dir, "loss_curve.png"))
    plt.show()

    plt.figure(figsize=(10, 6))
    plt.plot(history.history['mae'], label='Train MAE')
    plt.plot(history.history['val_mae'], label='Val MAE')
    plt.title("Kurva MAE Training")
    plt.xlabel("Epoch")
    plt.ylabel("MAE")
    plt.legend()
    if output_dir:
        plt.savefig(os.path.join(output_dir, "mae_curve.png"))
    plt.show()

# # FILE BARU 

# import logging
# from sklearn.preprocessing import MinMaxScaler
# from sklearn.model_selection import train_test_split
# import numpy as np
# import tensorflow as tf
# import random
# import matplotlib.pyplot as plt
# import os
# import re

# # Setup logger
# logger = logging.getLogger(__name__)

# def detect_precip_var(data):
#     """
#     Deteksi nama variabel yang merujuk ke precipitation.
#     Args:
#         data: xr.DataArray atau xr.Dataset.
#     Returns:
#         str: Nama variabel yang terdeteksi (misal: 'precip', 'pr', 'rain').
#     """
#     # Daftar nama yang mungkin untuk precipitation
#     possible_names = ['precip', 'pr', 'precipitation', 'rain', 'ch']
    
#     # Jika data adalah DataArray, cek nama variabelnya
#     if hasattr(data, 'name') and data.name in possible_names:
#         return data.name
    
#     # Jika data adalah Dataset, cek nama variabel di dalamnya
#     elif hasattr(data, 'variables'):
#         for var in data.variables:
#             if var in possible_names:
#                 return var
    
#     # Jika tidak ditemukan, cari pola yang mirip
#     for var in possible_names:
#         if re.search(var, str(data), re.IGNORECASE):
#             return var
    
#     # Jika tetap tidak ditemukan, kembalikan default
#     logger.warning(f"Variabel precipitation tidak terdeteksi. Menggunakan default 'precip'.")
#     return 'precip'

# def prepare_training_data(x_data, y_data, time_seq, seed):
#     """
#     Persiapkan data training dengan scaling per parameter.
#     Return:
#         - X_train, X_test, y_train, y_test: Data yang sudah di-split.
#         - scaler_X: Dict scaler untuk input features (misal: {'ch': MinMaxScaler(), 'tp': MinMaxScaler()}).
#         - scaler_y: Dict scaler untuk target (misal: {'precipitation': MinMaxScaler()}).
#     """
#     np.random.seed(seed)
#     random.seed(seed)
#     tf.random.set_seed(seed)

#     # Pastikan jumlah sampel sama antara x_data dan y_data
#     min_samples = min(x_data.shape[0], y_data.shape[0])
#     x_np = np.nan_to_num(x_data.isel(time=slice(0, min_samples)).values)
#     y_np = np.nan_to_num(y_data.isel(time=slice(0, min_samples)).values)

#     # === Deteksi nama variabel precipitation ===
#     precip_var = detect_precip_var(y_data)
#     logger.info(f"Variabel precipitation terdeteksi: '{precip_var}'")

#     # === Scaling per parameter untuk input (x_data) ===
#     scaler_X = {}
#     X_scaled = np.zeros_like(x_np)
    
#     # Asumsikan x_data memiliki shape (time, lat, lon, parameter)
#     # Misal: parameter = ['ch', 'tp', 'pr']
#     for i in range(x_np.shape[-1]):  # Loop per parameter
#         param_name = x_data.dims[-1] if len(x_data.dims) == 4 else f'param_{i}'
#         scaler_X[param_name] = MinMaxScaler()
        
#         # Reshape untuk scaling: (time, lat*lon) -> scale -> reshape kembali
#         param_data = x_np[..., i].reshape(x_np.shape[0], -1)
#         X_scaled[..., i] = scaler_X[param_name].fit_transform(param_data).reshape(x_np.shape[0], x_np.shape[1], x_np.shape[2])
        
#         # Log info scaler
#         logger.info(f"Scaler untuk input parameter '{param_name}': {scaler_X[param_name].__class__.__name__} "
#                     f"(min={scaler_X[param_name].min_[0]:.2f}, max={scaler_X[param_name].scale_[0]:.2f})")

#     # === Scaling untuk target (y_data) ===
#     scaler_y = {}
#     y_scaled = np.zeros_like(y_np)
    
#     # Gunakan nama variabel yang terdeteksi
#     scaler_y[precip_var] = MinMaxScaler(feature_range=(-1, 1))
    
#     # Reshape untuk scaling: (time, lat*lon) -> scale -> reshape kembali
#     y_scaled = scaler_y[precip_var].fit_transform(y_np.reshape(y_np.shape[0], -1)).reshape(y_np.shape)
    
#     # Log info scaler
#     logger.info(f"Scaler untuk target '{precip_var}': {scaler_y[precip_var].__class__.__name__} "
#                 f"(min={scaler_y[precip_var].min_[0]:.2f}, max={scaler_y[precip_var].scale_[0]:.2f})")

#     # === Buat sequence data ===
#     X_seq = np.array([X_scaled[i:i+time_seq] for i in range(len(X_scaled) - time_seq)])
#     y_seq = np.array([y_scaled[i+time_seq-1] for i in range(len(y_scaled) - time_seq)])

#     # Split data
#     X_train, X_test, y_train, y_test = train_test_split(
#         X_seq, y_seq, test_size=0.2, random_state=seed, shuffle=False
#     )

#     return X_train, X_test, y_train, y_test, scaler_X, scaler_y


# def plot_training_history(history, output_dir=None):
#     """Plot kurva loss dan MAE."""
#     if output_dir: 
#         os.makedirs(output_dir, exist_ok=True)
    
#     # Plot Loss
#     plt.figure(figsize=(10, 6))
#     plt.plot(history.history['loss'], label='Train Loss')
#     plt.plot(history.history['val_loss'], label='Val Loss')
#     plt.title("Kurva Loss Training")
#     plt.xlabel("Epoch")
#     plt.ylabel("Loss")
#     plt.legend()
#     if output_dir:
#         plt.savefig(os.path.join(output_dir, "loss_curve.png"))
#     plt.show()

#     # Plot MAE
#     plt.figure(figsize=(10, 6))
#     plt.plot(history.history['mae'], label='Train MAE')
#     plt.plot(history.history['val_mae'], label='Val MAE')
#     plt.title("Kurva MAE Training")
#     plt.xlabel("Epoch")
#     plt.ylabel("MAE")
#     plt.legend()
#     if output_dir:
#         plt.savefig(os.path.join(output_dir, "mae_curve.png"))
#     plt.show()