"""Taxonomies ``nsb-v1`` and ``nsb-v2``: heads, gating rules, loss masks and the gated output (pure Python).

Spec: docs/sentiment-labelling.md §2.1-2.2 (options and gates) and docs/classifier-plan.md §4.2 (heads and masks).
The bundle's ``taxonomy.json`` (written by ``naver-psy export-training``) holds field and option order; the rules
live here, and a ``taxonomy.json`` that states different ones is refused, not silently ignored.

Canonical ``taxonomy.json`` (what the export writes and every consumer reads)::

    {"version": "nsb-v1",
     "fields": {"relevance": ["market", "politics", "off_topic", "unclear"], ..., "emotion_valence": [...]},
     "absence": {"relevance": "off_topic", "target": "none", ...},
     "derived": {"emotion_valence": {"from": "emotion_fine", "map": {"euphoria": "positive", ...}}},
     "relevant_option": "market", "gated_by_relevance": [...], "commercial_ad_options": [...],
     "gate_order": [...], "field_weights": {"direction": 2.0}, "class_weights": {}, "macro_f1_exclude": {}}

Gates at inference (labelling §2.2, in this order; forbidden options get probability 0, the rest is renormalised):
1. top(relevance) != market -> target, direction, sarcasm, behavior, pnl, emotion_fine, emotion_valence, crowd = absence;
2. top(target) == none -> direction = none;
3./4. sarcasm and commercial do not restrict other heads (commercial is handled by the aggregation weights);
5. emotion_valence leads: emotion_fine is restricted to the options that map to top(emotion_valence).

Loss masks (classifier-plan §4.2): gate-decided absences are excluded, so each head is a conditional:
target and the mood heads only when relevance = market; direction only when market and target != none.

``nsb-v2`` (labelling §2.5; the nsb-v1 labels mapped at export, ``naver_psy/labelling/taxonomy_v2.py``) has the same
fields, absences and gates; ``sarcasm`` is {ironic, joke_meme, none} and ``commercial`` is {commercial, none}. Two
training/evaluation settings travel with the taxonomy (version defaults, overridable in ``taxonomy.json`` and on the
command line): ``class_weights`` (per-option loss multipliers; nsb-v2: pnl ``breakeven`` 3.0) and
``macro_f1_exclude`` (options left out of the headline macro-F1, still reported per class; nsb-v2: emotion_fine
``mixed``).
"""
from __future__ import annotations

import json
from pathlib import Path
from typing import Any, Iterable, Mapping

VERSION = "nsb-v1"
RELEVANT = "market"
VALENCE_FIELD = "emotion_valence"
FINE_FIELD = "emotion_fine"
FINE_TO_VALENCE = {
    "euphoria": "positive", "optimism": "positive", "fomo": "positive", "relief": "positive",
    "anxiety": "negative", "fear_panic": "negative", "anger": "negative", "regret": "negative",
    "resignation": "negative", "contempt_mockery": "negative", "mixed": "mixed", "not_expressed": "none",
}
# The nine asked fields (labelling §2.1) plus the derived valence, in the order of the §2.3 record example.
NSB_V1_FIELDS: dict[str, list[str]] = {
    "relevance": ["market", "politics", "off_topic", "unclear"],
    "target": ["this_stock", "other_stock", "sector_theme", "market_overall", "none"],
    "direction": ["bullish", "bearish", "neutral_range", "mixed", "none"],
    "sarcasm": ["ironic", "joke_meme", "unclear", "none"],
    "behavior": ["buy_add", "sell_reduce", "hold_wait", "none"],
    "pnl": ["profit", "loss", "breakeven", "none"],
    "emotion_fine": ["euphoria", "optimism", "fomo", "relief", "anxiety", "fear_panic", "anger", "regret",
                     "resignation", "contempt_mockery", "mixed", "not_expressed"],
    "crowd": ["herding_call", "blame_actor", "crowd_reference", "none"],
    "commercial": ["leading_room", "pump_promo", "contact_link", "other_ad", "unclear", "none"],
    "emotion_valence": ["positive", "negative", "mixed", "none"],
}
NSB_V1_ABSENCE = {
    "relevance": "off_topic", "target": "none", "direction": "none", "sarcasm": "none", "behavior": "none",
    "pnl": "none", "emotion_fine": "not_expressed", "crowd": "none", "commercial": "none", "emotion_valence": "none",
}
GATED_BY_RELEVANCE = ("target", "direction", "sarcasm", "behavior", "pnl", "emotion_fine", "emotion_valence", "crowd")
GATE_ORDER = ("relevance", "target", "direction", "sarcasm", "commercial", "emotion_valence", "emotion_fine")
COMMERCIAL_AD = ("leading_room", "pump_promo", "contact_link", "other_ad")
PRODUCT_FIELDS = ("relevance", "direction", "commercial")
FIELD_WEIGHTS = {"direction": 2.0}

# nsb-v2: only sarcasm and commercial change (the asked-field lists are owned by naver_psy/labelling/taxonomy_v2.py;
# a test checks they agree). Absences are those of nsb-v1.
NSB_V2_FIELDS: dict[str, list[str]] = dict(NSB_V1_FIELDS, sarcasm=["ironic", "joke_meme", "none"],
                                           commercial=["commercial", "none"])
SPECS: dict[str, dict[str, list[str]]] = {"nsb-v1": NSB_V1_FIELDS, "nsb-v2": NSB_V2_FIELDS}
AD_OPTIONS: dict[str, tuple[str, ...]] = {"nsb-v1": COMMERCIAL_AD, "nsb-v2": ("commercial",)}
# For consumers that see option names but not the version (rollups, PPI): the names never collide across versions.
AD_OPTIONS_ANY = COMMERCIAL_AD + ("commercial",)
# Training/evaluation defaults per version (owner-approved for nsb-v2, 2026-10-19). breakeven 3.0 = sqrt of the
# mean pnl profit/loss count over the breakeven count in the teacher labels (sqrt(1051/112) = 3.06), rounded down.
CLASS_WEIGHTS: dict[str, dict[str, dict[str, float]]] = {"nsb-v1": {}, "nsb-v2": {"pnl": {"breakeven": 3.0}}}
MACRO_F1_EXCLUDE: dict[str, dict[str, list[str]]] = {"nsb-v1": {}, "nsb-v2": {"emotion_fine": ["mixed"]}}


class TaxonomyError(ValueError):
    """A taxonomy definition that does not match its version (nsb-v1 or nsb-v2)."""


class Taxonomy:
    """Field and option order of one taxonomy version, its rules, class weights and macro-F1 exclusions."""

    def __init__(self, fields: Mapping[str, Iterable[str]], absence: Mapping[str, str] | None = None,
                 version: str = VERSION, class_weights: Mapping[str, Mapping[str, float]] | None = None,
                 macro_f1_exclude: Mapping[str, Iterable[str]] | None = None) -> None:
        if version not in SPECS:
            raise TaxonomyError(f"unknown taxonomy version {version!r} (known: {sorted(SPECS)}); "
                                "versions are never mixed")
        spec = SPECS[version]
        self.version = version
        self.fields = {f: list(o) for f, o in fields.items()}
        self.names = list(self.fields)
        self.absence = dict(absence or NSB_V1_ABSENCE)
        if set(self.fields) != set(spec):
            raise TaxonomyError(f"fields differ from {version}: {sorted(set(self.fields) ^ set(spec))}")
        for f, options in self.fields.items():
            if len(set(options)) != len(options):
                raise TaxonomyError(f"duplicate option in {f}")
            if set(options) != set(spec[f]):
                raise TaxonomyError(f"options of {f} differ from {version}: {sorted(set(options) ^ set(spec[f]))}")
            if self.absence.get(f) != NSB_V1_ABSENCE[f]:
                raise TaxonomyError(f"absence option of {f} differs from {version}")
        cw = CLASS_WEIGHTS[version] if class_weights is None else class_weights
        self.class_weights = {f: {o: float(w) for o, w in ow.items()} for f, ow in cw.items() if ow}
        mx = MACRO_F1_EXCLUDE[version] if macro_f1_exclude is None else macro_f1_exclude
        self.macro_f1_exclude = {f: list(opts) for f, opts in mx.items() if opts}
        for name, table in (("class_weights", self.class_weights), ("macro_f1_exclude", self.macro_f1_exclude)):
            for f, opts in table.items():
                unknown = sorted(set(opts) - set(self.fields.get(f, ())))
                if f not in self.fields or unknown:
                    raise TaxonomyError(f"{name}: unknown field or options {f}: {unknown}")
        if any(not 0 < w < float("inf") for ow in self.class_weights.values() for w in ow.values()):
            raise TaxonomyError("class_weights must be positive and finite")

    def replace(self, *, class_weights: Mapping[str, Mapping[str, float]] | None = None,
                macro_f1_exclude: Mapping[str, Iterable[str]] | None = None) -> "Taxonomy":
        """The same taxonomy with other class weights or macro-F1 exclusions (None keeps the current ones)."""
        return Taxonomy(self.fields, self.absence, self.version,
                        self.class_weights if class_weights is None else class_weights,
                        self.macro_f1_exclude if macro_f1_exclude is None else macro_f1_exclude)

    @property
    def commercial_ad(self) -> tuple[str, ...]:
        """The commercial options that count as advertising (q = 1 - their probability, classifier-plan §7.2)."""
        return AD_OPTIONS[self.version]

    # ------------------------------------------------------------------ (de)serialisation
    @classmethod
    def default(cls, version: str = VERSION) -> "Taxonomy":
        if version not in SPECS:
            raise TaxonomyError(f"unknown taxonomy version {version!r}")
        return cls(SPECS[version], NSB_V1_ABSENCE, version)

    @classmethod
    def load(cls, path: str | Path) -> "Taxonomy":
        """Read a canonical taxonomy.json (a bundle or run directory is accepted too)."""
        path = Path(path)
        if path.is_dir():
            path = path / "taxonomy.json"
        return cls.from_json(json.loads(path.read_text(encoding="utf-8")))

    @classmethod
    def from_json(cls, value: Mapping[str, Any]) -> "Taxonomy":
        """Canonical form only; ``normalise_source`` handles other shapes (the labelling definition)."""
        version = value.get("version")
        derived = (value.get("derived") or {}).get(VALENCE_FIELD, {})
        if derived and derived.get("map") != FINE_TO_VALENCE:
            raise TaxonomyError("taxonomy.json valence map differs from nsb-v1")
        if value.get("gate_order") not in (None, list(GATE_ORDER)):
            raise TaxonomyError("taxonomy.json gate order differs from nsb-v1")
        return cls(value["fields"], value.get("absence"), version, value.get("class_weights"),
                   value.get("macro_f1_exclude"))

    def to_json(self) -> dict[str, Any]:
        return {
            "version": self.version, "fields": self.fields, "absence": self.absence,
            "derived": {VALENCE_FIELD: {"from": FINE_FIELD, "map": FINE_TO_VALENCE}},
            "relevant_option": RELEVANT, "gated_by_relevance": list(GATED_BY_RELEVANCE),
            "commercial_ad_options": list(self.commercial_ad), "product_fields": list(PRODUCT_FIELDS),
            "gate_order": list(GATE_ORDER), "field_weights": FIELD_WEIGHTS,
            "class_weights": self.class_weights, "macro_f1_exclude": self.macro_f1_exclude,
        }

    # ------------------------------------------------------------------ shapes
    def sizes(self) -> list[int]:
        return [len(self.fields[f]) for f in self.names]

    def flat_size(self) -> int:
        return sum(self.sizes())

    def weights(self) -> list[float]:
        return [FIELD_WEIGHTS.get(f, 1.0) for f in self.names]

    def class_weight_vectors(self) -> list[list[float] | None]:
        """Per field in head order: the option weights for the cross-entropy (1.0 when unlisted), None if unweighted."""
        return [[self.class_weights[f].get(o, 1.0) for o in self.fields[f]] if f in self.class_weights else None
                for f in self.names]

    def index(self, field: str, option: str | None) -> int:
        """Option index, or -1 for a masked target."""
        return -1 if option is None else self.fields[field].index(option)

    def encode(self, targets: Mapping[str, str | None], mask: Iterable[str] = ()) -> list[int]:
        """Label row for the loss: option index per field, -1 where masked or missing."""
        masked = set(mask)
        return [-1 if f in masked else self.index(f, targets.get(f)) for f in self.names]

    def flatten(self, logits: Mapping[str, list[float]]) -> list[float]:
        """Per-field logits -> one flat list in taxonomy order (``post_predictions.logits``)."""
        out: list[float] = []
        for f in self.names:
            row = logits[f]
            if len(row) != len(self.fields[f]):
                raise ValueError(f"{f}: {len(row)} logits for {len(self.fields[f])} options")
            out.extend(float(x) for x in row)
        return out

    def unflatten(self, flat: list[float]) -> dict[str, list[float]]:
        if len(flat) != self.flat_size():
            raise ValueError(f"{len(flat)} logits for {self.flat_size()} options")
        out, i = {}, 0
        for f in self.names:
            k = len(self.fields[f])
            out[f] = [float(x) for x in flat[i:i + k]]
            i += k
        return out


def normalise_source(value: Mapping[str, Any]) -> Taxonomy:
    """Build a Taxonomy from the labelling definition (``naver_psy/labelling/taxonomy.json``), whatever its shape.

    Accepted: ``fields`` as ``{name: [options]}`` or ``[{"name", "options", "absence"?}]``, options as strings or
    ``{"name"|"value"|"id": ...}`` objects; version under ``version`` or ``taxonomy_version``. The derived
    ``emotion_valence`` head is appended when the definition lists only the nine asked fields. ``class_weights`` and
    ``macro_f1_exclude`` are taken when present (else the version defaults).
    """
    version = value.get("version") or value.get("taxonomy_version") or value.get("id")
    raw = value.get("fields") if "fields" in value else value.get("heads")
    if raw is None:
        raise TaxonomyError("taxonomy definition has no fields")

    def opt(o: Any) -> str:
        if isinstance(o, str):
            return o
        for key in ("name", "value", "id", "option"):
            if isinstance(o, Mapping) and key in o:
                return str(o[key])
        raise TaxonomyError(f"cannot read option {o!r}")

    fields: dict[str, list[str]] = {}
    if isinstance(raw, Mapping):
        for name, spec in raw.items():
            options = spec.get("options") if isinstance(spec, Mapping) else spec
            fields[name] = [opt(o) for o in options]
    else:
        for spec in raw:
            fields[spec.get("name") or spec.get("field")] = [opt(o) for o in spec["options"]]
    if VALENCE_FIELD not in fields:
        fields[VALENCE_FIELD] = list(NSB_V1_FIELDS[VALENCE_FIELD])
    derived = value.get("derived") or {}
    if isinstance(derived, list):  # the labelling definition lists derived heads as [{name, from, map, ...}]
        derived = {d.get("name"): d for d in derived}
    vmap = value.get("emotion_valence_map") or value.get("valence_map") or (derived.get(VALENCE_FIELD) or {}).get("map")
    if vmap is not None and dict(vmap) != FINE_TO_VALENCE:
        raise TaxonomyError("valence map differs from nsb-v1")
    return Taxonomy(fields, NSB_V1_ABSENCE, version or VERSION, value.get("class_weights"),
                    value.get("macro_f1_exclude"))


# ---------------------------------------------------------------------- labels: rules, targets and masks
def valence_of(fine: str | None) -> str | None:
    return None if fine is None else FINE_TO_VALENCE[fine]


def violations(labels: Mapping[str, str | None]) -> list[str]:
    """Hard rules (labelling §2.2 rules 1, 2, 5) a label set breaks; a null value is not checked."""
    out = []
    rel = labels.get("relevance")
    if rel is not None and rel != RELEVANT:
        for f in GATED_BY_RELEVANCE:
            if labels.get(f) is not None and labels[f] != NSB_V1_ABSENCE[f]:
                out.append(f"not_market_{f}")
    if labels.get("target") == "none" and labels.get("direction") not in (None, "none"):
        out.append("target_none_direction")
    fine, val = labels.get(FINE_FIELD), labels.get(VALENCE_FIELD)
    if fine is not None and val is not None and FINE_TO_VALENCE[fine] != val:
        out.append("fine_valence")
    return out


def gate_fill(labels: Mapping[str, str | None]) -> dict[str, str | None]:
    """Apply the label-side gates: fields a gate decides are set to their absence; valence is derived."""
    t = dict(labels)
    if t.get("relevance") is not None and t["relevance"] != RELEVANT:
        for f in GATED_BY_RELEVANCE:
            t[f] = NSB_V1_ABSENCE[f]
    if t.get("target") == "none":
        t["direction"] = "none"
    if t.get(FINE_FIELD) is not None:
        t[VALENCE_FIELD] = FINE_TO_VALENCE[t[FINE_FIELD]]
    return t


def loss_mask(labels: Mapping[str, str | None], unresolved: Iterable[str] = ()) -> list[str]:
    """Fields excluded from the loss: gate-decided absences, plus unresolved (null) fields."""
    masked = set(unresolved)
    rel = labels.get("relevance")
    if rel is None:
        masked.update(GATED_BY_RELEVANCE)
    elif rel != RELEVANT:
        masked.update(GATED_BY_RELEVANCE)
    elif labels.get("target") in (None, "none"):
        masked.add("direction")
    masked.update(f for f, v in labels.items() if v is None)
    return sorted(f for f in masked if f in NSB_V1_FIELDS)


# ---------------------------------------------------------------------- distributions: gated output
def top(dist: Mapping[str, float]) -> str:
    """Highest-probability option; ties go to the earlier option in taxonomy order."""
    best = None
    for k, v in dist.items():
        if best is None or v > dist[best]:
            best = k
    if best is None:
        raise ValueError("empty distribution")
    return best


def _restrict(dist: Mapping[str, float], allowed: Iterable[str]) -> dict[str, float]:
    allowed = set(allowed)
    kept = {k: (v if k in allowed else 0.0) for k, v in dist.items()}
    total = sum(kept.values())
    if total <= 0:  # every allowed option had probability 0: uniform over the allowed ones
        return {k: (1.0 / len(allowed) if k in allowed else 0.0) for k in dist}
    return {k: v / total for k, v in kept.items()}


def _only(dist: Mapping[str, float], option: str) -> dict[str, float]:
    return {k: (1.0 if k == option else 0.0) for k in dist}


def gate(probs: Mapping[str, Mapping[str, float]]) -> dict[str, dict[str, float]]:
    """The labelling §2.2 gates applied in GATE_ORDER to per-field distributions (option order kept)."""
    out = {f: dict(d) for f, d in probs.items()}
    if top(out["relevance"]) != RELEVANT:
        for f in GATED_BY_RELEVANCE:
            out[f] = _only(out[f], NSB_V1_ABSENCE[f])
        return out
    if top(out["target"]) == "none":
        out["direction"] = _only(out["direction"], "none")
    val = top(out[VALENCE_FIELD])
    out[FINE_FIELD] = _restrict(out[FINE_FIELD], [o for o, v in FINE_TO_VALENCE.items() if v == val])
    return out


def argmax(probs: Mapping[str, Mapping[str, float]]) -> dict[str, str]:
    return {f: top(d) for f, d in probs.items()}