"""Evidence selection for the Solomon layer (v1.1): the trained relevance head, with word overlap as a labelled fallback.

At evidence='support' (and the verification levels built on it) the layer's selector is:

  * TrainedEvidenceSelector, when an evidence head is configured. It reads what the answer ALREADY computed and
    makes no model call: the answer branch's own final hidden state (question branch, adapter active) and the
    document prefill's hidden states, mean-pooled over each sentence unit (solomon.units). The head scores every
    unit; a unit is selected iff max(unit probabilities) >= absent AND its probability >= select, with (select,
    absent) fitted per answer type on dev. Each candidate of a multi-label question has its own branch, so each
    label gets its own spans. method = 'trained_relevance_head'.
  * LexicalFallback otherwise (no head configured, an image document, a task the thresholds do not cover, or a
    state without prefix states). It is the v1 word-overlap selector, method = 'lexical_overlap_fallback', and the
    response says why the fallback ran.

RANKED POINTERS (the v1.1 serving policy, 'scope10-evidence-ranked-v1'). The post-hoc refit head (layer-42 MLP plus a
word-overlap residual: logit = head_logit + alpha * lexical, alpha from the head's own config) has no reliable
"no evidence" signal, so no threshold is applied. Each answer unit instead gets its TOP-K sentences (k = 3) by
probability, each with its score, as ranked pointers (method 'trained_relevance_head_ranked'). The answer's own
state is the absence signal: a unit whose collapsed four-state readout is 'not stated' (letter C), or a
reserved-option choice whose full readout resolves to the reserved 'does not state this' option, gets NO spans.

Neither establishes faithfulness: a relevance score says a sentence is where the answer's evidence is likely to be,
not that the answer was caused by it. `faithfulness_established` stays False.
"""
import hashlib
import json
import math
import re
from pathlib import Path

import numpy as np

METHOD = 'trained_relevance_head'
FALLBACK_METHOD = 'lexical_overlap_fallback'
VERIFICATION = 'model_relevance_only'
THRESHOLD_POLICY = 'scope10-evidence-thresholds-v1'
RANKED_POLICY = 'scope10-evidence-ranked-v1'
RANKED_METHOD = 'trained_relevance_head_ranked'
RANKED_SEMANTICS = ('ranked pointers: the top sentences by a trained relevance score (a post-hoc refit head plus a '
                    'word-overlap residual), highest first. Experimental. A score says where the evidence is likely to '
                    'be; it does not say the answer was caused by it (faithfulness not established), and it is not a '
                    'probability that the sentence is correct evidence. No spans are returned when the answer is not stated.')
NOUL_NOT_STATED = 2            # four-state letters A yes, B no, C not stated, D conflicting
ZERO_COST = {'calls': 0, 'input_tokens': 0, 'branches': 0}


def _sha(path):
    return hashlib.sha256(Path(path).read_bytes()).hexdigest()


def _digest(text):
    return hashlib.sha256(text.encode()).hexdigest()


def _forward():
    import solomon.evidence_head as EH
    return EH


def unit_token_spans(offsets, doc_start, units, limit):
    """Char units of the document -> [lo, hi) prompt-token spans clipped to the cached prefix (None if empty).
    Same rule as the trainer: a token belongs to a unit iff its char range overlaps the unit's and is non-empty."""
    offsets = np.asarray(offsets[:limit], dtype=np.int64).reshape(-1, 2)
    a, b = offsets[:, 0], offsets[:, 1]
    spans = []
    for s, e in units:
        lo_c, hi_c = doc_start + s, doc_start + e
        hit = np.nonzero((b > lo_c) & (a < hi_c) & (b > a))[0]
        spans.append((int(hit[0]), int(hit[-1]) + 1) if len(hit) else None)
    return spans


# The word-overlap residual. Byte-for-byte the feature the refit head was trained with
# (ops/scope10_evidence_refit.question_of / STOP / lexical); tests pin the equality.
STOP = set('the a an of to in is are was were be by for on and or with this that what which does do did it its as at from any document state states'.split())


def question_of(block):
    m = re.search(r'\nQuestion: (.*?)(?:\nOptions:|\nAnswer)', block, re.S)
    return m.group(1).strip() if m else block


def lexical(question, doc, units):
    """Word overlap with the question, stop words removed, document-IDF weighted, scaled to max 1 per row."""
    tokens = set(re.findall(r'\w+', question.casefold())) - STOP
    U = [set(re.findall(r'\w+', doc[s:e].casefold())) for s, e in units]
    df = {w: sum(w in u for u in U) for w in tokens}
    sc = np.array([sum(np.log((len(U) + 1) / (df[w] + .5)) for w in u & tokens) for u in U], np.float32)
    return sc / sc.max() if sc.max() > 0 else sc


def answer_not_stated(row):
    """True when the answer branch itself says the document does not state this: a four-state Noul whose top letter is
    C, or a reserved-option (|R) choice whose full readout (listed + 2 reserved) tops at the first reserved slot."""
    logits = np.asarray(row.get('letter_logits') if row.get('letter_logits') is not None else [], np.float64)
    if logits.ndim != 1 or not len(logits):
        return False
    key = row.get('head_key') or ''
    if key.endswith('/state4') and len(logits) == 4:
        return int(np.argmax(logits)) == NOUL_NOT_STATED
    if key.endswith('/choiceR') and len(logits) >= 3:
        return int(np.argmax(logits)) == len(logits) - 2
    return False


def keep_units(probs, select, absent):
    """The fitted rule: nothing is selected unless the best unit clears `absent`; then every unit >= `select`."""
    if not probs or max(probs) < absent:
        return []
    return [i for i, p in enumerate(probs) if p >= select]


class LexicalFallback:
    """The v1 word-overlap selector, labelled as the fallback it is."""
    task_aware = True

    def __init__(self, reason='no evidence head configured'):
        self.reason = reason
        self.last_reason = reason

    def describe(self):
        return {'method': FALLBACK_METHOD, 'reason': self.reason, 'faithfulness_established': False}

    def plan(self, text, questions, task=None):
        return {'calls': 0, 'text': text, 'questions': list(questions), 'task': task}

    def execute(self, plan, reason=None):
        from solomon.retrieval import lexical_select
        out = lexical_select(plan['text'], plan['questions'])
        self.last_reason = reason or self.reason
        return {**out, 'method': FALLBACK_METHOD, 'verification': 'retrieval_only', 'faithfulness_established': False,
                'cost': dict(ZERO_COST)}

    def __call__(self, text, questions):
        return self.execute(self.plan(text, questions))


class TrainedEvidenceSelector:
    """The trained evidence head as the layer's selector. `proxy` is the layer's engine proxy (its answer rows carry
    the branch hidden state and the warm document state they were asked on)."""
    task_aware = True

    def __init__(self, proxy, head, thresholds, *, head_sha256=None, thresholds_sha256=None, device=None):
        policy = thresholds.get('policy') if isinstance(thresholds, dict) else None
        tasks = (thresholds or {}).get('tasks') if isinstance(thresholds, dict) else None
        self.ranked = policy == RANKED_POLICY
        self.config = dict(getattr(head, 'config', {}) or {})
        if self.ranked:
            k = thresholds.get('top_k')
            if isinstance(k, bool) or not isinstance(k, int) or not 1 <= k <= 20:
                raise ValueError('ranked evidence policy needs an integer top_k in [1, 20]')
            if not isinstance(tasks, list) or not tasks or any(not isinstance(t, str) for t in tasks):
                raise ValueError('ranked evidence policy needs a nonempty list of answer types')
            if thresholds.get('suppress_when_answer_not_stated') is not True:
                raise ValueError('ranked evidence policy must suppress spans when the answer is not stated')
            residual = self.config.get('lexical_residual')
            if residual is not None:
                alpha = residual.get('alpha') if isinstance(residual, dict) else None
                if isinstance(alpha, bool) or not isinstance(alpha, (int, float)) or not math.isfinite(alpha):
                    raise ValueError('evidence head lexical_residual.alpha must be a finite number')
            self.top_k = k
            self.alpha = None if residual is None else float(residual['alpha'])
            self.proxy, self.head, self.device = proxy, head, device
            self.tasks = {t: {} for t in tasks}
            self.layer_index = self.config.get('layer_index') if self.config.get('layer') == 'mid' else None
            self.head_sha256, self.thresholds_sha256 = head_sha256, thresholds_sha256
            self.fallback = LexicalFallback('trained head not applicable')
            self.last_reason = None
            self.last_ranking = None
            return
        if policy != THRESHOLD_POLICY or not isinstance(tasks, dict) or not tasks:
            raise ValueError('evidence thresholds must be a ' + THRESHOLD_POLICY + ' or ' + RANKED_POLICY + ' policy')
        self.top_k, self.alpha, self.last_ranking = None, None, None
        for task, t in tasks.items():
            s, a = t.get('select'), t.get('absent')
            if any(isinstance(v, bool) or not isinstance(v, (int, float)) or not 0 <= v <= 1 for v in (s, a)) or s > a:
                raise ValueError(f'evidence thresholds for {task}: need 0 <= select <= absent <= 1')
        self.proxy, self.head, self.device = proxy, head, device
        self.tasks = {t: {'select': float(v['select']), 'absent': float(v['absent'])} for t, v in tasks.items()}
        self.layer_index = self.config.get('layer_index') if self.config.get('layer') == 'mid' else None
        self.head_sha256, self.thresholds_sha256 = head_sha256, thresholds_sha256
        self.fallback = LexicalFallback('trained head not applicable')
        self.last_reason = None

    @property
    def tap_layers(self):
        return () if self.layer_index is None else (int(self.layer_index),)

    def describe(self):
        if self.ranked:
            return {'method': RANKED_METHOD, 'policy': RANKED_POLICY, 'variant': self.config.get('variant', 'mlp'),
                    'arch': self.config.get('arch'), 'layer': self.config.get('layer', 'final'),
                    'layer_index': self.layer_index, 'lexical_residual_alpha': self.alpha,
                    'head_sha256': self.head_sha256, 'policy_sha256': self.thresholds_sha256, 'top_k': self.top_k,
                    'answer_types': sorted(self.tasks),
                    'rule': 'per answer unit: the top_k sentences by sigmoid(head_logit + alpha * lexical), with scores; '
                            'no spans when the answer unit is not stated (four-state C, or the reserved not-stated option)',
                    'fallback': FALLBACK_METHOD + ' (image documents, answer types outside the policy, missing prefix states)',
                    'faithfulness_established': False}
        return {'method': METHOD, 'variant': self.config.get('variant', 'mlp'), 'arch': self.config.get('arch'),
                'layer': self.config.get('layer', 'final'), 'head_sha256': self.head_sha256,
                'thresholds_sha256': self.thresholds_sha256, 'per_task_thresholds': dict(self.tasks),
                'rule': 'unit selected iff max(unit probabilities) >= absent and unit probability >= select',
                'fallback': FALLBACK_METHOD + ' (image documents, tasks without fitted thresholds, missing prefix states)',
                'faithfulness_established': False}

    def plan(self, text, questions, task=None):
        return {'calls': 0, 'text': text, 'questions': list(questions), 'task': task}

    def __call__(self, text, questions):
        return self.execute(self.plan(text, questions))

    def _fallback(self, plan, reason):
        self.last_reason = reason
        self.last_ranking = None
        return self.fallback.execute(plan, reason=reason)

    def execute(self, plan):
        text, questions, task = plan['text'], plan['questions'], plan.get('task')
        rows = list(getattr(self.proxy, 'rows', []))
        if task not in self.tasks:
            return self._fallback(plan, f'no fitted evidence thresholds for answer type {task}')
        if len(rows) != len(questions):
            return self._fallback(plan, 'answer branches do not match the question units')
        state = rows[0].get('state') or {}
        if not isinstance(state, dict) or state.get('doc_hidden') is None or not state.get('prompt'):
            return self._fallback(plan, 'the document state carries no prefix hidden states')
        marker = 'Document:\n' + text
        prompt = state['prompt']
        if marker not in prompt:
            return self._fallback(plan, 'document text not found in the rendered prompt')
        from solomon.units import units as sentence_units
        unit_list = sentence_units(text)
        if not unit_list:
            return self._fallback(plan, 'the document has no sentence units')
        tokenizer = self.proxy.t
        encoded = tokenizer(prompt, add_special_tokens=False, return_offsets_mapping=True)
        limit = int(state['prefix_tokens'])
        prefix_ids = state.get('prefix_ids')
        if prefix_ids is not None and list(encoded['input_ids'][:limit]) != [int(x) for x in prefix_ids[0].tolist()]:
            return self._fallback(plan, 'offset tokenisation differs from the cached prefix')
        spans = unit_token_spans(encoded['offset_mapping'], prompt.index(marker) + len('Document:\n'), unit_list, limit)
        if self.layer_index is None:
            doc = state['doc_hidden']
        else:
            doc = (state.get('doc_layers') or {}).get(str(self.layer_index))
            if doc is None:
                return self._fallback(plan, f'layer {self.layer_index} prefix states were not captured')
        EH = _forward()
        import torch
        thresholds = self.tasks[task]
        per_unit, source = [], _digest(text)
        self.last_ranking = None
        for i, row in enumerate(rows):
            branch = row.get('hidden') if self.layer_index is None else (row.get('taps') or {}).get(str(self.layer_index))
            if branch is None:
                return self._fallback(plan, 'the answer branch carries no hidden state for the evidence head')
            q = torch.as_tensor(np.asarray(branch, np.float32), device=doc.device)
            if self.ranked:
                per_unit.append({'unit': i, 'evidence': self._ranked(EH, torch, row, q, doc, text, unit_list, spans, source),
                                 'suppressed': answer_not_stated(row)})
                continue
            probs = EH.score_units_variant(self.head, q, doc, spans, missing=0.)
            chosen = keep_units(probs, thresholds['select'], thresholds['absent'])
            per_unit.append({'unit': i, 'evidence': [{'start': unit_list[j][0], 'end': unit_list[j][1],
                                                      'text': text[unit_list[j][0]:unit_list[j][1]], 'source_sha256': source,
                                                      'relevance': float(probs[j])} for j in chosen]})
        union = {}
        for u in per_unit:
            for s in u['evidence']:
                union[(s['start'], s['end'])] = {k: s[k] for k in ('start', 'end', 'text', 'source_sha256')}
        self.last_reason = None
        if self.ranked:
            # Ranked pointers keep their scores here; the chain below re-sorts spans by position and drops extra keys.
            self.last_ranking = [{'unit': u['unit'], 'suppressed': u.get('suppressed', False),
                                  'ranked': [dict(s) for s in u['evidence']]} for u in per_unit]
            per_unit = [{'unit': u['unit'], 'evidence': [{k: s[k] for k in ('start', 'end', 'text', 'source_sha256')}
                                                         for s in sorted(u['evidence'], key=lambda s: s['start'])]}
                        for u in per_unit]
        out = {'evidence': [union[k] for k in sorted(union)], 'units': per_unit, 'method': RANKED_METHOD if self.ranked else METHOD,
               'verification': VERIFICATION, 'faithfulness_established': False, 'cost': dict(ZERO_COST)}
        if self.head_sha256:
            out['relevance_weights_sha256'] = self.head_sha256
        return out

    def _ranked(self, EH, torch, row, q, doc, text, unit_list, spans, source):
        """Top-k sentences of one answer unit, each with its score; [] when the answer unit is not stated."""
        if answer_not_stated(row):
            return []
        kept = [j for j, sp in enumerate(spans) if sp is not None]
        if not kept:
            return []
        with torch.no_grad():
            logits = EH.unit_logits(self.head, q, doc, [spans[j] for j in kept]).float().cpu().numpy().astype(np.float64)
        if self.alpha is not None:
            lex = lexical(question_of(row.get('block') or ''), text, unit_list)
            logits = logits + self.alpha * np.asarray([lex[j] for j in kept], np.float64)
        probs = 1.0 / (1.0 + np.exp(-np.clip(logits, -500, 500)))
        order = sorted(range(len(kept)), key=lambda i: (-probs[i], kept[i]))[:self.top_k]
        return [{'start': unit_list[kept[i]][0], 'end': unit_list[kept[i]][1],
                 'text': text[unit_list[kept[i]][0]:unit_list[kept[i]][1]], 'source_sha256': source,
                 'score': float(probs[i]), 'rank': r + 1} for r, i in enumerate(order)]


def load_thresholds(path):
    raw = Path(path).read_bytes()
    return json.loads(raw), hashlib.sha256(raw).hexdigest()


def load_selector(proxy, head_prefix, thresholds_path, device='cuda'):
    """<prefix>.safetensors + <prefix>.json (the head) and the per-task thresholds JSON -> TrainedEvidenceSelector."""
    EH = _forward()
    head = EH.load(head_prefix, device=device)
    thresholds, thresholds_sha = load_thresholds(thresholds_path)
    return TrainedEvidenceSelector(proxy, head, thresholds, head_sha256=_sha(str(head_prefix) + '.safetensors'),
                                   thresholds_sha256=thresholds_sha, device=device)