"""Calibrated typed decisions over audio, images or text by Abacus.AI: one forward pass, a probability per option, no generated text.

    from smaug_decide import SmaugDecide
    model = SmaugDecide.from_pretrained('abacusai/Smaug-Decisions')
    model('clip.wav', 'Has the speaker finished their turn?')                      # -> {'Yes': 0.93, 'No': 0.07}
    model('clip.wav', 'Which genre is playing?', ['rock', 'jazz', 'classical'])    # -> one-answer: probabilities sum to 1
    model('clip.wav', 'Which sounds are present?', ['speech', 'music', 'alarm'], multi=True)   # -> select-all: one probability each
    model('photo.jpg', 'Is there a person in the picture?')                        # images and plain text work the same way
    model('The order never arrived.', 'How upset is the customer?', ['calm', 'annoyed', 'angry'])

The readout is the model's own LM-head logits for the option letters A..T at the answer position (softcapped), softmax for one-answer
questions and sigmoid(z - b) for select-all questions; `b` is the learned select-all threshold stored in readout.json / head.pt.
Runs on NVIDIA GPUs (bfloat16), Apple Silicon (MPS) and CPU; the device is picked automatically.
"""
import json
import os

import numpy as np
import soundfile as sf
import torch
from huggingface_hub import snapshot_download
from PIL import Image
from transformers import AutoModelForMultimodalLM, AutoProcessor


LETTERS = 'ABCDEFGHIJKLMNOPQRST'
SAMPLE_RATE, MAX_AUDIO_SECONDS = 16000, 30.0            # the model hears 16 kHz mono; longer clips keep their last 30 s
MAX_TOKENS, AUDIO_TOKENS_PER_SECOND, IMAGE_TOKENS = 3072, 25, 560   # images at 560 soft tokens, as trained and evaluated
PREAMBLE = ('You make decisions for apps. Read the rules, the context and the input, then answer the question by choosing '
            'from the lettered options.')
IMAGE_SUFFIXES = ('.jpg', '.jpeg', '.png', '.webp', '.gif', '.bmp')
AUDIO_SUFFIXES = ('.wav', '.flac', '.ogg', '.mp3', '.m4a', '.aiff', '.aif')
WEIGHT_FILES = ['*.json', '*.jinja', 'model*.safetensors', 'head.pt']   # the checkpoint at the repo root


def trim(s, n):
    if len(s) <= n:
        return s
    return s[:int(n * 0.7)] + ' ... ' + s[-int(n * 0.3):]


def render(question, options, descriptions, multi, media_tag, rules, context, text_input, char_budget):
    """The training prompt, verbatim: preamble, optional rules / context, the input, the lettered options, the answer instruction."""
    parts = [PREAMBLE]
    if rules:
        parts.append(f'Rules: {trim(rules, char_budget // 2)}')
    if context:
        parts.append(f'Context: {trim(context, char_budget)}')
    if media_tag:
        parts.append(f'Input: {media_tag}')
    elif text_input:
        parts.append(f'Input: {trim(text_input, char_budget)}')
    lines = [f'Question: {question}' + (' Select all that apply.' if multi else '')]
    for k, (opt, desc) in enumerate(zip(options, descriptions)):
        lines.append(f'{LETTERS[k]}) {opt}' + (f' — {desc}' if desc else ''))
    parts.append('\n'.join(lines))
    parts.append('Answer with the letters of all options that apply.' if multi else 'Answer with the letter of the best option.')
    return '<bos><|turn>user\n' + '\n'.join(parts) + '<turn|>\n<|turn>model\n'


def load_audio(x):
    """A path, or a (samples, sample_rate) pair, or a 1-D float array at 16 kHz -> mono float32 at 16 kHz, last 30 s."""
    if isinstance(x, str):
        w, sr = sf.read(x, dtype='float32')
    elif isinstance(x, tuple):
        w, sr = x
    else:
        w, sr = x, SAMPLE_RATE
    w = np.asarray(w, dtype=np.float32)
    if w.ndim > 1:
        w = w.mean(1)
    if sr != SAMPLE_RATE:   # linear resampling is enough for a reference script; use a proper resampler for production audio
        n = int(round(len(w) * SAMPLE_RATE / sr))
        w = np.interp(np.linspace(0.0, len(w) - 1, n), np.arange(len(w)), w).astype(np.float32)
    return w[-int(MAX_AUDIO_SECONDS * SAMPLE_RATE):]


class SmaugDecide:
    def __init__(self, model, processor, readout, device):
        self.model, self.proc, self.device = model.eval(), processor, device
        tok = processor.tokenizer
        self.letter_ids = torch.tensor([tok(L, add_special_tokens=False)['input_ids'][0] for L in LETTERS], device=device)
        self.cap, self.b = readout.get('final_logit_softcapping'), float(readout['multi_bias_b'])

    @classmethod
    def from_pretrained(cls, name_or_path, device=None):
        path = name_or_path if os.path.isdir(name_or_path) else snapshot_download(name_or_path, allow_patterns=WEIGHT_FILES)
        device = device or ('cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu')
        model = AutoModelForMultimodalLM.from_pretrained(path, dtype=torch.bfloat16).to(device)
        readout = json.load(open(os.path.join(path, 'readout.json')))
        return cls(model, AutoProcessor.from_pretrained(path), readout, device)

    @torch.no_grad()
    def __call__(self, inp, question, options=('Yes', 'No'), descriptions=None, multi=False, rules=None, context=None):
        """inp: an audio path / (samples, sr) / 16 kHz array, an image path / PIL image, or a plain string. Returns {option: probability}."""
        options = list(options)
        assert 2 <= len(options) <= len(LETTERS), f'2 to {len(LETTERS)} options'
        descriptions = list(descriptions) if descriptions else [None] * len(options)
        audio = image = text_input = None
        if isinstance(inp, Image.Image) or (isinstance(inp, str) and inp.lower().endswith(IMAGE_SUFFIXES)):
            image = (inp if isinstance(inp, Image.Image) else Image.open(inp)).convert('RGB')
        elif isinstance(inp, str) and not inp.lower().endswith(AUDIO_SUFFIXES) and not os.path.exists(inp):
            text_input = inp
        else:
            audio = load_audio(inp)
        media_tokens = len(audio) / SAMPLE_RATE * AUDIO_TOKENS_PER_SECOND + 2 if audio is not None else (IMAGE_TOKENS + 2 if image is not None else 0)
        tag = '<|audio|>' if audio is not None else '<|image|>' if image is not None else None
        budget = int(max(400, (MAX_TOKENS - media_tokens - 80 - 10 * len(options)) * 3.5))
        text = render(question, options, descriptions, multi, tag, rules, context, text_input, budget)
        kw = {}
        if audio is not None:
            kw['audio'] = [audio]
        if image is not None:
            kw['images'] = [image]; kw['images_kwargs'] = {'max_soft_tokens': IMAGE_TOKENS}
        enc = self.proc(text=[text], return_tensors='pt', add_special_tokens=False, **kw).to(self.device)
        h = self.model.model(**enc, use_cache=False).last_hidden_state[0, -1].float()          # hidden state at the answer position
        z = h @ self.model.lm_head.weight[self.letter_ids[:len(options)]].float().T          # LM-head logits of the letters only
        if self.cap:
            z = torch.tanh(z / self.cap) * self.cap
        p = torch.sigmoid(z - self.b) if multi else torch.softmax(z, -1)
        return dict(zip(options, p.tolist()))