"""
data_prep.py — Pipeline de Préparation des Données
====================================================
PFA : Système de Détection d'Anomalies dans les Transactions Financières

Auteurs  : Othmane & Ramzi
Dataset  : Kaggle Credit Card Fraud Detection (creditcard.csv)
Sortie   : Datasets transformés dans ./data/ (format .pkl)

Usage    : python data_prep.py
"""

import os
import joblib
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from imblearn.over_sampling import SMOTE


# ─────────────────────────────────────────────────────────────
# Configuration
# ─────────────────────────────────────────────────────────────

DATA_DIR    = os.path.join(os.path.dirname(__file__), "data")
CSV_PATH    = os.path.join(os.path.dirname(__file__), "creditcard.csv")
RANDOM_STATE = 42
TEST_SIZE    = 0.2


# ─────────────────────────────────────────────────────────────
# 1. Chargement des données
# ─────────────────────────────────────────────────────────────

def load_data(path: str) -> pd.DataFrame:
    """Charge le CSV et affiche les statistiques de base."""
    print(f"[1/5] Chargement du dataset : {path}")
    df = pd.read_csv(path)
    
    print(f"      Shape        : {df.shape}")
    print(f"      Colonnes     : {list(df.columns)}")
    print(f"      Valeurs NaN  : {df.isnull().sum().sum()}")
    
    n_fraud   = (df["Class"] == 1).sum()
    n_legit   = (df["Class"] == 0).sum()
    pct_fraud = n_fraud / len(df) * 100
    print(f"\n      Distribution des classes :")
    print(f"        Légitimes (0) : {n_legit:,}  ({100 - pct_fraud:.2f}%)")
    print(f"        Fraudes   (1) : {n_fraud:,}  ({pct_fraud:.4f}%)")
    
    return df


# ─────────────────────────────────────────────────────────────
# 2. Préprocessing
# ─────────────────────────────────────────────────────────────

def preprocess(df: pd.DataFrame) -> tuple[pd.DataFrame, pd.Series, StandardScaler]:
    """
    Normalise les colonnes 'Time' et 'Amount' avec StandardScaler.
    Les features V1–V28 sont déjà normalisées (résultat de PCA sur le dataset original).
    
    Retourne : X (features), y (labels), scaler (pour réutilisation en inférence)
    """
    print("\n[2/5] Préprocessing — Normalisation de Time et Amount")
    
    df = df.copy()
    
    scaler = StandardScaler()
    df[["Time", "Amount"]] = scaler.fit_transform(df[["Time", "Amount"]])
    
    X = df.drop(columns=["Class"])
    y = df["Class"]
    
    print(f"      Features     : {X.shape[1]} colonnes")
    print(f"      Target       : 'Class' (0=Légitime, 1=Fraude)")
    
    return X, y, scaler


# ─────────────────────────────────────────────────────────────
# 3. Split Train / Test
# ─────────────────────────────────────────────────────────────

def split_data(
    X: pd.DataFrame,
    y: pd.Series,
    test_size: float = TEST_SIZE,
    random_state: int = RANDOM_STATE,
) -> tuple:
    """
    Split stratifié 80/20.
    La stratification garantit que la proportion de fraudes est préservée
    dans les deux splits, malgré le déséquilibre extrême du dataset.
    """
    print(f"\n[3/5] Split Train/Test (stratifié, test_size={test_size})")
    
    X_train, X_test, y_train, y_test = train_test_split(
        X, y,
        test_size=test_size,
        random_state=random_state,
        stratify=y,
    )
    
    print(f"      Train : {X_train.shape[0]:,} samples  "
          f"(fraudes : {y_train.sum():,}  — {y_train.mean()*100:.4f}%)")
    print(f"      Test  : {X_test.shape[0]:,} samples  "
          f"(fraudes : {y_test.sum():,}  — {y_test.mean()*100:.4f}%)")
    
    return X_train, X_test, y_train, y_test


# ─────────────────────────────────────────────────────────────
# 4. SMOTE — Rééquilibrage (Train uniquement !)
# ─────────────────────────────────────────────────────────────

def apply_smote(
    X_train: pd.DataFrame,
    y_train: pd.Series,
    random_state: int = RANDOM_STATE,
) -> tuple[pd.DataFrame, pd.Series]:
    """
    Applique SMOTE sur le jeu d'entraînement UNIQUEMENT.

    Raison critique : Appliquer SMOTE sur le test set provoquerait du
    data leakage et des métriques artificiellement gonflées. Le test
    set doit rester dans sa distribution naturelle (déséquilibrée)
    pour simuler des conditions réelles de détection.
    """
    print("\n[4/5] Application de SMOTE (sur X_train/y_train uniquement)")
    print(f"      Avant SMOTE  : {y_train.value_counts().to_dict()}")
    
    smote = SMOTE(random_state=random_state)
    X_train_res, y_train_res = smote.fit_resample(X_train, y_train)
    
    print(f"      Après SMOTE  : {pd.Series(y_train_res).value_counts().to_dict()}")
    print(f"      Ratio        : 1:1 (parfaitement équilibré)")
    
    return X_train_res, y_train_res


# ─────────────────────────────────────────────────────────────
# 5. Sauvegarde des Artifacts
# ─────────────────────────────────────────────────────────────

def save_artifacts(
    X_train: pd.DataFrame | np.ndarray,
    X_test: pd.DataFrame | np.ndarray,
    y_train: pd.Series | np.ndarray,
    y_test: pd.Series | np.ndarray,
    scaler: StandardScaler,
) -> None:
    """Sauvegarde tous les artifacts dans ./data/ au format .pkl."""
    print(f"\n[5/5] Sauvegarde dans : {DATA_DIR}/")
    os.makedirs(DATA_DIR, exist_ok=True)
    
    artifacts = {
        "X_train_balanced.pkl": X_train,
        "X_test.pkl":           X_test,
        "y_train_balanced.pkl": y_train,
        "y_test.pkl":           y_test,
        "scaler.pkl":           scaler,
    }
    
    for filename, obj in artifacts.items():
        path = os.path.join(DATA_DIR, filename)
        joblib.dump(obj, path)
        size_mb = os.path.getsize(path) / 1_000_000
        print(f"      [OK] {filename:<30} ({size_mb:.2f} MB)")
    
    print("\n[DONE] Pipeline de donnees terminee avec succes.")


# ─────────────────────────────────────────────────────────────
# Main
# ─────────────────────────────────────────────────────────────

def main():
    print("=" * 60)
    print(" PFA — Pipeline de Préparation des Données")
    print("=" * 60)
    
    # 1. Chargement
    df = load_data(CSV_PATH)
    
    # 2. Préprocessing
    X, y, scaler = preprocess(df)
    
    # 3. Split
    X_train, X_test, y_train, y_test = split_data(X, y)
    
    # 4. SMOTE (train uniquement)
    X_train_balanced, y_train_balanced = apply_smote(X_train, y_train)
    
    # 5. Sauvegarde
    save_artifacts(X_train_balanced, X_test, y_train_balanced, y_test, scaler)
    
    print("\n--- Recapitulatif final ---")
    print(f"   X_train_balanced : {X_train_balanced.shape}")
    print(f"   X_test           : {X_test.shape}")
    print(f"   y_train_balanced : {pd.Series(y_train_balanced).value_counts().to_dict()}")
    print(f"   y_test           : {y_test.value_counts().to_dict()}")


if __name__ == "__main__":
    main()