#!/usr/bin/env python3
"""Fine-tune a two-head transformer: gold-relevant? and bullish/bearish.

Data (all labels are DeepSeek's live-prompt XAU score):
  * data/history.db   crawled headlines 2019-2026 + labels table
  * newsdata.db       the live archive (xau_score -> JSON XAU -> backfill)
Headlines are deduplicated on case-folded text (median xau). Anything in
the golden set is removed from every split.

Targets per headline, from xau:
  relevance  1 if |xau| >= 0.25, 0 if |xau| <= 0.10, ignored in between
  direction  sign(xau), only where relevance = 1

Splits:
  val    newsdata rows published on or after --val-from (2026-09-18):
         the unseen days of the 2026-09-22 report, never used to choose
         anything
  dev    a seeded random 5% of the rest: picks the best epoch and the
         relevance threshold
  train  everything else
``--final`` trains on train + dev + val for the epoch count chosen by
the last non-final run, for deployment.

Outputs models/xau_transformer/torch/ (state dict, tokenizer, meta.json);
``--final`` writes models/xau_transformer/torch_final/.
Export for CPU with transformer/export_onnx.py.

Usage:
    .venv-train/bin/python transformer/train.py
    .venv-train/bin/python transformer/train.py --final
"""

from __future__ import annotations

import argparse
import json
import math
import random
import sqlite3
import sys
import time
from collections import Counter
from datetime import datetime, timezone
from pathlib import Path
from statistics import median

import numpy as np
import torch
from torch import nn
from transformers import AutoModel, AutoTokenizer, get_linear_schedule_with_warmup

ML_NEWS_DIR = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ML_NEWS_DIR))
sys.path.insert(0, str(ML_NEWS_DIR / "eval"))

from news_transformer import (MAX_LENGTH, RELEVANT_MAX_IRRELEVANT,  # noqa: E402
                              RELEVANT_MIN, clean_headline)

HISTORY_DB = ML_NEWS_DIR / "data" / "history.db"
NEWSDATA_DB = ML_NEWS_DIR / "newsdata.db"
BACKFILL = ML_NEWS_DIR / "eval" / "results" / "backfill_labels.jsonl"
OUT_DIR = ML_NEWS_DIR / "models" / "xau_transformer" / "torch"
IGNORE = -1.0


class TwoHeadModel(nn.Module):
    """Mean-pooled encoder with a relevance logit and a direction logit."""

    def __init__(self, base: str, dropout: float = 0.1):
        super().__init__()
        # transformers 5 loads the checkpoint's stored dtype (fp16 for
        # deberta-v3-small) by default; fp16 master weights overflow under
        # AdamW. Keep fp32 weights; autocast does bf16 compute.
        self.encoder = AutoModel.from_pretrained(base, dtype=torch.float32)
        hidden = self.encoder.config.hidden_size
        self.dropout = nn.Dropout(dropout)
        self.relevance = nn.Linear(hidden, 1)
        self.direction = nn.Linear(hidden, 1)

    def forward(self, input_ids, attention_mask):
        states = self.encoder(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
        mask = attention_mask.unsqueeze(-1).to(states.dtype)
        pooled = (states * mask).sum(1) / mask.sum(1).clamp(min=1.0)
        pooled = self.dropout(pooled)
        return self.relevance(pooled).squeeze(-1), self.direction(pooled).squeeze(-1)


# ── data ─────────────────────────────────────────────────────────────

def golden_texts() -> set[str]:
    from eval_common import load_golden
    return {clean_headline(g["headline"]).casefold() for g in load_golden()}


def load_rows() -> list[dict]:
    raw: list[dict] = []
    if HISTORY_DB.is_file():
        uri = f"{HISTORY_DB.resolve().as_uri()}?mode=ro"
        with sqlite3.connect(uri, uri=True) as conn:
            has_labels = conn.execute("SELECT 1 FROM sqlite_master WHERE type = 'table'"
                                      " AND name = 'labels'").fetchone()
            for headline, published_at, xau in [] if not has_labels else conn.execute(
                    "SELECT h.headline, h.published_at, l.xau FROM headlines h"
                    " JOIN labels l ON l.hash = h.hash"):
                raw.append({"headline": headline, "published_at": published_at or "",
                            "xau": float(xau), "source": "history"})
    from eval_common import load_archive_rows
    for r in load_archive_rows(NEWSDATA_DB, BACKFILL if BACKFILL.is_file() else None):
        if r["score_source"] == "llm" and r["xau"] is not None and r["headline"]:
            raw.append({"headline": r["headline"],
                        "published_at": r["published_at"] or r["first_seen"] or "",
                        "xau": float(r["xau"]), "source": "newsdata"})

    golden = golden_texts()
    grouped: dict[str, dict] = {}
    dropped_golden = 0
    for r in raw:
        try:
            text = clean_headline(r["headline"])
        except ValueError:
            continue
        key = text.casefold()
        if key in golden:
            dropped_golden += 1
            continue
        g = grouped.setdefault(key, {"text": text, "xaus": [], "dates": [], "sources": set()})
        g["xaus"].append(r["xau"])
        g["dates"].append(r["published_at"])
        g["sources"].add(r["source"])
    rows = []
    for g in grouped.values():
        xau = float(median(g["xaus"]))
        a = abs(xau)
        rel = 1.0 if a >= RELEVANT_MIN else 0.0 if a <= RELEVANT_MAX_IRRELEVANT else IGNORE
        rows.append({
            "text": g["text"], "xau": xau, "published_at": max(g["dates"]),
            "source": "newsdata" if "newsdata" in g["sources"] else "history",
            "rel": rel, "dir": (1.0 if xau > 0 else 0.0) if rel == 1.0 else IGNORE,
        })
    print(f"{len(raw)} labelled rows -> {len(rows)} unique headlines "
          f"({dropped_golden} golden-set duplicates dropped)")
    return rows


def split(rows: list[dict], val_from: str, seed: int):
    val = [r for r in rows if r["source"] == "newsdata" and r["published_at"] >= val_from]
    rest = [r for r in rows if not (r["source"] == "newsdata" and r["published_at"] >= val_from)]
    rng = random.Random(seed)
    rng.shuffle(rest)
    n_dev = max(200, len(rest) // 20)
    return rest[n_dev:], rest[:n_dev], val


def describe(name: str, rows: list[dict]) -> dict:
    rel = Counter(r["rel"] for r in rows)
    d = Counter(r["dir"] for r in rows if r["dir"] != IGNORE)
    years = Counter(r["published_at"][:4] for r in rows)
    info = {"n": len(rows), "relevant": rel[1.0], "irrelevant": rel[0.0], "grey": rel[IGNORE],
            "bullish": d[1.0], "bearish": d[0.0],
            "sources": dict(Counter(r["source"] for r in rows)),
            "years": dict(sorted(years.items()))}
    print(f"  {name:<6} n={info['n']:>6}  relevant {info['relevant']:>6}  irrelevant "
          f"{info['irrelevant']:>6}  grey {info['grey']:>5}  bull {info['bullish']:>5}  "
          f"bear {info['bearish']:>5}  {info['sources']}")
    return info


# ── training / evaluation ────────────────────────────────────────────

def batches(rows, tokenizer, batch_size, shuffle, rng=None):
    order = list(range(len(rows)))
    if shuffle:
        rng.shuffle(order)
    for start in range(0, len(order), batch_size):
        chunk = [rows[i] for i in order[start:start + batch_size]]
        enc = tokenizer([r["text"] for r in chunk], padding=True, truncation=True,
                        max_length=MAX_LENGTH, return_tensors="pt")
        rel = torch.tensor([r["rel"] for r in chunk])
        dirs = torch.tensor([r["dir"] for r in chunk])
        yield enc["input_ids"], enc["attention_mask"], rel, dirs


def masked_bce(logits, targets, pos_weight):
    mask = targets != IGNORE
    if not mask.any():
        return logits.sum() * 0.0
    return nn.functional.binary_cross_entropy_with_logits(
        logits[mask].float(), targets[mask].float(), pos_weight=pos_weight)


@torch.no_grad()
def predict(model, rows, tokenizer, device, batch_size=256):
    model.eval()
    p_rel, p_bull = [], []
    for ids, mask, _, _ in batches(rows, tokenizer, batch_size, False):
        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=device.type == "cuda"):
            r, d = model(ids.to(device), mask.to(device))
        p_rel.append(torch.sigmoid(r.float()).cpu())
        p_bull.append(torch.sigmoid(d.float()).cpu())
    return torch.cat(p_rel).numpy(), torch.cat(p_bull).numpy()


def deepseek_class(xau: float, band: float = 0.05) -> int:
    return 1 if xau > band else -1 if xau < -band else 0


def metrics(rows, p_rel, p_bull, threshold: float) -> dict:
    rel_t = np.array([r["rel"] for r in rows])
    dir_t = np.array([r["dir"] for r in rows])
    m = rel_t != IGNORE
    pred_rel = p_rel >= threshold
    tp = int(((pred_rel == 1) & (rel_t == 1)).sum())
    fp = int(((pred_rel == 1) & (rel_t == 0)).sum())
    fn = int(((pred_rel == 0) & (rel_t == 1)).sum())
    prec = tp / (tp + fp) if tp + fp else 0.0
    rec = tp / (tp + fn) if tp + fn else 0.0
    dm = dir_t != IGNORE
    dir_acc = float(((p_bull[dm] >= 0.5) == (dir_t[dm] == 1)).mean()) if dm.any() else float("nan")
    pred_cls = np.where(pred_rel, np.where(p_bull >= 0.5, 1, -1), 0)
    ds = np.array([deepseek_class(r["xau"]) for r in rows])
    cls3 = float((pred_cls == ds).mean())
    from sklearn.metrics import f1_score
    return {
        "relevance_threshold": threshold,
        "relevance_precision": prec, "relevance_recall": rec,
        "relevance_f1": 2 * prec * rec / (prec + rec) if prec + rec else 0.0,
        "relevance_accuracy": float((pred_rel[m] == (rel_t[m] == 1)).mean()),
        "direction_accuracy_on_relevant": dir_acc,
        "three_class_vs_deepseek_0.05": cls3,
        "three_class_macro_f1": float(f1_score(ds, pred_cls, average="macro")),
        "n": len(rows),
    }


def best_threshold(rows, p_rel, p_bull) -> float:
    """Relevance cut that maximises 3-class agreement on dev."""
    grid = [round(x, 2) for x in np.arange(0.2, 0.81, 0.05)]
    return max(grid, key=lambda t: metrics(rows, p_rel, p_bull, t)["three_class_vs_deepseek_0.05"])


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("--base", default="microsoft/deberta-v3-small")
    ap.add_argument("--epochs", type=int, default=4)
    ap.add_argument("--batch-size", type=int, default=64)
    ap.add_argument("--lr", type=float, default=5e-5)
    ap.add_argument("--warmup", type=float, default=0.06)
    ap.add_argument("--weight-decay", type=float, default=0.01)
    ap.add_argument("--val-from", default="2026-09-18")
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--out", type=Path, default=OUT_DIR)
    ap.add_argument("--final", action="store_true",
                    help="train on every split for the epoch count chosen by the last run")
    args = ap.parse_args()

    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    rows = load_rows()
    train, dev, val = split(rows, args.val_from, args.seed)
    print("splits:")
    data_info = {"train": describe("train", train), "dev": describe("dev", dev),
                 "val": describe("val", val)}
    previous = {}
    if args.final:
        prev_path = args.out / "meta.json"
        if not prev_path.is_file():
            raise SystemExit("--final needs a previous non-final run (its meta.json)")
        previous = json.loads(prev_path.read_text())
        args.epochs = int(previous["best_epoch"])
        train = train + dev + val
        args.out = args.out.with_name("torch_final")
        print(f"final: training on all {len(train)} rows for {args.epochs} epochs -> {args.out}")

    tokenizer = AutoTokenizer.from_pretrained(args.base)
    model = TwoHeadModel(args.base).to(device)
    rel_counts = Counter(r["rel"] for r in train)
    dir_counts = Counter(r["dir"] for r in train)
    rel_pw = torch.tensor(rel_counts[0.0] / max(1, rel_counts[1.0]), device=device)
    dir_pw = torch.tensor(dir_counts[0.0] / max(1, dir_counts[1.0]), device=device)

    no_decay = ("bias", "LayerNorm.weight", "layernorm", "norm")
    params = [
        {"params": [p for n, p in model.named_parameters() if not any(k in n for k in no_decay)],
         "weight_decay": args.weight_decay},
        {"params": [p for n, p in model.named_parameters() if any(k in n for k in no_decay)],
         "weight_decay": 0.0},
    ]
    optimizer = torch.optim.AdamW(params, lr=args.lr)
    steps = math.ceil(len(train) / args.batch_size) * args.epochs
    scheduler = get_linear_schedule_with_warmup(optimizer, int(steps * args.warmup), steps)

    args.out.mkdir(parents=True, exist_ok=True)
    rng = random.Random(args.seed)
    history = []
    best = {"score": -1.0}
    started = time.time()
    for epoch in range(1, args.epochs + 1):
        model.train()
        total = n = 0
        for ids, mask, rel, dirs in batches(train, tokenizer, args.batch_size, True, rng):
            ids, mask, rel, dirs = ids.to(device), mask.to(device), rel.to(device), dirs.to(device)
            with torch.autocast("cuda", dtype=torch.bfloat16, enabled=device.type == "cuda"):
                r_logit, d_logit = model(ids, mask)
            loss = masked_bce(r_logit, rel, rel_pw) + masked_bce(d_logit, dirs, dir_pw)
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            scheduler.step()
            total += loss.item() * len(ids)
            n += len(ids)
        entry = {"epoch": epoch, "train_loss": total / n, "seconds": round(time.time() - started)}
        if not args.final:
            pr, pb = predict(model, dev, tokenizer, device)
            thr = best_threshold(dev, pr, pb)
            entry["dev"] = metrics(dev, pr, pb, thr)
            score = entry["dev"]["three_class_vs_deepseek_0.05"]
            if score > best["score"]:
                best = {"score": score, "epoch": epoch, "threshold": thr}
                torch.save(model.state_dict(), args.out / "model.pt")
            d = entry["dev"]
            print(f"epoch {epoch}: loss {entry['train_loss']:.4f}  dev 3-class "
                  f"{d['three_class_vs_deepseek_0.05']:.3f}  rel F1 {d['relevance_f1']:.3f}  "
                  f"dir acc {d['direction_accuracy_on_relevant']:.3f}  thr {thr}  "
                  f"({entry['seconds']}s)", flush=True)
        else:
            print(f"epoch {epoch}: loss {entry['train_loss']:.4f} ({entry['seconds']}s)", flush=True)
        history.append(entry)

    if args.final:
        torch.save(model.state_dict(), args.out / "model.pt")
        best = {"epoch": args.epochs, "threshold": previous["relevance_threshold"]}
        evaluation = previous.get("evaluation", {})
    else:
        model.load_state_dict(torch.load(args.out / "model.pt", map_location=device))
        evaluation = {}
        for name, part in (("dev", dev), ("val", val)):
            pr, pb = predict(model, part, tokenizer, device)
            evaluation[name] = metrics(part, pr, pb, best["threshold"])
        v = evaluation["val"]
        print(f"best epoch {best['epoch']}, relevance threshold {best['threshold']}")
        print(f"VAL (unseen {args.val_from}+, n={v['n']}): 3-class {v['three_class_vs_deepseek_0.05']:.3f}"
              f"  macro-F1 {v['three_class_macro_f1']:.3f}  rel P/R/F1 {v['relevance_precision']:.3f}/"
              f"{v['relevance_recall']:.3f}/{v['relevance_f1']:.3f}  dir acc "
              f"{v['direction_accuracy_on_relevant']:.3f}")

    tokenizer.save_pretrained(args.out)
    meta = {
        "base_model": args.base,
        "trained_at": datetime.now(timezone.utc).isoformat(),
        "final": args.final,
        "max_length": MAX_LENGTH,
        "relevant_min_abs_xau": RELEVANT_MIN,
        "irrelevant_max_abs_xau": RELEVANT_MAX_IRRELEVANT,
        "relevance_threshold": best["threshold"],
        "best_epoch": best["epoch"],
        "hyperparameters": {k: v for k, v in vars(args).items() if k not in ("out",)},
        "data": data_info,
        "history": history,
        "evaluation": evaluation,
        "evaluation_note": ("final model trained on val too; evaluation is from the "
                            "non-final run" if args.final else "val never used for selection"),
        "labels": "DeepSeek deepseek-v4-flash live prompt XAU score",
    }
    (args.out / "meta.json").write_text(json.dumps(meta, indent=2, default=str) + "\n")
    print(f"wrote {args.out}/model.pt and meta.json")
    return 0


if __name__ == "__main__":
    sys.exit(main())