"""Label schema, handbook hard rules and the gated output (pure Python, no torch).

Shared by the local export (``classifier_export.py``), the pod training code and the local evaluation
(``classifier_eval.py``). The schema comes from ``taxonomy.json``, written by the export from the M2
protocol (handbook-v4-dev1); nothing about fields or options is hard-coded elsewhere.

Hard rules (handbook §3, §8, Part C; 0 violations in the 32k Astra labels):
- not relevant (``not_relevant`` or ``unclear``) -> asset_domain not_applicable, direction_primary none;
  relevant -> asset_domain is not not_applicable;
- not evaluable -> direction_primary none, emotion_valence none;
- emotion_valence none <-> emotion_fine not_expressed, and every fine emotion maps to one valence.

Gated output: the rules are applied in GATE_ORDER; forbidden options get probability 0 and the rest of the
field is renormalised. Valence leads fine emotion (the reported field, with more data).
"""
from __future__ import annotations

import json
from pathlib import Path

VERSION = 'taxonomy-v1'
NOT_RELEVANT = ('not_relevant', 'unclear')
FINE_TO_VALENCE = {
    'euphoric': 'positive', 'optimistic': 'positive', 'fomo': 'positive', 'calm': 'positive',
    'uncertain': 'negative', 'anxious': 'negative', 'panic_or_capitulation': 'negative',
    'regret_or_frustration': 'negative', 'angry': 'negative', 'mixed': 'mixed', 'not_expressed': 'none',
}
GATE_ORDER = ('market_relevance', 'asset_domain', 'view_evaluability', 'direction_primary', 'emotion_valence',
              'emotion_fine')


class Taxonomy:
    def __init__(self, fields: dict[str, list[str]], absence: dict[str, str], meta: dict | None = None):
        self.fields = {f: list(o) for f, o in fields.items()}
        self.names = list(self.fields)
        self.absence = dict(absence)
        self.meta = dict(meta or {})
        for f, options in self.fields.items():
            if len(set(options)) != len(options) or self.absence.get(f) not in options:
                raise ValueError('Bad field: ' + f)
        if set(self.fields.get('emotion_fine', ())) != set(FINE_TO_VALENCE):
            raise ValueError('emotion_fine options differ from the valence map')

    @classmethod
    def load(cls, path) -> 'Taxonomy':
        path = Path(path)
        if path.is_dir():
            path = path / 'taxonomy.json'
        value = json.loads(path.read_text(encoding='utf-8'))
        if value.get('version') != VERSION:
            raise ValueError('Unknown taxonomy version: %r' % value.get('version'))
        # The rules live in this module; a taxonomy.json that states different ones is refused, not ignored.
        if (value.get('not_relevant') != list(NOT_RELEVANT) or value.get('fine_to_valence') != FINE_TO_VALENCE
                or value.get('gate_order') != list(GATE_ORDER)):
            raise ValueError('taxonomy.json rules differ from classifier.taxonomy')
        return cls(value['fields'], value['absence'], {k: v for k, v in value.items() if k not in ('fields', 'absence')})

    def to_json(self) -> dict:
        return dict(self.meta, version=VERSION, fields=self.fields, absence=self.absence,
                    not_relevant=list(NOT_RELEVANT), fine_to_valence=FINE_TO_VALENCE, gate_order=list(GATE_ORDER))

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

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

    def encode(self, targets: dict) -> list[int]:
        return [self.index(f, targets.get(f)) for f in self.names]


def rule_fill(targets: dict) -> dict:
    """Fill nulls that a hard rule decides; everything else (including all non-null labels) is kept."""
    t = dict(targets)
    rel, ev, val, fine = (t.get('market_relevance'), t.get('view_evaluability'), t.get('emotion_valence'),
                          t.get('emotion_fine'))
    if rel in NOT_RELEVANT:
        if t.get('asset_domain') is None:
            t['asset_domain'] = 'not_applicable'
        if t.get('direction_primary') is None:
            t['direction_primary'] = 'none'
    if ev == 'not_evaluable' and t.get('direction_primary') is None:
        t['direction_primary'] = 'none'
    if val is None:
        implied = {'none'} if ev == 'not_evaluable' else set()
        if fine is not None:
            implied.add(FINE_TO_VALENCE[fine])
        if len(implied) == 1:  # decided by the rules; a conflict (not evaluable + an emotion) stays null
            t['emotion_valence'] = val = implied.pop()
    if fine is None and t.get('emotion_valence') in ('none', 'mixed'):  # the two valences with one fine option
        t['emotion_fine'] = 'not_expressed' if t['emotion_valence'] == 'none' else 'mixed'
    return t


def violations(labels: dict) -> list[str]:
    """Hard rules a complete label set breaks (a rule with a null in it is not checked)."""
    out = []
    rel, dom, ev = labels.get('market_relevance'), labels.get('asset_domain'), labels.get('view_evaluability')
    direction, val, fine = labels.get('direction_primary'), labels.get('emotion_valence'), labels.get('emotion_fine')
    if rel in NOT_RELEVANT and dom is not None and dom != 'not_applicable':
        out.append('not_relevant_domain')
    if rel == 'relevant' and dom == 'not_applicable':
        out.append('relevant_domain')
    if rel in NOT_RELEVANT and direction is not None and direction != 'none':
        out.append('not_relevant_direction')
    if ev == 'not_evaluable' and direction is not None and direction != 'none':
        out.append('not_evaluable_direction')
    if ev == 'not_evaluable' and val is not None and val != 'none':
        out.append('not_evaluable_valence')
    if val is not None and fine is not None and FINE_TO_VALENCE[fine] != val:
        out.append('fine_valence')
    return out


def _restrict(dist: dict[str, float], allowed) -> 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: dict[str, float], option: str) -> dict[str, float]:
    return {k: (1.0 if k == option else 0.0) for k in dist}


def top(dist: dict[str, float]) -> str:
    """Highest-probability option; ties go to the earlier option in taxonomy order (dicts keep it)."""
    best = None
    for k, v in dist.items():
        if best is None or v > dist[best]:
            best = k
    return best


def gate(probs: dict[str, dict[str, float]]) -> dict[str, dict[str, float]]:
    """Hard rules applied in GATE_ORDER to per-field distributions (taxonomy option order is kept).

    The steps below follow GATE_ORDER exactly: each decision uses the already-gated fields before it."""
    out = {f: dict(d) for f, d in probs.items()}
    rel = top(out['market_relevance'])
    dom = out['asset_domain']
    out['asset_domain'] = _only(dom, 'not_applicable') if rel in NOT_RELEVANT else \
        _restrict(dom, [o for o in dom if o != 'not_applicable'])
    ev = top(out['view_evaluability'])
    if rel in NOT_RELEVANT or ev == 'not_evaluable':
        out['direction_primary'] = _only(out['direction_primary'], 'none')
    if ev == 'not_evaluable':
        out['emotion_valence'] = _only(out['emotion_valence'], 'none')
    val = top(out['emotion_valence'])
    out['emotion_fine'] = _restrict(out['emotion_fine'], [o for o, v in FINE_TO_VALENCE.items() if v == val])
    return out


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