# src/model_compare.py

import os
from glob import glob
from typing import Dict

import numpy as np
import pandas as pd
from dotenv import load_dotenv
from sklearn.metrics import (
    accuracy_score,
    confusion_matrix,
    f1_score,
    precision_score,
    recall_score,
)

load_dotenv()


# ---------------------------
# Helpers: robust CSV reading
# ---------------------------
def _clean_columns(df: pd.DataFrame) -> pd.DataFrame:
    df.columns = [str(c).replace("\ufeff", "").strip() for c in df.columns]
    return df


def _clean_text_col(df: pd.DataFrame, col: str = "text") -> pd.DataFrame:
    if col in df.columns:
        df[col] = (
            df[col]
            .astype(str)
            .str.replace("\ufeff", "", regex=False)
            .str.replace("\u200b", "", regex=False)  # zero-width space, just in case
            .str.strip()
        )
    return df


def _safe_read_csv(path: str) -> pd.DataFrame:
    """
    Read CSV in a way that survives UTF-8 BOM and keeps columns consistent.
    """
    df = pd.read_csv(path, encoding="utf-8-sig")
    df = _clean_columns(df)
    df = _clean_text_col(df, "text")
    df = _clean_text_col(df, "id")
    return df


def _extract_tag(path: str, prefix: str) -> str:
    # e.g. chat_with_probs_xlmr.csv -> xlmr
    base = os.path.basename(path)
    tag = base.replace(prefix, "").replace(".csv", "")
    return tag.strip("_").lower()


# ---------------------------
# Label loading + joining
# ---------------------------
def _load_labels(
    processed_dir: str, label_file: str = "text_all_clean.csv"
) -> pd.DataFrame:
    label_path = os.path.join(processed_dir, label_file)
    df = _safe_read_csv(label_path)

    print(f"[SUCCESS] Loaded labels: {os.path.basename(label_path)} | shape={df.shape}")
    print("[SUCCESS] Label columns:", list(df.columns))

    if "label" not in df.columns:
        raise ValueError(
            f"Label file must have 'label' column. Found: {list(df.columns)}"
        )

    if "id" not in df.columns and "text" not in df.columns:
        raise ValueError(
            "Label file must have either 'id' or 'text' column to join predictions."
        )

    keep = [c for c in ["id", "text", "label", "lang"] if c in df.columns]
    out = df[keep].copy()

    # Ensure label numeric 0/1
    out["label"] = pd.to_numeric(out["label"], errors="coerce")
    if out["label"].isna().any():
        bad = out[out["label"].isna()].head(5)
        raise ValueError(
            "Some label values are not numeric (0/1). Example bad rows:\n" f"{bad}"
        )

    out["label"] = out["label"].astype(int)
    return out


def _join_pred_with_labels(
    labels_df: pd.DataFrame, pred_df: pd.DataFrame
) -> pd.DataFrame:
    """
    Join strategy:
    1) if both have 'id' -> join on id
    2) else if both have 'text' -> join on text

    ✅ IMPORTANT FIX:
    Prediction files sometimes already contain label columns.
    We REMOVE those before merging to avoid label_x/label_y.
    """
    pred_df = _clean_columns(pred_df)
    pred_df = _clean_text_col(pred_df, "text")
    pred_df = _clean_text_col(pred_df, "id")

    # ✅ drop any label-like columns in prediction df to avoid merge suffixes
    drop_cols = [
        c for c in pred_df.columns if c.lower() in {"label", "label_x", "label_y"}
    ]
    if drop_cols:
        pred_df = pred_df.drop(columns=drop_cols)

    if "id" in labels_df.columns and "id" in pred_df.columns:
        out = pred_df.merge(labels_df[["id", "label"]], on="id", how="inner")
        return out

    if "text" in labels_df.columns and "text" in pred_df.columns:
        out = pred_df.merge(labels_df[["text", "label"]], on="text", how="inner")
        return out

    raise ValueError(
        "Cannot join predictions with labels. Need shared 'id' or shared 'text' column."
    )


# ---------------------------
# Metrics
# ---------------------------
def _metrics(y_true: np.ndarray, y_pred: np.ndarray) -> Dict[str, float]:
    return {
        "accuracy": float(accuracy_score(y_true, y_pred)),
        "precision": float(precision_score(y_true, y_pred, zero_division=0)),
        "recall": float(recall_score(y_true, y_pred, zero_division=0)),
        "f1": float(f1_score(y_true, y_pred, zero_division=0)),
    }


def _format_cm(y_true: np.ndarray, y_pred: np.ndarray) -> str:
    cm = confusion_matrix(y_true, y_pred, labels=[0, 1])
    tn, fp, fn, tp = cm.ravel()
    return f"tn={tn}, fp={fp}, fn={fn}, tp={tp}"


# ---------------------------
# Compare logic
# ---------------------------
def compare_models(
    processed_dir: str = "data/processed",
    reports_dir: str = "outputs/reports",
    chat_thr_default: float = 0.50,
) -> Dict[str, str]:
    """
    Expects:
      - data/processed/chat_with_probs_<tag>.csv with: chat_prob and (id or text)
      - data/processed/fusion_final_output_<tag>.csv with: final_risk_score and (id or text)
      - data/processed/text_all_clean.csv with: text,label,lang (or id,label)
    """
    os.makedirs(reports_dir, exist_ok=True)

    labels_df = _load_labels(processed_dir)

    chat_files = sorted(glob(os.path.join(processed_dir, "chat_with_probs_*.csv")))
    fusion_files = sorted(
        glob(os.path.join(processed_dir, "fusion_final_output_*.csv"))
    )

    if not chat_files:
        raise FileNotFoundError(f"No chat_with_probs_*.csv found in: {processed_dir}")
    if not fusion_files:
        raise FileNotFoundError(
            f"No fusion_final_output_*.csv found in: {processed_dir}"
        )

    # -------- CHAT COMPARISON --------
    chat_rows = []
    for path in chat_files:
        tag = _extract_tag(path, "chat_with_probs_")
        df_pred = _safe_read_csv(path)

        if "chat_prob" not in df_pred.columns:
            raise ValueError(
                f"{os.path.basename(path)} missing 'chat_prob'. Found: {list(df_pred.columns)}"
            )

        joined = _join_pred_with_labels(labels_df, df_pred)

        if len(joined) == 0:
            raise ValueError(
                f"Join produced 0 rows for chat file: {os.path.basename(path)}\n"
                f"Labels columns: {list(labels_df.columns)}\n"
                f"Pred columns: {list(df_pred.columns)}\n"
                "Fix: ensure both sides share 'id' or 'text' with same values."
            )

        if "label" not in joined.columns:
            raise KeyError(
                f"After join, 'label' is missing for chat file: {os.path.basename(path)}\n"
                f"Joined columns: {list(joined.columns)}"
            )

        y_true = joined["label"].astype(int).values

        # Optional per-model threshold.json
        thr = chat_thr_default
        thr_path = os.path.join(
            "outputs", "models", f"chat_brain_{tag}", "threshold.json"
        )
        if os.path.exists(thr_path):
            try:
                import json

                with open(thr_path, "r", encoding="utf-8") as f:
                    thr_obj = json.load(f)
                if "thr" in thr_obj:
                    thr = float(thr_obj["thr"])
            except Exception:
                pass

        y_pred = (joined["chat_prob"].astype(float).values >= thr).astype(int)
        m = _metrics(y_true, y_pred)

        chat_rows.append(
            {
                "model_tag": tag,
                "n": int(len(joined)),
                "threshold_used": float(thr),
                **m,
                "confusion": _format_cm(y_true, y_pred),
            }
        )

    chat_df = pd.DataFrame(chat_rows).sort_values(["recall", "f1"], ascending=False)
    chat_out = os.path.join(reports_dir, "model_comparison_chat.csv")
    chat_df.to_csv(chat_out, index=False, encoding="utf-8")

    # -------- FUSION COMPARISON --------
    fusion_rows = []
    for path in fusion_files:
        tag = _extract_tag(path, "fusion_final_output_")
        df_pred = _safe_read_csv(path)

        if "final_risk_score" not in df_pred.columns:
            raise ValueError(
                f"{os.path.basename(path)} missing 'final_risk_score'. Found: {list(df_pred.columns)}"
            )

        joined = _join_pred_with_labels(labels_df, df_pred)

        if len(joined) == 0:
            raise ValueError(
                f"Join produced 0 rows for fusion file: {os.path.basename(path)}\n"
                f"Labels columns: {list(labels_df.columns)}\n"
                f"Pred columns: {list(df_pred.columns)}\n"
                "Fix: ensure both sides share 'id' or 'text' with same values."
            )

        if "label" not in joined.columns:
            raise KeyError(
                f"After join, 'label' is missing for fusion file: {os.path.basename(path)}\n"
                f"Joined columns: {list(joined.columns)}"
            )

        y_true = joined["label"].astype(int).values

        # use same threshold as chat model if available else default
        thr = chat_thr_default
        thr_path = os.path.join(
            "outputs", "models", f"chat_brain_{tag}", "threshold.json"
        )
        if os.path.exists(thr_path):
            try:
                import json

                with open(thr_path, "r", encoding="utf-8") as f:
                    thr_obj = json.load(f)
                if "thr" in thr_obj:
                    thr = float(thr_obj["thr"])
            except Exception:
                pass

        y_pred = (joined["final_risk_score"].astype(float).values >= thr).astype(int)
        m = _metrics(y_true, y_pred)

        fusion_rows.append(
            {
                "model_tag": tag,
                "n": int(len(joined)),
                "threshold_used": float(thr),
                **m,
                "confusion": _format_cm(y_true, y_pred),
            }
        )

    fusion_df = pd.DataFrame(fusion_rows).sort_values(["recall", "f1"], ascending=False)
    fusion_out = os.path.join(reports_dir, "model_comparison_fusion.csv")
    fusion_df.to_csv(fusion_out, index=False, encoding="utf-8")

    # -------- SUMMARY --------
    summary = chat_df.merge(
        fusion_df,
        on="model_tag",
        how="outer",
        suffixes=("_chat", "_fusion"),
    )

    summary = summary.sort_values(
        ["recall_fusion", "f1_fusion", "recall_chat", "f1_chat"],
        ascending=False,
        na_position="last",
    )

    summary_out = os.path.join(reports_dir, "model_comparison_summary.csv")
    summary.to_csv(summary_out, index=False, encoding="utf-8")

    print("\n================= MODEL COMPARISON (CHAT) =================")
    print(
        chat_df[
            [
                "model_tag",
                "threshold_used",
                "recall",
                "f1",
                "precision",
                "accuracy",
                "confusion",
            ]
        ].to_string(index=False)
    )

    print("\n================= MODEL COMPARISON (FUSION) =================")
    print(
        fusion_df[
            [
                "model_tag",
                "threshold_used",
                "recall",
                "f1",
                "precision",
                "accuracy",
                "confusion",
            ]
        ].to_string(index=False)
    )

    print("\n[SUCCESS] Saved reports:")
    print("-", chat_out)
    print("-", fusion_out)
    print("-", summary_out)

    return {
        "chat_report": chat_out,
        "fusion_report": fusion_out,
        "summary_report": summary_out,
    }


def main():
    processed_dir = os.getenv("PROCESSED_DIR", "data/processed")
    reports_dir = os.getenv("REPORTS_DIR", "outputs/reports")
    compare_models(processed_dir=processed_dir, reports_dir=reports_dir)


if __name__ == "__main__":
    main()