#!/usr/bin/env python3
"""
What do cheap "proper pipeline" steps buy on the SAME data?

Re-runs the expanding daily folds of eval_archive.py with variants of
the text recipe and training, then trains each variant on all rows and
grades it on the 78-headline golden set (independent truth).

Variants
  shipped        headline + description, as trained today
  head           headline only
  head+clean     + Google News " - Publisher" suffix stripped
  head+clean+neg + negation marking: words after no/not/never/without/
                   fails/n't get a NOT_ prefix until the clause ends
                   (the classic bag-of-words negation trick)
  ...+weight     + training rows weighted by |DeepSeek score| so
                   confident labels count more than weak ones
  full+clean+neg headline + description with the cleaning and negation

Usage:
    venv/bin/python experiment_recipe.py
"""

from __future__ import annotations

import re
from collections import Counter
from typing import Any, Callable, Dict, List, Optional

from eval_archive import usable_rows
from eval_common import (CLASS_NAMES, DEFAULT_DB, LABEL_THRESHOLD_TRAIN,
                         NEUTRAL_BAND_GOLDEN, NAME_TO_CLASS, agreement_block,
                         fmt_pct, hr, load_golden, to_class)

NEG_CUES = re.compile(r"\b(?:no|not|never|without|fails?|neither|nor|cannot|"
                      r"stops short of|unlikely to)\b|n't\b", re.IGNORECASE)
CLAUSE_END = re.compile(r"[,.;:!?]|\b(?:as|after|amid|while|despite|but|and|"
                        r"ahead of|before)\b", re.IGNORECASE)
PUBLISHER_SUFFIX = re.compile(r"\s+[-–|]\s+[^-–|]{2,40}$")


def strip_publisher(text: str) -> str:
    return PUBLISHER_SUFFIX.sub("", text)


def mark_negation(text: str) -> str:
    """Prefix every word after a negation cue with NOT_ until the clause
    ends, so "no rate hike" and "rate hike" become different features."""
    out: List[str] = []
    pos = 0
    for m in NEG_CUES.finditer(text):
        end = CLAUSE_END.search(text, m.end())
        stop = end.start() if end else len(text)
        out.append(text[pos:m.end()])
        span = text[m.end():stop]
        out.append(re.sub(r"\b(\w+)", r"NOT_\1", span))
        pos = stop
    out.append(text[pos:])
    return "".join(out)


def make_text(row: Dict[str, Any], *, description: bool, clean: bool, neg: bool) -> str:
    from news_sentiment import article_text
    head = row["headline"]
    desc = row["description"] if description else None
    if clean:
        head = strip_publisher(head)
    if neg:
        head = mark_negation(head)
        if desc:
            desc = mark_negation(desc)
    return article_text(head, desc)


def fit(train_rows, train_labels, weights: Optional[List[float]], seed=42):
    from news_sentiment import build_pipeline
    pipe = build_pipeline(random_state=seed, min_df=2, max_features=50_000)
    if weights is None:
        pipe.fit(train_rows, train_labels)
    else:
        pipe.fit(train_rows, train_labels, classifier__sample_weight=weights)
    return pipe


def predict_cls(pipe, texts: List[str]) -> List[int]:
    classes = [str(c) for c in pipe.classes_]
    out = []
    for p in pipe.predict_proba(texts):
        by = {c: float(v) for c, v in zip(classes, p)}
        out.append(NAME_TO_CLASS[max(by, key=by.get)])
    return out


VARIANTS: Dict[str, Dict[str, Any]] = {
    "shipped (head+desc)":    dict(description=True,  clean=False, neg=False, weight=False),
    "head":                   dict(description=False, clean=False, neg=False, weight=False),
    "head+clean":             dict(description=False, clean=True,  neg=False, weight=False),
    "head+clean+neg":         dict(description=False, clean=True,  neg=True,  weight=False),
    "head+clean+neg+weight":  dict(description=False, clean=True,  neg=True,  weight=True),
    "full+clean+neg":         dict(description=True,  clean=True,  neg=True,  weight=False),
    "full+clean+neg+weight":  dict(description=True,  clean=True,  neg=True,  weight=True),
}


def main() -> int:
    thr = LABEL_THRESHOLD_TRAIN
    rows = usable_rows(DEFAULT_DB)
    days = sorted({r["day"] for r in rows})
    golden = load_golden()
    g_exp = [g["expected_xau_sign"] for g in golden]
    print(f"{len(rows)} rows, {len(days)} days; golden {len(golden)}")

    hr("Expanding daily folds (agreement with DeepSeek, OOF) and golden set (human truth)")
    print(f"  {'variant':<24} {'OOF 3-class':>11} {'OOF F1':>7} {'OOF sign':>9} | {'golden':>7} {'gold dir':>9}")
    results = {}
    for name, v in VARIANTS.items():
        exp_all: List[int] = []
        pred_all: List[int] = []
        for d in days:
            train = [r for r in rows if r["day"] < d]
            test = [r for r in rows if r["day"] == d]
            if len(train) < 40 or not test:
                continue
            labels = [CLASS_NAMES[to_class(r["xau"], thr)] for r in train]
            dist = Counter(labels)
            if len(dist) < 3 or min(dist.values()) < 2:
                continue
            tx = [make_text(r, description=v["description"], clean=v["clean"], neg=v["neg"])
                  for r in train]
            w = [0.5 + abs(r["xau"]) for r in train] if v["weight"] else None
            pipe = fit(tx, labels, w)
            te = [make_text(r, description=v["description"], clean=v["clean"], neg=v["neg"])
                  for r in test]
            pred_all.extend(predict_cls(pipe, te))
            exp_all.extend(to_class(r["xau"], thr) for r in test)
        oof = agreement_block(exp_all, pred_all)

        # final model on every row -> golden set (headline only there)
        labels = [CLASS_NAMES[to_class(r["xau"], thr)] for r in rows]
        tx = [make_text(r, description=v["description"], clean=v["clean"], neg=v["neg"])
              for r in rows]
        w = [0.5 + abs(r["xau"]) for r in rows] if v["weight"] else None
        pipe = fit(tx, labels, w)
        g_rows = [{"headline": g["headline"], "description": None} for g in golden]
        g_tx = [make_text(r, description=False, clean=v["clean"], neg=v["neg"]) for r in g_rows]
        g_pred = predict_cls(pipe, g_tx)
        gold = agreement_block(g_exp, g_pred)
        results[name] = {"oof": oof, "golden": gold}
        print(f"  {name:<24} {100 * oof['accuracy']:10.1f}% {oof['macro_f1']:7.3f} "
              f"{fmt_pct(oof['sign_hits'], oof['sign_n']):>9} | "
              f"{gold['hits']:>3}/{gold['n']:<3} {fmt_pct(gold['directional_hits'], gold['directional_n']):>9}")

    hr("Golden per category for the best headline-only variant vs shipped")
    best = "head+clean+neg+weight"
    for name in ("shipped (head+desc)", best):
        v = VARIANTS[name]
        labels = [CLASS_NAMES[to_class(r["xau"], thr)] for r in rows]
        tx = [make_text(r, description=v["description"], clean=v["clean"], neg=v["neg"]) for r in rows]
        w = [0.5 + abs(r["xau"]) for r in rows] if v["weight"] else None
        pipe = fit(tx, labels, w)
        g_tx = [make_text({"headline": g["headline"], "description": None},
                          description=False, clean=v["clean"], neg=v["neg"]) for g in golden]
        pred = predict_cls(pipe, g_tx)
        by_cat: Dict[str, List[int]] = {}
        for g, p in zip(golden, pred):
            by_cat.setdefault(g["category"], []).append(int(p == g["expected_xau_sign"]))
        print(f"  {name}: " + ", ".join(f"{c} {sum(h)}/{len(h)}" for c, h in sorted(by_cat.items())))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())