#!/usr/bin/env python3
"""exp16 — arm assembly A1/A2/A3/A3-abl + §D3 metrics (research/20 §D2-D4). miner-v1. $0.

  A1  = V-A                    freq × contrast × c-value
  A2  = V-B  = A1 + spread     translation-spread boost + dst-variants (confounded lower bound on records.json)
  A3  = V-C  = A2 + patterns   surname/title/topo/formant/rank-grade + Palladius; PROPOSES sub-floor terms
  A3-abl = V-C with λ=0        patterns WITHOUT the spread signal (isolates each signal's contribution)

Metric structure (§D3, §D4-а — signals reported SEPARATELY, no manufactured convergence):
  • recall @PROPOSED (candidate-set membership) — the BLINDNESS axis: "where is each arm blind?"
    A1/A2 share a candidate set (spread only re-ranks); A3/A3-abl share the pattern-extended set.
    This is where V-A (freq-floor-blind on f<3) and V-C (patterns close f<3) differ.
  • recall@top-K + pseudo-precision@K — the RANKING/precision axis + trade-off curve.
  • catastrophe screen (top-50 of the WINNER), threshold ±50% sensitivity.
  • spread's contribution (A2 vs A1, A3 vs A3-abl) is ISOLATED — on injected drafts it is ~0/negative
    (confound §D1); measured ecologically only on the cold-start slice.

Thresholds chosen on TUNING half (ch1-15), frozen (FROZEN), TEST half (ch16-25) reported (§D1).
"""
from __future__ import annotations

import json
from dataclasses import dataclass, field

import exp16_common as X
import detectors as D
import patterns as P
import spread as SP
import palladius as PAL

# ── frozen miner-v1 config (thresholds tuned on ch1-15; pinned in the freeze commit) ──────────────
FROZEN = dict(
    freq_floor=X.FREQ_FLOOR,          # 3
    subsume_alpha=D.SUBSUME_ALPHA,    # 0.80
    ngram_max=X.NGRAM_MAX,            # 6
    lam=0.0,                          # V-B spread coefficient — 0 on injected drafts (confound §D1); >0 on cold-start
    formant_min_partners=3,
    formant_min_over_rep=15.0,
    bonus=dict(name=140.0, place=140.0, title=120.0, term=80.0),
    pattern_source_weight=dict(surname=1.0, ordinal_title=1.0, title_suffix=0.9, topo_suffix=0.9,
                               rank_grade=1.0, formant_suffix=0.6, formant_prefix=0.5,
                               title_bare=0.4, palladius=0.8),
    subfloor_freq_scale=40.0,
    spread_freq_min=5,                # only compute spread for candidates with freq>=this (cheap; rare ill-defined)
    top_k=90,                         # operating point for recall@K/precision reporting
)


@dataclass
class ScoredCand:
    src: str
    score: float
    freq: int
    types: list = field(default_factory=list)
    evidence: list = field(default_factory=list)
    dst_variants: list = field(default_factory=list)
    spread: float = 0.0
    from_pattern: bool = False


class Arms:
    def __init__(self, chunks, contrast, cfg=FROZEN, use_spread=True):
        self.chunks = chunks
        self.C = contrast
        self.cfg = cfg
        self.va = D.VA(chunks, contrast, freq_floor=cfg["freq_floor"],
                       subsume_alpha=cfg["subsume_alpha"]).build()
        self.va_ranked = self.va.score_all()                    # A1
        self.pats, self.formants = P.pattern_candidates(
            chunks, self.va.cand_freq, contrast,
            min_partners=cfg["formant_min_partners"], min_over_rep=cfg["formant_min_over_rep"])
        self.sm = SP.SpreadModel(chunks) if use_spread else None

    def _spread_of(self, src, freq):
        if self.sm and freq >= self.cfg["spread_freq_min"]:
            return self.sm.spread(src)
        return 0.0

    def _dst_of(self, src, freq):
        if self.sm and freq >= self.cfg["spread_freq_min"]:
            return self.sm.dst_variants(src)[0]
        return []

    def arm_A1(self):
        return list(self.va_ranked)

    def arm_A2(self, lam=None):
        lam = self.cfg["lam"] if lam is None else lam
        out = []
        for c in self.va_ranked:
            sp = self._spread_of(c.src, c.freq)
            out.append(ScoredCand(src=c.src, score=c.score * (1 + lam * sp), freq=c.freq, spread=sp,
                                  dst_variants=self._dst_of(c.src, c.freq)))
        out.sort(key=lambda s: (-s.score, -s.freq, s.src))
        return out

    def arm_A3(self, lam=None):
        lam = self.cfg["lam"] if lam is None else lam
        cfg = self.cfg
        merged: dict[str, ScoredCand] = {}
        for c in self.va_ranked:
            sp = self._spread_of(c.src, c.freq) if lam else 0.0
            merged[c.src] = ScoredCand(src=c.src, score=c.score * (1 + lam * sp), freq=c.freq, spread=sp,
                                       dst_variants=self._dst_of(c.src, c.freq))
        for cand, info in self.pats.items():
            pw = max((cfg["pattern_source_weight"].get(ev.split(":")[0], 0.3) for ev in info["evidence"]),
                     default=0.3)
            typ = info["types"][0] if info["types"] else "term"
            bonus = cfg["bonus"].get(typ, cfg["bonus"]["term"]) * pw
            if cand in merged:
                merged[cand].score += bonus
                for t in info["types"]:
                    if t not in merged[cand].types:
                        merged[cand].types.append(t)
                merged[cand].evidence += info["evidence"][:3]
                merged[cand].from_pattern = True
            else:
                f = X.count_occurrences(cand, self.chunks)
                merged[cand] = ScoredCand(src=cand, score=bonus + cfg["subfloor_freq_scale"] * f * pw,
                                          freq=f, types=list(info["types"]), evidence=info["evidence"][:3],
                                          dst_variants=self._dst_of(cand, f), from_pattern=True)
        # Palladius ru-side confirmation (gated on capitalized-predominant name lemmas — blocks the
        # 'найти'=най+ти class of false positives)
        for sc in merged.values():
            for lm, d, _ in sc.dst_variants[:2]:
                if PAL.is_palladius_token(lm, min_syllables=2) and (self.sm and self.sm.is_name_lemma(lm)):
                    sc.score += cfg["bonus"]["name"] * cfg["pattern_source_weight"]["palladius"] * 0.5
                    if "name" not in sc.types:
                        sc.types.append("name")
                    sc.evidence.append(f"palladius:{lm}")
                    break
        out = list(merged.values())
        out.sort(key=lambda s: (-s.score, -s.freq, s.src))
        return out

    def arm_A3_abl(self):
        return self.arm_A3(lam=0.0)


# ── metrics ───────────────────────────────────────────────────────────────────────────────────────
def candidate_set(ranked):
    return {c.src for c in ranked}


def recall_by(ranked, gt, chunks, top_k=None, include_annotation=True):
    """recall by stratum & type. top_k=None => @PROPOSED (candidate-set membership = blindness axis)."""
    if top_k is None:
        cset = candidate_set(ranked)
        def hit(e):
            return any(s in cset for s in e.norm_surfaces)
    else:
        idx = {c.src: i for i, c in enumerate(ranked)}
        def hit(e):
            return any((idx.get(s) is not None and idx[s] < top_k) for s in e.norm_surfaces)
    buckets = {k: [] for k in ("overall", "f>=10", "f3-9", "f<3",
                               "name", "title", "place", "term", "nickname")}
    misses = {}
    for e in gt:
        f = X.gt_occurrences(e, chunks, include_annotation)
        h = hit(e)
        buckets["overall"].append(h)
        buckets[X.freq_stratum(f)].append(h)
        if e.typ in buckets:
            buckets[e.typ].append(h)
        if not h:
            misses.setdefault("overall", []).append(f"{e.src}({X.freq_stratum(f)})")
    return {k: (sum(v) / len(v) if v else None, len(v)) for k, v in buckets.items()}, misses


def pseudo_precision(ranked, gt, top_k):
    gtsurf = {s for e in gt for s in e.norm_surfaces}
    top = ranked[:top_k]
    return (sum(1 for c in top if c.src in gtsurf) / len(top)) if top else 0.0


def catastrophe_screen(ranked, top_k=50):
    idx = {c.src: i for i, c in enumerate(ranked)}
    res = {s: idx.get(X.norm(s)) for s in X.CATASTROPHE}
    return all(r is not None and r < top_k for r in res.values()), res


def gt_in_chapters(gt, chunks, chapters, include_annotation=True):
    sub = [c for c in chunks if c.chapter in chapters]
    return [e for e in gt if X.gt_occurrences(e, sub, include_annotation) >= 1]


def summarize(name, ranked, gt, chunks, K):
    rp, miss_p = recall_by(ranked, gt, chunks, top_k=None)      # @proposed
    rk, _ = recall_by(ranked, gt, chunks, top_k=K)              # @top-K
    ok, res = catastrophe_screen(ranked, 50)
    pp = pseudo_precision(ranked, gt, K)
    return dict(name=name, n=len(ranked), catastrophe_ok=ok, catastrophe=res, pseudo_prec_at_K=round(pp, 3),
                recall_proposed={k: (round(v[0], 3) if v[0] is not None else None, v[1]) for k, v in rp.items()},
                recall_at_K={k: (round(v[0], 3) if v[0] is not None else None, v[1]) for k, v in rk.items()},
                misses_proposed=miss_p.get("overall", []))


if __name__ == "__main__":
    chunks = X.load_chunks()
    gt = X.load_gt()
    C = X.Contrast()
    arms = Arms(chunks, C)
    A = {"A1": arms.arm_A1(), "A2": arms.arm_A2(lam=8.0), "A3": arms.arm_A3(lam=0.0),
         "A3-abl": arms.arm_A3_abl()}
    K = FROZEN["top_k"]
    print(f"=== exp16 arms — FULL slice (57 chunks); K={K} ===")
    print(f"detected formants: {sorted(arms.formants, key=lambda c: -arms.formants[c]['over_rep'])}\n")
    for name, ranked in A.items():
        s = summarize(name, ranked, gt, chunks, K)
        rp, rk = s["recall_proposed"], s["recall_at_K"]
        print(f"{name}: n={s['n']} catastrophe={'PASS' if s['catastrophe_ok'] else 'FAIL'} {s['catastrophe']}")
        print(f"   recall@PROPOSED: overall={rp['overall'][0]} f>=10={rp['f>=10'][0]} f3-9={rp['f3-9'][0]} f<3={rp['f<3'][0]}")
        print(f"   recall@top{K}:    overall={rk['overall'][0]} f>=10={rk['f>=10'][0]} f3-9={rk['f3-9'][0]} f<3={rk['f<3'][0]}  pseudo-prec={s['pseudo_prec_at_K']}")
        print(f"   by type @proposed: " + " ".join(f"{t}={rp[t][0]}({rp[t][1]})" for t in ("name","title","place","term","nickname") if rp[t][0] is not None))
        print()
