"""
Benchmark — Comparing Your RF Model Against Classic ML Baselines
================================================================
Trains and evaluates multiple fake-news classifiers on the SAME
locked test split so the comparison is completely fair.

Models compared
---------------
  1. Naive Bayes        (TF-IDF only)            — simple baseline
  2. Logistic Regression (TF-IDF only)           — strong linear baseline
  3. Linear SVM          (TF-IDF only)           — often best for text
  4. Random Forest       (TF-IDF only)           — RF without embeddings
  5. ★ YOUR MODEL ★   (TF-IDF + MiniLM + Stylo) — YOUR hybrid RF

Metrics reported (per model, per class + weighted avg)
-------------------------------------------------------
  • Accuracy
  • Precision  (Fake class)
  • Recall     (Fake class)
  • F1 Score   (weighted)
  • AUC-ROC

Output
------
  evaluation_results/benchmark_table.txt   — plaintext comparison table
  evaluation_results/benchmark_chart.png   — bar chart (Accuracy + F1)
  evaluation_results/benchmark_roc.png     — ROC curves for all models

Usage
-----
  python backend/benchmark.py
  python backend/benchmark.py --mode tagalog    # Tagalog dataset only
  python backend/benchmark.py --mode cebuano   # Cebuano dataset only
  python backend/benchmark.py --mode mixed      # All languages (default)
"""

import sys
import os
import re
import time
import argparse
import warnings

warnings.filterwarnings("ignore")

PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, PROJECT_ROOT)

import numpy as np
import pandas as pd

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches

from scipy.sparse import hstack, csr_matrix

from sklearn.model_selection import train_test_split
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import Pipeline

# Classifiers
from sklearn.naive_bayes import MultinomialNB
from sklearn.linear_model import LogisticRegression
from sklearn.svm import LinearSVC
from sklearn.ensemble import RandomForestClassifier
from sklearn.calibration import CalibratedClassifierCV

# Metrics
from sklearn.metrics import (
    accuracy_score,
    classification_report,
    confusion_matrix,
    roc_auc_score,
    roc_curve,
    f1_score,
    precision_score,
    recall_score,
)

# Re-use helpers from train.py (keep feature extraction identical)
from backend.train import (
    load_fake_news_dataset,
    preprocess,
    clean_text,
    extract_stylometric_features,
    get_minilm_model,
    STYLOMETRIC_FEATURE_NAMES,
)

# ── Paths ──────────────────────────────────────────────────────────────────
OUTPUT_DIR = os.path.join(PROJECT_ROOT, "evaluation_results")


# ───────────────────────────────────────────────────────────────────────────
# Feature builders
# ───────────────────────────────────────────────────────────────────────────


def build_tfidf_features(X_train, X_test):
    """Plain TF-IDF (used by all baseline models)."""
    tfidf = TfidfVectorizer(
        max_features=15_000,
        ngram_range=(1, 3),
        min_df=2,
        max_df=0.95,
        sublinear_tf=True,
    )
    X_tr = tfidf.fit_transform(X_train)
    X_te = tfidf.transform(X_test)
    return X_tr, X_te, tfidf


def build_hybrid_features(X_train, X_test):
    """TF-IDF + MiniLM embeddings + stylometric (your full pipeline)."""
    print("  [hybrid] Fitting TF-IDF …")
    tfidf = TfidfVectorizer(
        max_features=15_000,
        ngram_range=(1, 3),
        min_df=2,
        max_df=0.95,
        sublinear_tf=True,
    )
    X_tr_tfidf = tfidf.fit_transform(X_train)
    X_te_tfidf = tfidf.transform(X_test)

    print("  [hybrid] Encoding with MiniLM …")
    minilm = get_minilm_model()
    emb_train = minilm.encode(X_train, show_progress_bar=True, batch_size=64)
    emb_test  = minilm.encode(X_test,  show_progress_bar=True, batch_size=64)

    print(f"  [hybrid] Extracting {len(STYLOMETRIC_FEATURE_NAMES)} stylometric features …")
    stylo_train = np.array([extract_stylometric_features(t) for t in X_train])
    stylo_test  = np.array([extract_stylometric_features(t) for t in X_test])

    scaler = StandardScaler()
    stylo_train_sc = scaler.fit_transform(stylo_train)
    stylo_test_sc  = scaler.transform(stylo_test)

    X_tr = hstack([X_tr_tfidf, csr_matrix(emb_train), csr_matrix(stylo_train_sc)])
    X_te = hstack([X_te_tfidf, csr_matrix(emb_test),  csr_matrix(stylo_test_sc)])

    n_total = X_tr.shape[1]
    print(
        f"  [hybrid] Feature dimensions: {n_total} "
        f"(TF-IDF: {X_tr_tfidf.shape[1]} + MiniLM: 384 + Stylo: {len(STYLOMETRIC_FEATURE_NAMES)})"
    )
    return X_tr, X_te


# ───────────────────────────────────────────────────────────────────────────
# Compute metrics for one model
# ───────────────────────────────────────────────────────────────────────────


def evaluate(name, model, X_test, y_test, proba=None):
    """Return a metrics dict for one fitted model."""
    y_pred = model.predict(X_test)

    acc         = accuracy_score(y_test, y_pred)
    prec_fake   = precision_score(y_test, y_pred, pos_label=1, zero_division=0)
    rec_fake    = recall_score(y_test, y_pred, pos_label=1, zero_division=0)
    f1_weighted = f1_score(y_test, y_pred, average="weighted", zero_division=0)
    f1_fake     = f1_score(y_test, y_pred, pos_label=1, zero_division=0)

    # AUC-ROC (needs probability scores)
    if proba is not None:
        try:
            auc = roc_auc_score(y_test, proba[:, 1])
        except Exception:
            auc = float("nan")
    else:
        auc = float("nan")

    report = classification_report(
        y_test, y_pred,
        target_names=["Real", "Fake"],
        zero_division=0,
    )

    return {
        "name":         name,
        "accuracy":     acc,
        "precision":    prec_fake,
        "recall":       rec_fake,
        "f1_weighted":  f1_weighted,
        "f1_fake":      f1_fake,
        "auc":          auc,
        "report":       report,
        "y_pred":       y_pred,
        "proba":        proba,
    }


# ───────────────────────────────────────────────────────────────────────────
# Bar chart
# ───────────────────────────────────────────────────────────────────────────


def plot_bar_chart(results, output_path):
    """Side-by-side bar chart: Accuracy vs F1 (weighted) vs AUC-ROC."""
    names     = [r["name"] for r in results]
    accs      = [r["accuracy"]    for r in results]
    f1s       = [r["f1_weighted"] for r in results]
    aucs      = [r["auc"]         for r in results]

    x     = np.arange(len(names))
    width = 0.26

    fig, ax = plt.subplots(figsize=(max(10, len(names) * 2.2), 6))

    # Color highlight for YOUR model (last entry)
    bar_colors_acc = ["#2196F3"] * (len(names) - 1) + ["#E91E63"]
    bar_colors_f1  = ["#4CAF50"] * (len(names) - 1) + ["#FF5722"]
    bar_colors_auc = ["#9C27B0"] * (len(names) - 1) + ["#FF9800"]

    b1 = ax.bar(x - width, accs,  width, color=bar_colors_acc, edgecolor="black", lw=0.5, label="Accuracy")
    b2 = ax.bar(x,          f1s,   width, color=bar_colors_f1,  edgecolor="black", lw=0.5, label="F1 Weighted")
    b3 = ax.bar(x + width,  aucs,  width, color=bar_colors_auc, edgecolor="black", lw=0.5, label="AUC-ROC")

    # Value labels
    for bars in (b1, b2, b3):
        for bar in bars:
            h = bar.get_height()
            if not np.isnan(h):
                ax.text(
                    bar.get_x() + bar.get_width() / 2,
                    h + 0.005,
                    f"{h:.3f}",
                    ha="center", va="bottom", fontsize=7.5, fontweight="bold",
                )

    ax.set_xticks(x)
    ax.set_xticklabels(names, rotation=12, ha="right", fontsize=10)
    ax.set_ylim(0, 1.12)
    ax.set_ylabel("Score", fontsize=12)
    ax.set_title(
        "Benchmark: Your RF Model vs. Classic ML Baselines\n"
        "(Highlighted in pink/orange = Your Model)",
        fontsize=13, fontweight="bold",
    )
    ax.axhline(y=0.80, color="gray", linestyle="--", alpha=0.4, linewidth=1)
    ax.text(len(names) - 0.5, 0.805, "80% threshold", color="gray", fontsize=8)

    patch_yours = mpatches.Patch(color="#E91E63", label="★ Your Hybrid RF (Accuracy)")
    ax.legend(handles=[*ax.get_legend_handles_labels()[0][:3], patch_yours], fontsize=9)

    plt.tight_layout()
    fig.savefig(output_path, dpi=150, bbox_inches="tight")
    plt.close(fig)
    print(f"  Saved: {output_path}")


# ───────────────────────────────────────────────────────────────────────────
# ROC curve chart
# ───────────────────────────────────────────────────────────────────────────


def plot_roc_curves(results, y_test, output_path):
    """Overlay ROC curves for all models that have probability scores."""
    fig, ax = plt.subplots(figsize=(8, 6))

    COLORS = [
        "#2196F3", "#4CAF50", "#FF9800", "#9C27B0",
        "#E91E63", "#00BCD4", "#F44336", "#8BC34A",
    ]

    for i, r in enumerate(results):
        if r["proba"] is None or np.isnan(r["auc"]):
            continue
        fpr, tpr, _ = roc_curve(y_test, r["proba"][:, 1], pos_label=1)
        lw = 2.5 if "★" in r["name"] else 1.5
        dash = "-" if "★" in r["name"] else "--"
        ax.plot(
            fpr, tpr,
            color=COLORS[i % len(COLORS)],
            lw=lw, linestyle=dash,
            label=f"{r['name']}  (AUC={r['auc']:.3f})",
        )

    ax.plot([0, 1], [0, 1], "k:", lw=1, label="Random (AUC=0.500)")
    ax.set_xlim([0.0, 1.0])
    ax.set_ylim([0.0, 1.05])
    ax.set_xlabel("False Positive Rate", fontsize=12)
    ax.set_ylabel("True Positive Rate", fontsize=12)
    ax.set_title("ROC Curves — Fake-News Detection Benchmark", fontsize=13, fontweight="bold")
    ax.legend(loc="lower right", fontsize=9)
    ax.grid(alpha=0.3)

    plt.tight_layout()
    fig.savefig(output_path, dpi=150, bbox_inches="tight")
    plt.close(fig)
    print(f"  Saved: {output_path}")


# ───────────────────────────────────────────────────────────────────────────
# Plaintext summary table
# ───────────────────────────────────────────────────────────────────────────


def save_table(results, output_path, mode, n_train, n_test):
    """Write a neatly formatted comparison table to disk and stdout."""
    lines = []
    sep   = "=" * 90

    lines.append(sep)
    lines.append("  BENCHMARK RESULTS — Fake-News Detection (Filipino / Cebuano)")
    lines.append(f"  Mode: {mode.upper()}   |   Train: {n_train:,} samples   |   Test: {n_test:,} samples")
    lines.append(sep)
    lines.append("")

    header = (
        f"  {'Model':<35} {'Accuracy':>9} {'Prec(Fk)':>10} {'Rec(Fk)':>9} "
        f"{'F1(Wtd)':>9} {'F1(Fk)':>8} {'AUC-ROC':>9}"
    )
    lines.append(header)
    lines.append("  " + "-" * 88)

    for r in results:
        auc_str = f"{r['auc']:.4f}" if not np.isnan(r["auc"]) else "  N/A  "
        marker  = "  ★" if "★" in r["name"] else "   "
        row = (
            f"{marker} {r['name']:<32} "
            f"{r['accuracy']:>9.4f} "
            f"{r['precision']:>10.4f} "
            f"{r['recall']:>9.4f} "
            f"{r['f1_weighted']:>9.4f} "
            f"{r['f1_fake']:>8.4f} "
            f"{auc_str:>9}"
        )
        lines.append(row)

    lines.append("")
    lines.append(sep)
    lines.append("  DETAILED CLASSIFICATION REPORTS")
    lines.append(sep)

    for r in results:
        lines.append("")
        lines.append(f"  ── {r['name']} ──────────────────────────────────────────────")
        for ln in r["report"].splitlines():
            lines.append(f"  {ln}")

    text = "\n".join(lines)
    print(text)

    with open(output_path, "w", encoding="utf-8") as f:
        f.write(text)
    print(f"\n  Saved: {output_path}")


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


def main():
    parser = argparse.ArgumentParser(description="Benchmark fake-news classifiers.")
    parser.add_argument(
        "--mode",
        choices=["mixed", "tagalog", "cebuano"],
        default="mixed",
        help="Which language subset to benchmark on (default: mixed).",
    )
    parser.add_argument(
        "--test-size",
        type=float,
        default=0.20,
        help="Fraction of data to hold out as the locked test set (default: 0.20).",
    )
    parser.add_argument(
        "--skip-minilm",
        action="store_true",
        default=False,
        help="Skip your hybrid RF model (useful if MiniLM is too slow for a quick check).",
    )
    args = parser.parse_args()

    os.makedirs(OUTPUT_DIR, exist_ok=True)

    print("=" * 60)
    print("  FAKE-NEWS BENCHMARK")
    print(f"  Mode: {args.mode.upper()} | Test size: {args.test_size:.0%}")
    print("=" * 60)

    # ── 1. Load & preprocess dataset ─────────────────────────────────────
    tagalog_only = args.mode == "tagalog"
    cebuano_only = args.mode == "cebuano"

    df = load_fake_news_dataset(tagalog_only=tagalog_only, cebuano_only=cebuano_only)
    X_all, y_all = preprocess(df, undersample=False, oversample=True)

    # ── 2. Locked test split (same seed → reproducible) ──────────────────
    print(f"\nCreating locked test split ({1 - args.test_size:.0%} train / {args.test_size:.0%} test) …")
    X_train, X_test, y_train, y_test = train_test_split(
        X_all, y_all,
        test_size=args.test_size,
        random_state=42,
        stratify=y_all,
    )
    print(f"  Train: {len(X_train):,}  |  Test: {len(X_test):,}")
    print(f"  Test distribution — Real: {y_test.count(0):,}, Fake: {y_test.count(1):,}")

    y_train_arr = np.array(y_train)
    y_test_arr  = np.array(y_test)

    # ── 3. Build TF-IDF features (shared by baseline models) ─────────────
    print("\nBuilding TF-IDF features for baseline models …")
    X_tr_tfidf, X_te_tfidf, tfidf = build_tfidf_features(X_train, X_test)
    print(f"  TF-IDF shape: {X_tr_tfidf.shape}")

    # ── 4. Train and evaluate each model ─────────────────────────────────
    results = []

    # ── 4a. Naive Bayes (TF-IDF, no negative values — shift by min) ──────
    print("\n[1/5] Naive Bayes (TF-IDF) …")
    t0  = time.time()
    nb  = MultinomialNB(alpha=1.0)
    # MultinomialNB needs non-negative input; TF-IDF with sublinear_tf is ok
    nb.fit(X_tr_tfidf, y_train_arr)
    nb_proba = nb.predict_proba(X_te_tfidf)
    elapsed  = time.time() - t0
    res_nb   = evaluate("Naive Bayes (TF-IDF)", nb, X_te_tfidf, y_test_arr, nb_proba)
    res_nb["train_time"] = elapsed
    results.append(res_nb)
    print(f"  Accuracy: {res_nb['accuracy']:.4f}  |  F1 Weighted: {res_nb['f1_weighted']:.4f}  |  Time: {elapsed:.1f}s")

    # ── 4b. Logistic Regression (TF-IDF) ─────────────────────────────────
    print("\n[2/5] Logistic Regression (TF-IDF) …")
    t0  = time.time()
    lr  = LogisticRegression(
        max_iter=1000,
        class_weight="balanced",
        solver="saga",
        C=1.0,
        random_state=42,
        n_jobs=-1,
    )
    lr.fit(X_tr_tfidf, y_train_arr)
    lr_proba = lr.predict_proba(X_te_tfidf)
    elapsed  = time.time() - t0
    res_lr   = evaluate("Logistic Regression (TF-IDF)", lr, X_te_tfidf, y_test_arr, lr_proba)
    res_lr["train_time"] = elapsed
    results.append(res_lr)
    print(f"  Accuracy: {res_lr['accuracy']:.4f}  |  F1 Weighted: {res_lr['f1_weighted']:.4f}  |  Time: {elapsed:.1f}s")

    # ── 4c. Linear SVM (TF-IDF) ──────────────────────────────────────────
    print("\n[3/5] Linear SVM (TF-IDF) …")
    t0  = time.time()
    svm = CalibratedClassifierCV(
        LinearSVC(class_weight="balanced", max_iter=2000, random_state=42),
        cv=3,
    )
    svm.fit(X_tr_tfidf, y_train_arr)
    svm_proba = svm.predict_proba(X_te_tfidf)
    elapsed   = time.time() - t0
    res_svm   = evaluate("Linear SVM (TF-IDF)", svm, X_te_tfidf, y_test_arr, svm_proba)
    res_svm["train_time"] = elapsed
    results.append(res_svm)
    print(f"  Accuracy: {res_svm['accuracy']:.4f}  |  F1 Weighted: {res_svm['f1_weighted']:.4f}  |  Time: {elapsed:.1f}s")

    # ── 4d. Random Forest (TF-IDF only — no embeddings) ──────────────────
    print("\n[4/5] Random Forest (TF-IDF only — no embeddings) …")
    t0    = time.time()
    rf_tf = RandomForestClassifier(
        n_estimators=300,
        max_depth=20,
        min_samples_split=5,
        min_samples_leaf=3,
        class_weight="balanced",
        n_jobs=-1,
        random_state=42,
    )
    rf_tf.fit(X_tr_tfidf, y_train_arr)
    rf_tf_proba = rf_tf.predict_proba(X_te_tfidf)
    elapsed     = time.time() - t0
    res_rf_tf   = evaluate("Random Forest (TF-IDF only)", rf_tf, X_te_tfidf, y_test_arr, rf_tf_proba)
    res_rf_tf["train_time"] = elapsed
    results.append(res_rf_tf)
    print(f"  Accuracy: {res_rf_tf['accuracy']:.4f}  |  F1 Weighted: {res_rf_tf['f1_weighted']:.4f}  |  Time: {elapsed:.1f}s")

    # ── 4e. YOUR Hybrid RF (TF-IDF + MiniLM + Stylometric) ───────────────
    if not args.skip_minilm:
        print("\n[5/5] ★ YOUR Hybrid RF (TF-IDF + MiniLM + Stylometric) …")
        X_tr_hy, X_te_hy = build_hybrid_features(X_train, X_test)

        t0    = time.time()
        max_depth = 8 if cebuano_only else 20
        rf_hy = RandomForestClassifier(
            n_estimators=500,
            max_depth=max_depth,
            min_samples_split=5,
            min_samples_leaf=3,
            class_weight="balanced",
            n_jobs=-1,
            random_state=42,
        )
        rf_hy.fit(X_tr_hy, y_train_arr)
        rf_hy_proba = rf_hy.predict_proba(X_te_hy)
        elapsed     = time.time() - t0
        res_rf_hy   = evaluate("★ Hybrid RF (TF-IDF + MiniLM + Stylo)", rf_hy, X_te_hy, y_test_arr, rf_hy_proba)
        res_rf_hy["train_time"] = elapsed
        results.append(res_rf_hy)
        print(f"  Accuracy: {res_rf_hy['accuracy']:.4f}  |  F1 Weighted: {res_rf_hy['f1_weighted']:.4f}  |  Time: {elapsed:.1f}s")
    else:
        print("\n[5/5] Skipping Hybrid RF (--skip-minilm flag set).")

    # ── 5. Output table ───────────────────────────────────────────────────
    print("\n" + "=" * 60)
    print("  BENCHMARK SUMMARY")
    print("=" * 60)
    table_path = os.path.join(OUTPUT_DIR, f"benchmark_table_{args.mode}.txt")
    save_table(results, table_path, args.mode, len(X_train), len(X_test))

    # ── 6. Bar chart ──────────────────────────────────────────────────────
    chart_path = os.path.join(OUTPUT_DIR, f"benchmark_chart_{args.mode}.png")
    plot_bar_chart(results, chart_path)

    # ── 7. ROC curves ─────────────────────────────────────────────────────
    roc_path = os.path.join(OUTPUT_DIR, f"benchmark_roc_{args.mode}.png")
    plot_roc_curves(results, y_test_arr, roc_path)

    # ── 8. Train time summary ─────────────────────────────────────────────
    print("\n  Training times:")
    for r in results:
        t = r.get("train_time", 0)
        print(f"    {r['name']:<45} {t:>6.1f}s")

    print("\n" + "=" * 60)
    print("  BENCHMARK COMPLETE")
    print(f"  Results saved to: {OUTPUT_DIR}")
    print("=" * 60)


if __name__ == "__main__":
    main()