"""Saif v2 DI adapter. Run from the source checkout with scripts on PYTHONPATH.

saif08-core releases: identical to the v3 fast-path engine except for one layout-triggered calibration. If the package
carries saif_phishing_calibration.json with status "fitted" (fitted on the train-derived dev_phishing.jsonl only), a
request whose state is a dict with keys exactly PHISHING_STATE and that has a choice field with criteria keys exactly
{phishing, legitimate} (and, if present, an `is_phishing` noul twin with criteria keys {false, true}) gets, on each such
choice field only:  p(phishing) = sigmoid((z_phishing - z_legitimate) / T + b),  p(legitimate) = 1 - p(phishing),
with z the raw label logits (the per-type temperature is not applied to these fields; T and b were fitted on raw logits).
Every other field and every other request goes through the unchanged v3 path. Nothing keys on ids, names or text.
"""
import json
import math
from pathlib import Path

import torch
from decision_index.engines.base import Engine, Unsupported
from saif_v2_core import load
import singlepass
singlepass.install()


PHISHING_STATE = frozenset({'sender', 'from', 'subject', 'body', 'link_display_text', 'link_url'})
PHISHING_KEYS = frozenset({'phishing', 'legitimate'})


def criteria_keys(question):
    criteria = question.get('criteria')
    return frozenset(criteria) if isinstance(criteria, dict) else None


def phishing_fields(state, questions):
    """Names of the choice fields the calibration applies to, or () when the request does not match the layout."""
    if not isinstance(state, dict) or frozenset(state) != PHISHING_STATE:
        return ()
    twin = questions.get('is_phishing')
    if twin is not None and (twin.get('type') != 'noul' or criteria_keys(twin) not in (None, frozenset({'false', 'true'}))):
        return ()
    return tuple(name for name, question in questions.items()
                 if question.get('type') == 'choice' and criteria_keys(question) == PHISHING_KEYS)


def load_phishing(path):
    calibration = Path(path) / 'saif_phishing_calibration.json'
    if not calibration.is_file():
        return None
    values = json.loads(calibration.read_text())
    if values.get('status') != 'fitted':
        return None
    temperature, bias = float(values['temperature']), float(values['bias'])
    if not (math.isfinite(temperature) and temperature > 0 and math.isfinite(bias)):
        raise ValueError(f'invalid PhishNChips-layout calibration: {values}')
    return {'temperature': temperature, 'bias': bias, 'dev_sha256': values.get('dev', {}).get('sha256')}


class SaifEngine(Engine):
    name = 'saif-v2'

    def __init__(self, model=None, device='cuda', threads=4, read_twice=None, reading=None, **options):
        super().__init__(**options)
        if options:
            raise ValueError(f'Unknown options: {sorted(options)}')
        torch.set_num_threads(int(threads))
        self.device = torch.device(device)
        self.model = load(model or Path(__file__).parent, device)
        self.phishing = load_phishing(model or Path(__file__).parent)
        if reading is not None:
            if reading not in ('joint', 'per-question'):
                raise ValueError('reading must be joint or per-question')
            self.model.reading = reading
        if read_twice is not None:
            if read_twice not in (True, False, 'true', 'false'):
                raise ValueError('read_twice must be true or false')
            self.model.read_twice = read_twice in (True, 'true')
        self.provenance = {'K': 1, 'prefix_caching': False, 'truncation': False,
                           'attention': self.model.attention, 'reading': self.model.reading,
                           'read_twice': self.model.read_twice,
                           'serving_path': 'shared-prefix-segmented-bf16',
                           'base': self.model.base, 'revision': self.model.revision,
                           'phishing_layout_calibration': self.phishing}

    def decide(self, state, questions):
        fields = phishing_fields(state, questions) if self.phishing else ()
        if not fields:
            return self.model.decide(state, questions)
        return self.decide_phishing(state, questions, set(fields))

    @torch.no_grad()
    def decide_phishing(self, state, questions, fields):
        model = self.model
        model.model.eval()
        result = {}
        for name, keys, logits in model.logits(state, questions):
            if name in fields:
                index = {key: position for position, key in enumerate(keys)}
                margin = logits[index['phishing']] - logits[index['legitimate']]
                phishing = torch.sigmoid(margin / self.phishing['temperature'] + self.phishing['bias'])
                values = {'phishing': float(phishing), 'legitimate': float(1 - phishing)}
                result[name] = {key: values[key] for key in keys}
            else:
                temperature = model.temperatures.get(questions[name]['type'], 1.)
                result[name] = dict(zip(keys, (logits / temperature).softmax(-1).cpu().tolist()))
        return result

    def __call__(self, state, questions, **inputs):
        if inputs:
            raise Unsupported('Text only')
        try:
            probabilities = self.decide(state, questions)
            answers = {}
            for name, distribution in probabilities.items():
                kind = questions[name]['type']
                answer = {'type': kind, 'probabilities': distribution}
                if kind == 'noul':
                    answer['noul'] = distribution['true']
                elif kind == 'choice':
                    answer['choice'] = max(distribution, key=distribution.get)
                else:
                    answer['score'] = sum(int(key) * value for key, value in distribution.items())
                answers[name] = answer
            return {'model': self.name, 'answers': answers}, None
        except ValueError as error:
            if 'exceeds context' in str(error):
                raise Unsupported(str(error)) from error
            raise

    def synchronize(self):
        if self.device.type == 'cuda':
            torch.cuda.synchronize(self.device)

    def runtime(self):
        return {'torch': torch.__version__, 'device': str(self.device), 'K': 1}

    def close(self):
        del self.model