"""Question rendering, calibrated answers and audio preprocessing."""

import json
import math
from typing import Literal

import numpy as np
import torch
from pydantic import BaseModel, ConfigDict, Field

MARKER = "\n<spev_score>\n"


class ChoiceQuestion(BaseModel):
    model_config = ConfigDict(extra="forbid", strict=True)
    type: Literal["choice"] = "choice"
    instructions: str
    criteria: dict[str, str] = Field(min_length=2, max_length=255)


class ChoiceQuestions(BaseModel):
    model_config = ConfigDict(extra="forbid", strict=True)
    questions: dict[str, ChoiceQuestion] = Field(min_length=1)


def literal(value):
    return (
        json.dumps(value, ensure_ascii=False)
        .replace("<", "\\u003c")
        .replace(">", "\\u003e")
    )


def render(state, questions):
    if not questions:
        raise ValueError("At least one question is required")
    schema = []
    mapping = []
    for key, question in questions.items():
        criteria = question["criteria"]
        if not 2 <= len(criteria) <= 255:
            raise ValueError("Each question requires 2–255 criteria")
        schema.append(
            {
                "instructions": question["instructions"],
                "criteria": list(criteria.values()),
            }
        )
        mapping.extend((key, name) for name in criteria)
    # IDs only route outputs. Descriptions carry the criterion semantics.
    text = "Use the audio and state as evidence for the supplied questions.\n"
    text += "State: " + literal(state) + "\nQuestions: " + literal(schema) + "\n"
    for question in questions.values():
        for description in question["criteria"].values():
            text += (
                literal(
                    {"question": question["instructions"], "criterion": description}
                )
                + MARKER
            )
    return text, mapping


def locate_markers(ids, marker, expected):
    positions = [
        i + len(marker) - 1
        for i in range(len(ids) - len(marker) + 1)
        if ids[i : i + len(marker)] == marker
    ]
    if len(positions) != expected:
        raise ValueError("Scoring anchor count mismatch")
    return positions


def answers(logits, mapping, temperature=1.0):
    if not math.isfinite(temperature) or temperature <= 0:
        raise ValueError("Temperature must be positive and finite")
    if not torch.isfinite(logits).all():
        raise ValueError("Nonfinite logits")
    grouped = {}
    for i, (question, criterion) in enumerate(mapping):
        grouped.setdefault(question, []).append((criterion, i))
    result = {}
    for key, options in grouped.items():
        probabilities = (
            (logits[[i for _, i in options]].double() / temperature).softmax(0).tolist()
        )
        p = dict(zip([name for name, _ in options], probabilities))
        choice = min(p, key=lambda name: (-p[name], name))
        result[key] = {
            "choice": choice,
            "probabilities": p,
            "confidence": (p[choice] - 1 / len(p)) / (1 - 1 / len(p)),
        }
    return result


def prepare_audio(audio, sample_rate):
    from scipy.signal import resample_poly

    wave = np.asarray(audio, dtype=np.float32)
    if wave.ndim == 2:
        wave = wave.mean(axis=1)
    if wave.ndim != 1 or not len(wave) or not np.isfinite(wave).all():
        raise ValueError(
            "Audio must be finite mono samples or frames-by-channels samples"
        )
    if (
        not isinstance(sample_rate, int)
        or isinstance(sample_rate, bool)
        or sample_rate <= 0
    ):
        raise ValueError("sample_rate must be a positive integer")
    if sample_rate != 16000:
        divisor = math.gcd(sample_rate, 16000)
        wave = resample_poly(wave, 16000 // divisor, sample_rate // divisor).astype(
            np.float32
        )
    return wave