"""Request planning shared by the local runtime and the server: option-order averaging, the joint-read rule and
the pieces of self-gated thinking (gate question, think prompt, re-read with the thought)."""
from __future__ import annotations

import json
import math

from .protocol import DEFAULT_INSTRUCTIONS, describe, options

THINK_SYSTEM = "You reason briefly about decisions. Think in at most 60 words, only what decides the answer."
THOUGHT_NOTE = "\n\nReasoning notes (may contain mistakes): "
GATE_QUESTION = ("Would a careful person need to work this out step by step (calculation, logic puzzle, code tracing, "
                 "multi-hop reasoning) rather than just read, recognise or classify it?")

# thinking presets: "auto" = Decision Index setting (thinks only when unsure and the question needs reasoning),
# True = think on nearly every question
AUTO = {"tau": 0.6, "gate": 0.2, "alpha": 0.7, "cap": 96}
ALWAYS = {"tau": 0.97, "gate": 0.0, "alpha": 0.5, "cap": 320}


def orders(n: int, k: int) -> list[list[int]]:
    """Identity, reversed, then cyclic shifts of n options; the first k distinct orders."""
    out = [list(range(n)), list(range(n))[::-1]] + [[(i + s) % n for i in range(n)] for s in range(1, n)]
    seen, uniq = set(), []
    for o in out:
        if tuple(o) not in seen:
            seen.add(tuple(o))
            uniq.append(o)
    return uniq[:k]


def facet_joint(questions: dict) -> bool:
    """Read jointly when several questions apply the same criteria to different items with plain-text instructions
    (select-all candidates, aspect/sentiment pairs, per-entity labels); otherwise one question per request."""
    v = list(questions.values())
    return (len(v) > 1 and all(isinstance(q.get("instructions"), str) for q in v) and len({q.get("type") for q in v}) == 1
            and len({json.dumps(q.get("criteria"), sort_keys=True) for q in v}) == 1)


def permuted(q: dict, order: list[int]) -> dict:
    crit = q["criteria"]
    keys, vals = list(crit), list(crit.values())
    return {**q, "criteria": {keys[i]: vals[i] for i in order}}


def distribution(a: dict, keys: list[str]) -> list[float]:
    if a["type"] == "noul":
        return [a["noul"] if x == "true" else 1 - a["noul"] for x in keys]
    return [a["probabilities"][x] for x in keys]


def geo_mean(ps: list[list[float]]) -> list[float]:
    m = [math.exp(sum(math.log(max(p[i], 1e-9)) for p in ps) / len(ps)) for i in range(len(ps[0]))]
    z = sum(m)
    return [x / z for x in m]


def gate_request(state, q: dict) -> dict:
    """Cheap yes/no request: does question q need step-by-step reasoning? Uses a short state excerpt."""
    _, desc = options(q)
    text = ("Question: " + describe(q.get("instructions") or "Choose the best matching option.") + "\nOptions:\n" +
            "\n".join(str(x)[:200] for x in desc[:12]) + "\n\nContext excerpt:\n" + describe(state)[:1200])
    return {"state": text, "questions": {"g": {"type": "noul", "instructions": GATE_QUESTION}}}


def think_body(enc, request: dict) -> str:
    """User text of the think pass (the format the think adapter was trained on): state, one question, coded options."""
    (key, q), = request["questions"].items()
    _, desc = options(q)
    return ("State:\n" + enc._text(request["state"]) + "\n\nQuestion:\n" + enc._text(q.get("instructions") or DEFAULT_INSTRUCTIONS)
            + "\nOptions:\n" + "\n".join(f"{c}: {enc._text(t)}" for c, t in zip(enc.codes, desc)))


def with_thought(q: dict, thought: str) -> dict:
    q = dict(q)
    q["instructions"] = describe(q.get("instructions") or "Choose the best matching option.") + THOUGHT_NOTE + thought.strip()
    return q