"""Dinah-1 with two experts (code and calculator).

For every request the gate looks only at the content, the same way for every request:
  1. code expert: "what does this function return" questions are settled by running the function (`tools.py`);
  2. calculator expert: numeric multiple-choice word problems are solved by a calculator loop in which a Dinah
     trained on calculator decisions (`calculator/model.safetensors`, a second 150M encoder) picks each calculation (which number, which operation, when to
     stop), a beam search keeps several paths, and the option that collects the most probability among finished
     paths wins;
  3. everything else, and whatever the experts leave open, is answered by the Dinah-1 encoder at the root of
     this repository (PyTorch, or int8 ONNX on CPU).
An expert answers only what it settles; nothing is truncated and every option gets a probability.

    from experts import Dinah01
    model = Dinah01.from_pretrained("path/or/repo")
    model.predict(state, {"answer": {"type": "choice", "criteria": {...}, "instructions": "..."}})
"""
from __future__ import annotations

import copy
import math
import operator
import re
from collections import OrderedDict, defaultdict
from pathlib import Path
from typing import Any

import numpy as np

from dinah import QUESTION_TYPES, Dinah, DinahONNX, _resolve
from tools import TOOLS

# ---------------------------------------------------------------- calculator expert
NUM = re.compile(r"(?<![\w.])-?\d[\d,]*\.?\d*")
NUMERIC_OPTION = re.compile(r"^\s*\$?\s*-?\d[\d,]*(\.\d+)?\s*%?\s*$")
WORDS = {"half": 0.5, "twice": 2, "double": 2, "triple": 3, "one": 1, "two": 2, "three": 3, "four": 4, "five": 5, "six": 6,
         "seven": 7, "eight": 8, "nine": 9, "ten": 10, "eleven": 11, "twelve": 12, "dozen": 12, "quarter": 0.25,
         "hundred": 100, "thousand": 1000}
CONSTANTS = [0.5, 1, 2, 3, 4, 5, 7, 10, 12, 24, 30, 52, 60, 100, 365, 1000]
OPS = {"+": operator.add, "-": operator.sub, "×": operator.mul, "÷": operator.truediv}
I_FIRST = ("Solve the problem step by step with a calculator. Choose the first number of the next calculation, "
           "or choose done if the last result already answers the question.")
I_OPERATION = "Solve the problem step by step with a calculator. Choose the next calculation."
DONE = "done: the last result is the answer"
MAX_STEPS = 10


def _r6(x):
    return round(float(x), 6)


def _fmt(x):
    x = _r6(x)
    return str(int(x)) if x.is_integer() else f"{x:g}"


def _numbers(text):
    seen, out = set(), []
    def add(v):
        v = _r6(v)
        if v not in seen:
            seen.add(v); out.append(v)
    for m in NUM.finditer(text):
        try:
            add(float(m.group().replace(",", "").rstrip(".")))
        except ValueError:
            pass
    low = text.lower()
    for w, v in WORDS.items():
        if re.search(rf"\b{w}\b", low):
            add(v)
    for m in re.findall(r"(\d+(?:\.\d+)?)\s*%", text):
        add(float(m) / 100)
    return out


class _Scratchpad:
    """Known numbers (problem, constants, results) and the calculations so far, as the calculator expert was trained."""

    def __init__(self, problem):
        self.problem, self.vals, self.origin, self.steps, self.last = problem, [], [], [], None
        for v in _numbers(problem):
            self._add(v, "question")
        for c in CONSTANTS:
            self._add(c, "constant")

    def _add(self, v, origin):
        v = _r6(v)
        if v not in self.vals:
            self.vals.append(v); self.origin.append(origin)

    def apply(self, a, op, b):
        if op == "÷" and b == 0:
            return None
        r = _r6(OPS[op](a, b))
        self.steps.append(f"{_fmt(a)} {op} {_fmt(b)} = {_fmt(r)}")
        self._add(r, f"= {_fmt(a)} {op} {_fmt(b)}")
        self.last = r
        return r

    def state(self, first=None):
        s = {"question": self.problem, "calculations so far": list(self.steps) or ["none"]}
        if first is not None:
            s["first number"] = _fmt(first)
        return s

    def first_options(self):
        return [f"{_fmt(v)} ({o})" for v, o in zip(self.vals, self.origin)] + ([DONE] if self.last is not None else [])

    def operation_options(self, a):
        return [(f"{_fmt(a)} {s} {_fmt(b)}", s, b) for b in self.vals for s in ("+", "-", "×", "÷")]


class Calculator:
    def __init__(self, model: Dinah, beam=4, k_first=2, k_op=3, p_answer=0.9, cache=4096):
        self.model, self.beam, self.k_first, self.k_op, self.p = model, beam, k_first, k_op, p_answer
        self.cache, self.max_cache = OrderedDict(), cache

    def _logprobs(self, requests):
        """Raw log-probabilities (no temperature) for [(state, instructions, options)]."""
        enc = [(*self.model._encode(i, s, o), QUESTION_TYPES["choice"]) for s, i, o in requests]
        out = []
        for start in range(0, len(enc), 32):
            logits, _ = self.model._forward(self.model._batch(enc[start:start + 32]))
            for j, (_, pos, _) in enumerate(enc[start:start + 32]):
                z = logits[j, :len(pos)].astype(np.float64); z = z - z.max()
                out.append(z - math.log(np.exp(z).sum()))
        return out

    def hypotheses(self, problem):
        if problem in self.cache:
            self.cache.move_to_end(problem)
            return self.cache[problem]
        try:
            hs = self._search(problem)
        except ValueError:   # too many operation options for the context window: the expert abstains
            hs = []
        self.cache[problem] = hs
        if len(self.cache) > self.max_cache:
            self.cache.popitem(last=False)
        return hs

    def _search(self, problem):
        alive, done = [(0.0, _Scratchpad(problem))], []
        for _ in range(MAX_STEPS):
            if not alive:
                break
            opts1 = [sp.first_options() for _, sp in alive]
            l1 = self._logprobs([(sp.state(), I_FIRST, o) for (_, sp), o in zip(alive, opts1)])
            second = []
            for (lp, sp), o, l in zip(alive, opts1, l1):
                for c in np.argsort(-l)[:self.k_first]:
                    if o[c] == DONE:
                        done.append((lp + l[c], sp.last))
                    else:
                        second.append((lp + l[c], sp, sp.vals[c]))
            if not second:
                alive = []
                break
            opts2 = [sp.operation_options(a) for _, sp, a in second]
            l2 = self._logprobs([(sp.state(a), I_OPERATION, [t for t, _, _ in o]) for (_, sp, a), o in zip(second, opts2)])
            new = []
            for (lp, sp, a), o, l in zip(second, opts2, l2):
                for c in np.argsort(-l)[:self.k_op]:
                    _, s, b = o[c]
                    sp2 = copy.deepcopy(sp)
                    if sp2.apply(a, s, b) is not None:
                        new.append((lp + l[c], sp2))
            alive = sorted(new, key=lambda x: -x[0])[:self.beam]
        done += [(lp, sp.last) for lp, sp in alive]
        return sorted([(lp, r) for lp, r in done if r is not None], reverse=True)

    def __call__(self, state, questions):
        problem = state if isinstance(state, str) else (
            max((v for v in state.values() if isinstance(v, str) and re.search(r"\d", v)), key=len, default=None)
            if isinstance(state, dict) else None)
        if not problem or len(problem) < 40 or len(_numbers(problem)) < 2:
            return {}
        out = {}
        for key, q in questions.items():
            if q.get("type") != "choice":
                continue
            keys = list(q["criteria"])
            options = [q["criteria"][k] for k in keys]
            if len(options) < 3 or not all(isinstance(o, str) and NUMERIC_OPTION.match(o) for o in options):
                continue
            values = [float(o.replace("$", "").replace(",", "").replace("%", "").strip()) for o in options]
            mass = defaultdict(float)
            for lp, r in self.hypotheses(problem):
                hit = [i for i, v in enumerate(values) if abs(r - v) < 1e-4]
                if hit:
                    mass[hit[0]] += math.exp(lp)
            if not mass:
                continue
            best = max(mass, key=mass.get)
            rest = (1.0 - self.p) / (len(keys) - 1)
            out[key] = {"type": "choice", "choice": keys[best], "confidence": self.p,
                        "probabilities": {k: (self.p if i == best else rest) for i, k in enumerate(keys)}}
        return out


def _calculator(directory: Path, device: str | None, dtype: str | None) -> Dinah:
    """The calculator expert is stored as a rank-64 difference from the root encoder (calculator/delta.safetensors,
    fp16): W = W_root + A @ B for large matrices, W_root + delta for the rest. Rebuilt in memory at load time."""
    from safetensors.torch import load_file
    model = Dinah(directory, device, dtype)                      # starts from the root encoder
    full = directory / "calculator" / "model.safetensors"
    if full.exists():                                            # Dinah-1: the calculator is a complete model (no difference to rebuild)
        model.net.load_state_dict(load_file(str(full)), strict=True)
        return model
    delta = load_file(str(directory / "calculator" / "delta.safetensors"))
    state = {k: v.float().cpu() for k, v in model.net.state_dict().items()}
    for k in state:
        if k + ".A" in delta:
            state[k] = state[k] + delta[k + ".A"].float() @ delta[k + ".B"].float()
        elif k in delta:
            state[k] = state[k] + delta[k].float()
    model.net.load_state_dict(state, strict=True)
    return model


class Dinah01:
    FILES = ["*.json", "*.py", "model.safetensors", "onnx/model_int8.onnx", "calculator/*"]

    def __init__(self, directory: Path, device: str | None = None, dtype: str | None = None,
                 tools: str = "code,calculator", backend: str = "torch"):
        self.encoder = DinahONNX(directory) if backend == "onnx" else Dinah(directory, device, dtype)
        self.max_length = self.encoder.max_length
        names = [] if str(tools).lower() in ("", "0", "none", "off") else [t.strip() for t in str(tools).split(",")]
        self.experts = [Calculator(_calculator(directory, device, dtype)) if n == "calculator" else TOOLS[n] for n in names]
        self.tools, self.backend = names, backend

    @classmethod
    def from_pretrained(cls, path_or_repo: str | Path = ".", device: str | None = None, dtype: str | None = None,
                        revision: str | None = None, tools: str = "code,calculator", backend: str = "torch") -> "Dinah01":
        return cls(_resolve(path_or_repo, cls.FILES, revision), device, dtype, tools, backend)

    def predict(self, state: Any, questions: dict[str, dict[str, Any]], batch_size: int = 16) -> dict[str, dict[str, Any]]:
        """Answers {key: answer} for one request (a shared state and several typed questions)."""
        settled = {}
        for expert in self.experts:
            settled.update(expert(state, {k: q for k, q in questions.items() if k not in settled}))
        rest = [k for k in questions if k not in settled]
        out = dict(zip(rest, self.encoder.predict([{**questions[k], "state": state} for k in rest], batch_size))) if rest else {}
        out.update(settled)
        return {k: out[k] for k in questions}