eval_allqs.py

# %%
from __future__ import annotations

import json
import re
from dataclasses import dataclass
from pathlib import Path

import numpy as np
import pandas as pd
import seaborn as sns
import pingouin as pg
import matplotlib.pyplot as plt

from scipy.stats import pearsonr, spearmanr, linregress, t
from sklearn.metrics import (
    accuracy_score,
    balanced_accuracy_score,
    recall_score,
    precision_score,
    confusion_matrix,
    f1_score,
    matthews_corrcoef,
    roc_auc_score,
)

# %%
# -----------------------
# CONFIG
# -----------------------
@dataclass(frozen=True)
class Config:
    base: Path = Path("")
    raw_root: Path = Path("outputs_allqs") / "raw_outputs"
    out_root: Path = Path("analysis_out_allqs")
    gt_csv_name: str = "all_scores.csv"
    status_mode: str = "complete"   # "complete" | "complete_partial" | "all"

    @property
    def gt_csv(self) -> Path:
        return self.base / "allqs" / self.gt_csv_name

    @property
    def plots_dir(self) -> Path:
        return self.out_root / "plots"


CFG = Config()
CFG.out_root.mkdir(parents=True, exist_ok=True)
CFG.plots_dir.mkdir(parents=True, exist_ok=True)

print("GT_CSV:", CFG.gt_csv)
print("RAW_ROOT:", CFG.raw_root.resolve())
print("STATUS_MODE:", CFG.status_mode)

# %%
# -----------------------
# CONSTANTS / REGEX
# -----------------------
Q_RE = re.compile(r"^Q([1-9]|1[0-5])$", flags=re.IGNORECASE)

FENCED_JSON_RE = re.compile(
    r"```(?:json)?\s*(\{.*?\})\s*```",
    flags=re.IGNORECASE | re.DOTALL,
)
SEV_RE = re.compile(r'(?i)\bseverity\b\s*["\']?\s*[:=]\s*([0-6])')
FREQ_RE = re.compile(r'(?i)\bfrequency\b\s*["\']?\s*[:=]\s*([0-6])')


# %%
# -----------------------
# BASIC HELPERS
# -----------------------
def to_num(x):
    return pd.to_numeric(x, errors="coerce")


def normalize_model_name(x) -> str:
    if isinstance(x, (tuple, list)):
        return str(x[0]).strip() if len(x) else ""
    if pd.isna(x):
        return ""
    return str(x).strip()


def session_to_int(value):
    if value is None or (isinstance(value, float) and np.isnan(value)):
        return np.nan
    match = re.search(r"(\d+)", str(value))
    return float(match.group(1)) if match else np.nan


def normalize_status(value) -> str:
    if value is None or (isinstance(value, float) and np.isnan(value)):
        return ""
    s = str(value).strip()
    if not s:
        return ""
    s_lower = s.lower()
    if s_lower.startswith("comp"):
        return "Complete"
    if s_lower.startswith("part"):
        return "Partial"
    if s_lower == "no":
        return "No"
    return s


def pick_status_col(df: pd.DataFrame) -> str | None:
    for col in ("status_from_csv", "status", "Status"):
        if col in df.columns:
            return col
    return None


def filter_by_status(df: pd.DataFrame, mode: str) -> tuple[pd.DataFrame, str | None]:
    mode = str(mode).strip().lower()
    status_col = pick_status_col(df)

    if status_col is None or mode == "all":
        return df.copy(), status_col

    out = df.copy()
    out[status_col] = out[status_col].astype(str).str.strip()

    if mode == "complete":
        out = out[out[status_col].eq("Complete")]
    elif mode in {"complete_partial", "complete+partial", "cp"}:
        out = out[out[status_col].isin(["Complete", "Partial"])]
    else:
        raise ValueError("STATUS_MODE must be one of: complete / complete_partial / all")

    return out, status_col


def ensure_numeric(df: pd.DataFrame, cols: list[str]) -> pd.DataFrame:
    out = df.copy()
    for col in cols:
        if col in out.columns:
            out[col] = to_num(out[col])
    return out


# %%
# -----------------------
# STATS HELPERS
# -----------------------
def safe_pearson(x, y) -> tuple[float, float]:
    d = pd.DataFrame({"x": to_num(x), "y": to_num(y)}).dropna()
    if len(d) < 3:
        return np.nan, np.nan
    r, p = pearsonr(d["x"].to_numpy(), d["y"].to_numpy())
    return float(r), float(p)


def icc_a1_two_col(true_s, pred_s) -> float:
    """
    ICC(A,1) = ICC2 in pingouin:
    two-way random, absolute agreement, single measure.
    """
    d = pd.DataFrame({"true": to_num(true_s), "pred": to_num(pred_s)}).dropna()
    if len(d) < 3:
        return np.nan

    long = d.melt(var_name="rater", value_name="score")
    long["target"] = np.tile(np.arange(len(d)), 2)

    icc_tbl = pg.intraclass_corr(
        data=long,
        targets="target",
        raters="rater",
        ratings="score",
    )
    row = icc_tbl.loc[icc_tbl["Type"].str.upper().eq("ICC2")]
    return float(row["ICC"].iloc[0]) if not row.empty else np.nan


def auc_from_score(y_true, score) -> float:
    y_true = pd.Series(y_true)
    score = pd.Series(score)

    mask = y_true.notna() & score.notna()
    y = y_true[mask].astype(int)
    s = score[mask].astype(float)

    if y.nunique() < 2:
        return np.nan

    try:
        return float(roc_auc_score(y, s))
    except ValueError:
        return np.nan


def chr_binary(severity, frequency, aps_min_sev=3, aps_min_freq=3) -> np.ndarray:
    """
    CHR-P gate:
    - positive if severity == 6
    - OR severity >= 3 and frequency >= 3
    """
    sev = pd.Series(to_num(severity))
    freq = pd.Series(to_num(frequency))
    out = sev.eq(6) | (sev.ge(aps_min_sev) & freq.ge(aps_min_freq))
    return out.fillna(False).astype(int).to_numpy()


def metrics_from_confusion(y_true, y_pred) -> dict:
    mask = pd.notna(y_true) & pd.notna(y_pred)
    yt = pd.Series(y_true)[mask].astype(int)
    yp = pd.Series(y_pred)[mask].astype(int)

    if len(yt) == 0 or yt.nunique() < 2:
        return {
            "accuracy": np.nan,
            "accuracy_bal": np.nan,
            "sensitivity": np.nan,
            "specificity": np.nan,
            "ppv": np.nan,
            "npv": np.nan,
            "f1": np.nan,
            "mcc": np.nan,
            "tn": np.nan,
            "fp": np.nan,
            "fn": np.nan,
            "tp": np.nan,
        }

    acc = accuracy_score(yt, yp)
    acc_bal = balanced_accuracy_score(yt, yp)
    sens = recall_score(yt, yp, pos_label=1, zero_division=0)
    spec = recall_score(yt, yp, pos_label=0, zero_division=0)
    ppv = precision_score(yt, yp, pos_label=1, zero_division=0)

    tn, fp, fn, tp = confusion_matrix(yt, yp, labels=[0, 1]).ravel()
    npv = tn / (tn + fn) if (tn + fn) > 0 else np.nan
    f1 = f1_score(yt, yp, pos_label=1, zero_division=0)
    mcc = matthews_corrcoef(yt, yp)

    return {
        "accuracy": float(acc),
        "accuracy_bal": float(acc_bal),
        "sensitivity": float(sens),
        "specificity": float(spec),
        "ppv": float(ppv),
        "npv": float(npv) if pd.notna(npv) else np.nan,
        "f1": float(f1),
        "mcc": float(mcc),
        "tn": int(tn),
        "fp": int(fp),
        "fn": int(fn),
        "tp": int(tp),
    }


# %%
# -----------------------
# JSON REPAIR HELPERS
# -----------------------
def extract_balanced_json_block(text: str) -> str | None:
    if not isinstance(text, str) or "{" not in text:
        return None

    start = text.find("{")
    if start < 0:
        return None

    depth = 0
    in_str = False
    escaped = False

    for i in range(start, len(text)):
        ch = text[i]

        if in_str:
            if escaped:
                escaped = False
            elif ch == "\\":
                escaped = True
            elif ch == '"':
                in_str = False
            continue

        if ch == '"':
            in_str = True
        elif ch == "{":
            depth += 1
        elif ch == "}":
            depth -= 1
            if depth == 0:
                return text[start:i + 1]

    return None


def coerce_0_6_int(x):
    try:
        xi = int(x)
    except Exception:
        return None
    return xi if 0 <= xi <= 6 else None


def repair_preds_from_summary(summary: str) -> dict:
    out = {
        "severity_pred": None,
        "frequency_pred": None,
        "_repair_method": None,
        "_repair_success": False,
    }

    if not isinstance(summary, str) or not summary.strip():
        return out

    s = summary.strip()

    # 1) fenced JSON
    m = FENCED_JSON_RE.search(s)
    if m:
        try:
            obj = json.loads(m.group(1).strip())
            sev = coerce_0_6_int(obj.get("severity"))
            freq = coerce_0_6_int(obj.get("frequency"))
            if sev is not None and freq is not None:
                out.update({
                    "severity_pred": sev,
                    "frequency_pred": freq,
                    "_repair_method": "fenced_json",
                    "_repair_success": True,
                })
                return out
        except Exception:
            pass

    # 2) balanced JSON block
    block = extract_balanced_json_block(s)
    if block:
        try:
            obj = json.loads(block)
            sev = coerce_0_6_int(obj.get("severity"))
            freq = coerce_0_6_int(obj.get("frequency"))
            if sev is not None and freq is not None:
                out.update({
                    "severity_pred": sev,
                    "frequency_pred": freq,
                    "_repair_method": "balanced_json",
                    "_repair_success": True,
                })
                return out
        except Exception:
            pass

    # 3) regex fallback
    sev_match = SEV_RE.search(s)
    freq_match = FREQ_RE.search(s)
    if sev_match and freq_match:
        sev = coerce_0_6_int(sev_match.group(1))
        freq = coerce_0_6_int(freq_match.group(1))
        if sev is not None and freq is not None:
            out.update({
                "severity_pred": sev,
                "frequency_pred": freq,
                "_repair_method": "regex",
                "_repair_success": True,
            })

    return out


# %%
# -----------------------
# LOAD PREDICTIONS
# -----------------------
def load_predictions_all_models(
    raw_root: Path,
    repair_from_summary: bool = True,
    overwrite_existing: bool = False,
    keep_unparsed: bool = True,
) -> pd.DataFrame:
    raw_root = Path(raw_root)
    if not raw_root.exists():
        raise FileNotFoundError(f"RAW_ROOT not found: {raw_root.resolve()}")

    model_dirs = sorted(p for p in raw_root.iterdir() if p.is_dir())
    if not model_dirs:
        raise ValueError(f"No model subfolders under: {raw_root.resolve()}")

    rows = []

    for model_dir in model_dirs:
        jsonl_files = sorted(model_dir.glob("*.jsonl"))
        if not jsonl_files:
            print(f"[warn] no jsonl in {model_dir.name}")
            continue

        for fp in jsonl_files:
            with fp.open("r", encoding="utf-8") as f:
                for raw_line in f:
                    line = raw_line.strip()
                    if not line:
                        continue

                    try:
                        obj = json.loads(line)
                    except Exception:
                        if not keep_unparsed:
                            continue
                        obj = {
                            "ran": False,
                            "severity": None,
                            "frequency": None,
                            "summary": f"[unparseable_jsonl_line] {line[:2000]}",
                        }

                    obj["_model_dir"] = model_dir.name
                    obj["_jsonl_file"] = fp.name
                    rows.append(obj)

    df = pd.DataFrame(rows)
    if df.empty:
        raise ValueError("No JSON rows loaded.")

    df["model"] = (
        df["model_id"].astype(str)
        if "model_id" in df.columns
        else df["_model_dir"].astype(str)
    )

    if "prompt_name" in df.columns:
        df["prompt"] = df["prompt_name"].astype(str)
    elif "prompt" not in df.columns:
        df["prompt"] = ""

    if "severity_pred" not in df.columns and "severity" in df.columns:
        df = df.rename(columns={"severity": "severity_pred"})
    if "frequency_pred" not in df.columns and "frequency" in df.columns:
        df = df.rename(columns={"frequency": "frequency_pred"})

    df["question"] = df.get("question", "").astype(str).str.upper().str.strip()
    df = df[df["question"].apply(lambda x: bool(Q_RE.match(str(x))))].copy()

    df["src_subject_id"] = df.get("src_subject_id", "").astype(str).str.strip()
    df["visit"] = df.get("visit", "").astype(str).str.strip()

    if "status" in df.columns:
        df["status"] = df["status"].map(normalize_status)

    df = ensure_numeric(df, ["severity_pred", "frequency_pred", "V5"])

    if "session_num" not in df.columns and "V6" in df.columns:
        df["session_num"] = df["V6"].apply(session_to_int)

    if "ran" in df.columns:
        df = df[df["ran"].eq(True)].copy()

    if repair_from_summary:
        if "summary" not in df.columns:
            df["summary"] = ""

        repaired = df["summary"].apply(repair_preds_from_summary).apply(pd.Series)
        df["_repaired_from_summary"] = repaired["_repair_success"].fillna(False).astype(bool)
        df["_repair_method"] = repaired["_repair_method"]

        if overwrite_existing:
            mask = df["_repaired_from_summary"]
            df.loc[mask, "severity_pred"] = repaired.loc[mask, "severity_pred"].astype(float)
            df.loc[mask, "frequency_pred"] = repaired.loc[mask, "frequency_pred"].astype(float)
        else:
            sev_mask = df["severity_pred"].isna() & repaired["severity_pred"].notna()
            freq_mask = df["frequency_pred"].isna() & repaired["frequency_pred"].notna()
            df.loc[sev_mask, "severity_pred"] = repaired.loc[sev_mask, "severity_pred"].astype(float)
            df.loc[freq_mask, "frequency_pred"] = repaired.loc[freq_mask, "frequency_pred"].astype(float)

    key_cols = ["src_subject_id", "visit", "question", "model", "prompt"]
    key_cols = [c for c in key_cols if c in df.columns]
    df = df.drop_duplicates(subset=key_cols, keep="last")

    keep_cols = [
        "model", "prompt", "src_subject_id", "visit", "question", "status",
        "severity_pred", "frequency_pred", "summary",
        "V5", "V6", "session_num", "instance_id", "transcript_file",
        "_model_dir", "_jsonl_file",
        "_repaired_from_summary", "_repair_method",
    ]
    keep_cols = [c for c in keep_cols if c in df.columns]

    df["model"] = df["model"].map(normalize_model_name)
    return df[keep_cols].copy()


df_pred_all = load_predictions_all_models(CFG.raw_root)
print(df_pred_all.shape)
df_pred_all.head()

# %%
# -----------------------
# LOAD GROUND TRUTH
# -----------------------
def load_ground_truth_long_from_csv(gt_csv: Path) -> pd.DataFrame:
    gt = pd.read_csv(gt_csv)

    required = {"src_subject_id", "visit"}
    if not required.issubset(gt.columns):
        raise ValueError(f"GT CSV must contain {required}, found: {set(gt.columns)}")

    gt["src_subject_id"] = gt["src_subject_id"].astype(str).str.strip()
    gt["visit"] = gt["visit"].astype(str).str.strip()

    if "interview_date" in gt.columns:
        gt["interview_date"] = pd.to_datetime(gt["interview_date"], dayfirst=True, errors="coerce")

    rows = []
    for i in range(1, 16):
        sev_col = f"chrpsychs_scr_{i}d1"
        freq_col = f"chrpsychs_scr_{i}d2"

        if sev_col not in gt.columns or freq_col not in gt.columns:
            raise ValueError(f"Missing columns for Q{i}: {sev_col}, {freq_col}")

        cols = ["src_subject_id", "visit", sev_col, freq_col]
        if "interview_date" in gt.columns:
            cols.insert(2, "interview_date")

        tmp = gt[cols].copy().rename(columns={
            sev_col: "severity",
            freq_col: "frequency",
        })
        tmp["question"] = f"Q{i}"
        rows.append(tmp)

    out = pd.concat(rows, ignore_index=True)
    out["question"] = out["question"].astype(str).str.upper()
    out["severity"] = to_num(out["severity"])
    out["frequency"] = to_num(out["frequency"])

    out.loc[(out["severity"].notna()) & ~out["severity"].between(0, 6), "severity"] = np.nan
    out.loc[(out["frequency"].notna()) & ~out["frequency"].between(0, 6), "frequency"] = np.nan

    out = out[~(out["severity"].isna() & out["frequency"].isna())].copy()
    return out


df_gt = load_ground_truth_long_from_csv(CFG.gt_csv)
print(df_gt.shape)
df_gt.head()

# %%
# -----------------------
# MERGE / PREP
# -----------------------
df_eval = df_pred_all.merge(
    df_gt,
    on=["src_subject_id", "visit", "question"],
    how="inner",
    suffixes=("_pred", ""),
)

df_eval = ensure_numeric(df_eval, ["severity_pred", "frequency_pred", "severity", "frequency"])
df_eval = df_eval.drop_duplicates().reset_index(drop=True)

# fill GT frequency=0 when severity=0 and frequency missing
mask = df_eval["severity"].eq(0) & df_eval["frequency"].isna()
df_eval.loc[mask, "frequency"] = 0.0

df_eval_filt, status_col = filter_by_status(df_eval, CFG.status_mode)

print("df_eval:", df_eval.shape)
print("df_eval_filt:", df_eval_filt.shape)
print("status_col:", status_col)

# %%
# -----------------------
# EVALUATION
# -----------------------
def add_chr_columns(df: pd.DataFrame) -> pd.DataFrame:
    out = df.copy()
    out["chr_true"] = chr_binary(out["severity"], out["frequency"])
    out["chr_pred"] = chr_binary(out["severity_pred"], out["frequency_pred"])
    return out


def count_complete_pairs(df: pd.DataFrame) -> int:
    return int(
        df["severity"].notna().sum() >= 0 and (
            df["severity"].notna()
            & df["severity_pred"].notna()
            & df["frequency"].notna()
            & df["frequency_pred"].notna()
        ).sum()
    )


def eval_item_level(df_in: pd.DataFrame) -> tuple[pd.DataFrame, pd.DataFrame]:
    df = add_chr_columns(df_in)

    overall_rows = []
    question_rows = []

    for model_name, dsub in df.groupby("model", dropna=False):
        model_name = normalize_model_name(model_name)

        r_sev, p_sev = safe_pearson(dsub["severity"], dsub["severity_pred"])
        r_frq, p_frq = safe_pearson(dsub["frequency"], dsub["frequency_pred"])
        icc_sev = icc_a1_two_col(dsub["severity"], dsub["severity_pred"])
        icc_frq = icc_a1_two_col(dsub["frequency"], dsub["frequency_pred"])
        cm = metrics_from_confusion(dsub["chr_true"], dsub["chr_pred"])

        n_rows = len(dsub)
        n_valid_labels = int(dsub["severity"].notna().sum() & dsub["frequency"].notna().sum())
        n_complete_pairs = int(
            (
                dsub["severity"].notna()
                & dsub["severity_pred"].notna()
                & dsub["frequency"].notna()
                & dsub["frequency_pred"].notna()
            ).sum()
        )
        n_repair = int(dsub.get("_repaired_from_summary", False).fillna(False).astype(bool).sum())
        n_complete_pairs_raw = max(n_complete_pairs - n_repair, 0)

        overall_rows.append({
            "model": model_name,
            "n_rows": int(n_rows),
            "n_valid_labels": int(
                (dsub["severity"].notna() & dsub["frequency"].notna()).sum()
            ),
            "n_complete_pairs": n_complete_pairs,
            "n_repair": n_repair,
            "n_complete_pairs_raw": n_complete_pairs_raw,
            "ratio_complete_pairs": (
                n_complete_pairs / max((dsub["severity"].notna() & dsub["frequency"].notna()).sum(), 1)
            ),
            "ratio_complete_pairs_raw": (
                n_complete_pairs_raw / max((dsub["severity"].notna() & dsub["frequency"].notna()).sum(), 1)
            ),
            "pearson_severity": r_sev,
            "pearson_severity_p": p_sev,
            "pearson_frequency": r_frq,
            "pearson_frequency_p": p_frq,
            "icc_severity": icc_sev,
            "icc_frequency": icc_frq,
            "auc_sev_pred": auc_from_score(dsub["chr_true"], dsub["severity_pred"]),
            "auc_sev_frq": auc_from_score(dsub["chr_true"], dsub["frequency_pred"]),
            **cm,
        })

        for question, dq in dsub.groupby("question", dropna=False):
            r_sev_q, _ = safe_pearson(dq["severity"], dq["severity_pred"])
            r_frq_q, _ = safe_pearson(dq["frequency"], dq["frequency_pred"])
            icc_sev_q = icc_a1_two_col(dq["severity"], dq["severity_pred"])
            icc_frq_q = icc_a1_two_col(dq["frequency"], dq["frequency_pred"])
            cm_q = metrics_from_confusion(dq["chr_true"], dq["chr_pred"])

            question_rows.append({
                "model": model_name,
                "question": question,
                "n_rows": int(len(dq)),
                "n_complete_pairs": int(
                    (
                        dq["severity"].notna()
                        & dq["severity_pred"].notna()
                        & dq["frequency"].notna()
                        & dq["frequency_pred"].notna()
                    ).sum()
                ),
                "pearson_severity": r_sev_q,
                "pearson_frequency": r_frq_q,
                "icc_severity": icc_sev_q,
                "icc_frequency": icc_frq_q,
                "auc_sev_pred": auc_from_score(dq["chr_true"], dq["severity_pred"]),
                **cm_q,
            })

    overall = pd.DataFrame(overall_rows).sort_values("model").reset_index(drop=True)
    by_question = pd.DataFrame(question_rows).sort_values(["model", "question"]).reset_index(drop=True)
    return overall, by_question


def eval_subject_level(df_in: pd.DataFrame) -> tuple[pd.DataFrame, pd.DataFrame]:
    df = add_chr_columns(df_in)

    subj = (
        df.groupby(["model", "src_subject_id", "visit"], dropna=False)
        .agg(
            n_items=("question", "nunique"),
            chr_any_true=("chr_true", "max"),
            chr_any_pred=("chr_pred", "max"),
        )
        .reset_index()
    )

    rows = []
    for model_name, dsub in subj.groupby("model", dropna=False):
        model_name = normalize_model_name(model_name)
        cm = metrics_from_confusion(dsub["chr_any_true"], dsub["chr_any_pred"])

        rows.append({
            "model": model_name,
            "n_subject_visits": int(len(dsub)),
            "mean_items_per_subject_visit": float(dsub["n_items"].mean()),
            **cm,
        })

    summary = pd.DataFrame(rows).sort_values("model").reset_index(drop=True)
    return subj, summary


summary_item_overall, summary_item_by_q = eval_item_level(df_eval_filt)
subj_table, summary_subject = eval_subject_level(df_eval_filt)

print(summary_subject[["model", "accuracy", "sensitivity", "specificity", "ppv", "f1", "mcc"]]
      .to_string(index=False, float_format=lambda x: f"{x:.3f}"))