#!/usr/bin/env python3
"""Auditoria pública dos números do card do Veredicto/Trium.

Recalcula, a partir dos preds crus da régua (753 casos, split seed 42,
T=1), TODOS os números de qualidade publicados no model card, incluindo os
bootstraps pareados (10k reamostragens por caso, seed 42).

Uso: python3 auditoria_card.py --dir <pasta_com_os_preds>

Espera em --dir os quatro arquivos de preds (formato: uma linha JSON por
resposta com campos caso/q/rotulo/pred/option_ids):
  preds_v6.jsonl  preds_b2.jsonl  preds_e1.jsonl  preds_v4.jsonl
e o ouro em ../../data/processed/ouro_gold.jsonl (override: --ouro).

Composições: Trium = B2(categoria) + E1(decisão) + v4(urgência);
CompA = v6(cat/dec) + v4(urg).
"""
import argparse
import json
import random
import sys
from pathlib import Path

TAREFAS = ("categoria", "decisao", "urgencia")


def load_preds(path):
    d = {}
    with open(path, encoding="utf-8") as f:
        for line in f:
            r = json.loads(line)
            d[(r["caso"], r["q"])] = int(r["pred"] == r["rotulo"])
    return d


def load_raw(path):
    """(caso,q) -> índice ouro do rótulo, na ordem de option_ids do próprio arquivo."""
    d = {}
    with open(path, encoding="utf-8") as f:
        for line in f:
            r = json.loads(line)
            d[(r["caso"], r["q"])] = int(r["rotulo"])
    return d


def load_pred_raw(path):
    """(caso,q) -> índice predito, na ordem de option_ids do próprio arquivo."""
    d = {}
    with open(path, encoding="utf-8") as f:
        for line in f:
            r = json.loads(line)
            d[(r["caso"], r["q"])] = int(r["pred"])
    return d


def load_ops(path):
    """(caso,q) -> lista de option_ids (ordem em que pred/rotulo indexam)."""
    d = {}
    with open(path, encoding="utf-8") as f:
        for line in f:
            r = json.loads(line)
            d[(r["caso"], r["q"])] = r["option_ids"]
    return d


def load_subset(path):
    """(caso,q) -> subset ('nucleo' = hard core arbitrado por humano)."""
    d = {}
    with open(path, encoding="utf-8") as f:
        for line in f:
            r = json.loads(line)
            d[(r["caso"], r["q"])] = r.get("subset", "nao_informado")
    return d


def load_ouro(path, seed=42):
    """rótulos do teste (mesmo split da régua) como hit-flags de um oráculo."""
    here = Path(__file__).resolve().parent
    # avaliar_veredicto/veredicto_contrato: ao lado deste script (pacote HF)
    # ou no scripts/ do repo (quando publicado a partir do checkout).
    for cand in (here, here.parents[2] / "scripts"):
        if (cand / "avaliar_veredicto.py").exists():
            sys.path.insert(0, str(cand))
            break
    from avaliar_veredicto import chave_caso, split_por_caso
    rows = [json.loads(l) for l in open(path, encoding="utf-8")]
    _, teste = split_por_caso(rows, seed)
    gold = {}
    for row in rows:
        ck = chave_caso(row["estado"])
        if ck not in teste:
            continue
        i = row["pergunta"]["instrucao"]
        t = "categoria" if "categoria" in i else ("decisao" if "decis" in i else "urgencia")
        # NO ouro, rotulo indexa a ordem canônica de negócio de opcoes (não a
        # alfabética — conferido contra a medição histórica do JEV).
        ops = [o.strip() for o in row["pergunta"]["opcoes"]]
        gold[(ck, t)] = ops[row["rotulo"]]
    return gold


def per_caso(hits, casos):
    """media de acertos por caso (3 respostas)."""
    return [sum(hits[(c, t)] for t in TAREFAS) / 3 for c in casos]


def bootstrap_diff(va, vb, n=10_000, seed=42):
    """IC95 da diferença de médias (a-b) reamostrando casos; recebe vetores por caso."""
    rng = random.Random(seed)
    m = len(va)
    deltas = []
    for _ in range(n):
        sa = sb = 0.0
        for _ in range(m):
            k = rng.randrange(m)
            sa += va[k]
            sb += vb[k]
        deltas.append((sa - sb) / m)
    deltas.sort()
    return deltas[int(0.025 * n)] * 100, deltas[int(0.975 * n)] * 100


def acc(casos, fonte):
    """fonte: dict tarefa→hits (composição roteada) OU um único dict de hits (mono)."""
    roteado = all(isinstance(k, str) for k in fonte)
    out = {}
    for t in TAREFAS:
        src = fonte[t] if roteado else fonte
        out[t] = sum(src[(c, t)] for c in casos) / len(casos)
    out["geral"] = sum(out[t] for t in TAREFAS) / 3
    return out


def macro_f1_categoria(raw_cat, pred_cat, ops):
    """macro-F1 de categoria com a mesma regra da linhagem Decisum (avaliar_v4_bench):
    média do F1 sobre as classes que aparecem no ouro OU nas predições do teste
    (no teste atual: 12 classes; fraude_suspeita e reembolso têm n ínfimo no ouro)."""
    cls = {}
    def slot(lab):
        return cls.setdefault(lab, [0, 0, 0])
    for key, gold_idx in raw_cat.items():
        real, pred = ops[key][gold_idx], ops[key][pred_cat[key]]
        if real == pred:
            slot(real)[0] += 1
        else:
            slot(real)[2] += 1
            slot(pred)[1] += 1
    f1s = [2 * tp / max(1, 2 * tp + fp + fn) for tp, fp, fn in cls.values()]
    return sum(f1s) / len(f1s), len(f1s)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--dir", required=True)
    ap.add_argument("--ouro", default=None)
    args = ap.parse_args()
    d = Path(args.dir)
    v6 = load_preds(d / "preds_v6.jsonl")
    b2 = load_preds(d / "preds_b2.jsonl")
    e1 = load_preds(d / "preds_e1.jsonl")
    v4 = load_preds(d / "preds_v4.jsonl")
    keys = set(v6)
    assert set(b2) == keys and set(e1) == keys and set(v4) == keys, "preds desalinhados"
    casos = sorted({c for c, _ in keys})
    print(f"# {len(casos)} casos, {len(keys)} respostas")

    subset = load_subset(d / "preds_v6.jsonl")
    assert all(subset[k] in ("nucleo", "aleatorio") for k in keys), "subsets inesperados"
    nucleo = [k for k in keys if subset[k] == "nucleo"]
    ops = load_ops(d / "preds_v6.jsonl")

    def hard(src):
        """acc no hard core (subset 'nucleo'); src = roteado (tarefa->hits) ou mono."""
        roteado = all(isinstance(k, str) for k in src)
        return sum((src[k[1]] if roteado else src)[k] for k in nucleo) / len(nucleo)

    def macro_f1(src_raw_cat):
        """macro-F1 de categoria no critério da linhagem Decisum (ver macro_f1_
        categoria); src_raw_cat: (caso,'categoria') -> índice predito."""
        raw_cat = load_raw(d / "preds_v6.jsonl")
        raw_cat = {k: v for k, v in raw_cat.items() if k[1] == "categoria"}
        return macro_f1_categoria(raw_cat, src_raw_cat,
                                  {k: ops[k] for k in raw_cat})

    casos_n = sorted({c for c, _ in nucleo})

    def hard_diff(nome, other_hits):
        """Δ de hard-core (média por caso do núcleo) com IC bootstrap pareado."""
        ft = [sum(tri[(c, t)] for t in TAREFAS) / 3 for c in casos_n]
        ot = [sum(other_hits[(c, t)] for t in TAREFAS) / 3 for c in casos_n]
        lo, hi = bootstrap_diff(ft, ot)
        dm = 100 * (sum(ft) - sum(ot)) / len(casos_n)
        print(f"hard core Trium vs {nome}: Δ {dm:+.2f} p.p. IC95 [{lo:+.2f}; {hi:+.2f}]")

    base_path = d / "preds_base_clef.jsonl"
    base = load_preds(base_path) if base_path.exists() else None
    # No preds_base_clef.jsonl, rotulo/pred indexam a ordem canônica de NEGÓCIO
    # das opções (validado contra o ouro: reproduz 0.5817/0.2988/0.5511 exatos) —
    # diferente dos demais preds, que indexam option_ids (alfabética).
    BIZ = {"categoria": ["reembolso", "troca_devolucao", "fraude_suspeita", "prazo_envio",
                         "pagamento", "frete", "produto_defeito", "chargeback",
                         "duvida_comercial", "outro", "incidente", "whatsapp_meta",
                         "bug_produto", "integracao", "conta_config", "plano_creditos",
                         "duvida_produto"]}

    def macro_f1_names(pred_names):
        """macro-F1 de categoria sobre (caso,'categoria')->nome previsto."""
        raw_cat = {k: v for k, v in load_raw(d / "preds_v6.jsonl").items() if k[1] == "categoria"}
        gold_names = {k: load_ops(d / "preds_v6.jsonl")[k][v] for k, v in raw_cat.items()}
        cls = {}
        for k, pred in pred_names.items():
            real = gold_names[k]
            cls.setdefault(real, [0, 0, 0])
            cls.setdefault(pred, [0, 0, 0])
            if real == pred:
                cls[real][0] += 1
            else:
                cls[real][2] += 1
                cls[pred][1] += 1
        f1 = [2 * tp / max(1, 2 * tp + fp + fn) for tp, fp, fn in cls.values()]
        return sum(f1) / len(f1), len(f1)

    trium = {"categoria": b2, "decisao": e1, "urgencia": v4}
    compa = {"categoria": v6, "decisao": v6, "urgencia": v4}
    modelos = [("TRIUM", trium), ("CompA", compa),
               ("v6", v6), ("b2", b2), ("e1", e1), ("v4", v4)]
    if base is not None:
        assert set(base) == keys, "preds da base desalinhados"
        modelos.append(("base", base))
    for nome, fonte in modelos:
        a = acc(casos, fonte)
        line = (f"{nome:6s} cat {a['categoria']:.4f} dec {a['decisao']:.4f} "
                f"urg {a['urgencia']:.4f} geral {a['geral']:.4f} "
                f"hard {hard(fonte):.4f}")
        # macro-F1 de categoria: TRIUM usa a cabeça de categoria (b2); a base
        # indexa a ordem de negócio (BIZ); demais indexam option_ids (ops).
        cat_file = {"TRIUM": "preds_b2.jsonl", "v6": "preds_v6.jsonl",
                    "b2": "preds_b2.jsonl", "e1": "preds_e1.jsonl"}.get(nome)
        if cat_file is not None:
            src_pred = load_pred_raw(d / cat_file)
            names = {k: ops[k][p] for k, p in src_pred.items() if k[1] == "categoria"}
            m, n_cls = macro_f1_names(names)
            line += f" macroF1 {m:.4f}({n_cls})"
        elif nome == "base":
            base_pred = load_pred_raw(base_path)
            names = {k: BIZ["categoria"][p] for k, p in base_pred.items()
                     if k[1] == "categoria"}
            m, n_cls = macro_f1_names(names)
            line += f" macroF1 {m:.4f}({n_cls})"
        print(line)

    def compose(fonte):
        if all(isinstance(k, str) for k in fonte):
            return {(c, t): fonte[t][(c, t)] for c in casos for t in TAREFAS}
        return fonte

    tri = compose(trium)
    vtri = per_caso(tri, casos)
    for rot, other_src in (("vs v6", v6), ("vs CompA", compa)):
        voth = per_caso(compose(other_src), casos)
        lo, hi = bootstrap_diff(vtri, voth)
        dm = 100 * (sum(vtri) - sum(voth)) / len(casos)
        print(f"bootstrap Trium {rot}: Δ {dm:+.2f} p.p. IC95 [{lo:+.2f}; {hi:+.2f}]")
    if base is not None:
        vbase = per_caso(base, casos)
        lo, hi = bootstrap_diff(vtri, vbase)
        dm = 100 * (sum(vtri) - sum(vbase)) / len(casos)
        print(f"bootstrap Trium vs base (clef-flash zero-shot): "
              f"Δ {dm:+.2f} p.p. IC95 [{lo:+.2f}; {hi:+.2f}]")
        hard_diff("base", base)

    jev_path = d / "jev_ouro.jsonl"
    if jev_path.exists():
        # Preds brutos (índices na ordem alfabética de option_ids) para poder
        # pontuar o JEV (que publica texto) sem precisar do ouro.
        raw = load_raw(d / "preds_v6.jsonl")
        jev_texto = {}
        for line in open(jev_path, encoding="utf-8"):
            r = json.loads(line)
            for t in TAREFAS:
                jev_texto[(r["caso"], t)] = str(r[t]).strip()
        if args.ouro:
            gold = load_ouro(args.ouro)
            jev = {k: int(jev_texto[k] == gold[k]) for k in keys}
        else:
            # Sem o ouro local: as opções em cada preds_v*.jsonl vêm em
            # option_ids (ordem alfabética) — basta indexar o texto do JEV aí.
            ops = load_ops(d / "preds_v6.jsonl")
            jev = {k: int(ops[k].index(jev_texto[k]) == raw[k]) for k in keys}
        a = acc(casos, jev)
        m_jev, n_jev = macro_f1_names({k: jev_texto[k] for k in keys if k[1] == "categoria"})
        print(f"JEV    cat {a['categoria']:.4f} dec {a['decisao']:.4f} "
              f"urg {a['urgencia']:.4f} geral {a['geral']:.4f} "
              f"hard {hard(jev):.4f} macroF1 {m_jev:.4f}({n_jev})")
        vjev = per_caso(jev, casos)
        lo, hi = bootstrap_diff(vtri, vjev)
        dm = 100 * (sum(vtri) - sum(vjev)) / len(casos)
        print(f"bootstrap Trium vs JEV: Δ {dm:+.2f} p.p. IC95 [{lo:+.2f}; {hi:+.2f}]")
        hard_diff("JEV", jev)
        for t in TAREFAS:
            ft = [tri[(c, t)] for c in casos]
            jt = [jev[(c, t)] for c in casos]
            lo, hi = bootstrap_diff(ft, jt)
            dm = 100 * (sum(ft) - sum(jt)) / len(casos)
            print(f"  {t:10s} Δ {dm:+.2f} IC95 [{lo:+.2f}; {hi:+.2f}]")


if __name__ == "__main__":
    main()