"""
Treino de modelo de classificação multiclasse: Real (0) vs IA (1) vs CGI (2).

Treina XGBoost + Random Forest, compara, salva o melhor.

Uso:
  python ml/train_model.py
  python ml/train_model.py --input dataset_treino.csv --output ml/model.joblib
"""

import argparse
import sys
from pathlib import Path

import numpy as np
import pandas as pd
from sklearn.model_selection import StratifiedKFold, cross_val_score
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import Pipeline
from sklearn.metrics import classification_report, confusion_matrix
import joblib

# Tenta importar XGBoost — fallback para GradientBoosting se não disponível
try:
    from xgboost import XGBClassifier
    HAS_XGB = True
except ImportError:
    HAS_XGB = False
    print("[AVISO] xgboost não instalado. Usando GradientBoosting do sklearn como fallback.")


# Colunas que NÃO são features (metadados)
META_COLUMNS = {"video_file", "frame_index", "label"}


def load_dataset(csv_path: str) -> tuple[pd.DataFrame, np.ndarray, np.ndarray, list[str]]:
    """Carrega CSV e separa features / labels."""
    df = pd.read_csv(csv_path)

    print(f"Dataset carregado: {len(df)} amostras")
    print(f"Distribuição de classes:")
    label_names = {0: "Real", 1: "IA", 2: "CGI"}
    for label, count in df["label"].value_counts().sort_index().items():
        print(f"  {label} ({label_names.get(label, '?')}): {count}")

    feature_cols = [c for c in df.columns if c not in META_COLUMNS]
    X = df[feature_cols].values.astype(np.float32)
    y = df["label"].values.astype(np.int32)

    # Checa NaN/Inf
    nan_mask = np.isnan(X) | np.isinf(X)
    if nan_mask.any():
        nan_count = nan_mask.sum()
        print(f"[AVISO] {nan_count} valores NaN/Inf encontrados. Substituindo por 0.")
        X = np.nan_to_num(X, nan=0.0, posinf=0.0, neginf=0.0)

    return df, X, y, feature_cols


def build_models() -> dict:
    """Retorna dict de modelos candidatos."""
    models = {}

    # Random Forest — robusto, bom baseline
    models["RandomForest"] = Pipeline([
        ("scaler", StandardScaler()),
        ("clf", RandomForestClassifier(
            n_estimators=300,
            max_depth=12,
            min_samples_leaf=5,
            class_weight="balanced",
            random_state=42,
            n_jobs=-1,
        )),
    ])

    # XGBoost ou GradientBoosting
    if HAS_XGB:
        models["XGBoost"] = Pipeline([
            ("scaler", StandardScaler()),
            ("clf", XGBClassifier(
                n_estimators=500,
                max_depth=8,
                learning_rate=0.05,
                subsample=0.8,
                colsample_bytree=0.8,
                objective="multi:softprob",
                eval_metric="mlogloss",
                random_state=42,
                n_jobs=-1,
            )),
        ])
    else:
        models["GradientBoosting"] = Pipeline([
            ("scaler", StandardScaler()),
            ("clf", GradientBoostingClassifier(
                n_estimators=300,
                max_depth=6,
                learning_rate=0.05,
                subsample=0.8,
                random_state=42,
            )),
        ])

    return models


def evaluate_models(models: dict, X: np.ndarray, y: np.ndarray) -> str:
    """Avalia modelos com cross-validation e retorna o nome do melhor."""
    cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
    results = {}

    for name, model in models.items():
        print(f"\n--- {name} ---")
        scores = cross_val_score(model, X, y, cv=cv, scoring="f1_macro", n_jobs=-1)
        mean_score = float(np.mean(scores))
        std_score = float(np.std(scores))
        results[name] = mean_score
        print(f"  F1-macro (5-fold): {mean_score:.4f} (+/- {std_score:.4f})")
        print(f"  Scores por fold: {[f'{s:.4f}' for s in scores]}")

    best = max(results, key=results.get)
    print(f"\nMelhor modelo: {best} (F1={results[best]:.4f})")
    return best


def train_final(model, X: np.ndarray, y: np.ndarray, feature_cols: list[str]):
    """Treina o modelo final em todo o dataset e exibe relatório."""
    model.fit(X, y)

    y_pred = model.predict(X)
    label_names = ["Real", "IA", "CGI"]

    print("\n" + "=" * 50)
    print("RELATÓRIO FINAL (treino completo)")
    print("=" * 50)
    print(classification_report(y, y_pred, target_names=label_names))

    print("Matriz de Confusão:")
    cm = confusion_matrix(y, y_pred)
    print(f"{'':>8} {'Real':>8} {'IA':>8} {'CGI':>8}")
    for i, row in enumerate(cm):
        print(f"{label_names[i]:>8} {row[0]:>8} {row[1]:>8} {row[2]:>8}")

    # Feature importance
    clf = model.named_steps["clf"]
    if hasattr(clf, "feature_importances_"):
        importances = clf.feature_importances_
        sorted_idx = np.argsort(importances)[::-1]
        print("\nTop 10 features mais importantes:")
        for rank, idx in enumerate(sorted_idx[:10]):
            print(f"  {rank+1}. {feature_cols[idx]}: {importances[idx]:.4f}")

    return model


def main():
    parser = argparse.ArgumentParser(description="Treina modelo de classificação multiclasse")
    parser.add_argument("--input", default="dataset_treino.csv", help="CSV de entrada")
    parser.add_argument("--output", default="ml/model.joblib", help="Arquivo do modelo salvo")
    args = parser.parse_args()

    if not Path(args.input).exists():
        print(f"Arquivo não encontrado: {args.input}")
        print("Execute primeiro: python ml/extract_features.py")
        sys.exit(1)

    df, X, y, feature_cols = load_dataset(args.input)

    if len(np.unique(y)) < 2:
        print("Dataset precisa de pelo menos 2 classes. Adicione mais vídeos.")
        sys.exit(1)

    print(f"\nFeatures: {len(feature_cols)}")
    print(f"Shape: {X.shape}")

    # Avalia modelos
    models = build_models()
    best_name = evaluate_models(models, X, y)

    # Treina modelo final
    best_model = models[best_name]
    train_final(best_model, X, y, feature_cols)

    # Salva
    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    save_data = {
        "model": best_model,
        "feature_columns": feature_cols,
        "label_map": {0: "Real", 1: "IA", 2: "CGI"},
        "model_name": best_name,
    }
    joblib.dump(save_data, output_path)
    print(f"\nModelo salvo: {output_path}")
    print(f"Para carregar: joblib.load('{output_path}')")


if __name__ == "__main__":
    main()