"""Learned three-action baseline with explicit uncertainty abstention."""
from dataclasses import asdict
import json
from pathlib import Path
import re
import torch
from .actions import Action, Decision, Kind
from .features import encode, KINDS
from .model import PointerPolicy


class LearnedPolicy:
    def __init__(self, checkpoint, quantized=False, confidence_threshold=.85):
        root = Path(checkpoint)
        self.model = PointerPolicy.from_pretrained(root).eval()
        self.vocab = json.loads((root/'vocab.json').read_text(encoding='utf-8'))
        self.temperatures = json.loads((root/'calibration.json').read_text())['temperatures']
        self.threshold = confidence_threshold
        if quantized:
            self.model = torch.ao.quantization.quantize_dynamic(self.model,{torch.nn.Linear},dtype=torch.qint8)

    @torch.inference_mode()
    def predict(self, goal, state, ticket):
        row = dict(goal=goal,elements=[asdict(e) for e in state.elements])
        inputs,_,_,maps = encode([row],self.vocab)
        if not maps[0]:
            return Decision(ticket,Action(Kind.ASK_USER,('No eligible DOM target; a fallback is required.',)))
        a,t = self.model(*inputs)
        ap = (a/self.temperatures[0]).softmax(-1)[0]
        tp = (t/self.temperatures[1]).softmax(-1)[0]
        ai,ti = int(ap.argmax()),int(tp.argmax())
        ac,tc = float(ap[ai]),float(tp[ti])
        if min(ac,tc) < self.threshold:
            return Decision(ticket,Action(Kind.ASK_USER,('The policy is uncertain about this action or target.',)),ac,tc)
        kind = Kind(KINDS[ai])
        element = state.elements[maps[0][ti]]
        args = (element.ref,)
        if kind in {Kind.TYPE,Kind.SELECT}:
            # Literal copying, not inferred website/workflow logic. General span prediction is pending.
            literals = re.findall(r'"([^"\n]*)"',goal)
            if len(literals) != 1:
                return Decision(ticket,Action(Kind.ASK_USER,('This baseline needs one quoted value to copy.',)),ac,tc)
            args += (literals[0],)
        return Decision(ticket,Action(kind,args),ac,tc)