"""Run the fused Qwen3.5/OmniJev 5-bit model with its decision head."""
from __future__ import annotations

import argparse
import json
import statistics
import time
from pathlib import Path

import mlx.core as mx
import numpy as np
from PIL import Image


def _load_runtime(model_path: Path):
    from mlx_vlm import load

    mx.set_cache_limit(256 * 1024**2)
    model, processor = load(str(model_path))
    model.set_dtype(mx.float32)
    if model.config.model_type != "qwen3_5":
        raise ValueError(f"Expected qwen3_5 model, got {model.config.model_type}")
    return Qwen35Adapter(model, processor)


class Qwen35Adapter:
    def __init__(self, model, processor):
        self.model = model
        self.processor = processor
        self.tokenizer = processor.tokenizer

    def execution_context(self):
        from mlx_vlm.generate.common import wired_limit

        return wired_limit(self.model)

    def prepare(self, prompt: str, image: Image.Image):
        from mlx_vlm.utils import prepare_inputs

        inputs = prepare_inputs(self.processor, images=[image], prompts=prompt)
        if inputs["input_ids"].shape[-1] > 6000:
            raise ValueError("image and prompt exceed 6000 input tokens")
        return inputs

    def prefill(self, inputs):
        ids = inputs["input_ids"]
        features = self.model.get_input_embeddings(
            ids,
            pixel_values=inputs.get("pixel_values"),
            image_grid_thw=inputs.get("image_grid_thw"),
            mask=inputs.get("attention_mask"),
        )
        mx.eval(features.inputs_embeds, features.position_ids)
        cache = self.model.language_model.make_cache()
        out = self.model.language_model(
            ids,
            inputs_embeds=features.inputs_embeds,
            cache=cache,
            position_ids=features.position_ids,
            rope_deltas=features.rope_deltas,
            skip_logits=True,
            return_hidden=True,
        )
        mx.eval(out.hidden_states[-1], *[v for c in cache for v in c.state if v is not None])
        return cache, int(features.position_ids.max().item()) + 1

    def fork(self, cache, count):
        from mlx_vlm.models.cache import ArraysCache, KVCache

        result = []
        for entry in cache:
            if type(entry) not in (ArraysCache, KVCache):
                raise TypeError(f"unsupported cache: {type(entry).__name__}")
            state = [None if x is None else mx.repeat(x, count, axis=0) for x in entry.state]
            result.append(type(entry).from_state(state, entry.meta_state))
        mx.eval(*[v for c in result for v in c.state if v is not None])
        return result

    def suffix_hidden(self, cache, next_position, rows, picks):
        branches = self.fork(cache, len(rows))
        length = max(map(len, rows))
        pad = self.tokenizer.pad_token_id or self.tokenizer.eos_token_id
        ids = mx.array([row + [pad] * (length - len(row)) for row in rows])
        positions = mx.broadcast_to(
            (mx.arange(length) + next_position)[None, None, :],
            (3, len(rows), length),
        )
        out = self.model.language_model(
            ids,
            cache=branches,
            position_ids=positions,
            skip_logits=True,
            return_hidden=True,
        )
        hidden = out.hidden_states[-1]
        return [hidden[i, mx.array(row)] for i, row in enumerate(picks)]


class OmniHead:
    def __init__(self, path: Path):
        with np.load(path) as values:
            self.weights = {key: value for key, value in values.items()}
        metadata_path = path.with_name("head_meta.json")
        metadata = json.loads(metadata_path.read_text()) if metadata_path.exists() else {}
        self.choice_temperature = float(metadata.get("temperatures", {}).get("choice", 1.0))

    @staticmethod
    def _linear(x, weight, bias):
        return x @ mx.array(weight.T) + mx.array(bias)

    @staticmethod
    def _gelu(x):
        return 0.5 * x * (1 + mx.tanh(0.7978845608 * (x + 0.044715 * x * x * x)))

    def logits(self, option_hidden, question_hidden, features):
        w = self.weights
        option = self._linear(
            self._gelu(self._linear(option_hidden, w["opt__0__weight"], w["opt__0__bias"])),
            w["opt__2__weight"], w["opt__2__bias"],
        )
        question = self._linear(
            self._gelu(self._linear(question_hidden, w["que__0__weight"], w["que__0__bias"])),
            w["que__2__weight"], w["que__2__bias"],
        )
        gate = mx.tanh(question)
        fused = option * gate[None, :] + self._linear(
            features, w["feat__weight"], w["feat__bias"]
        )
        tau = mx.exp(mx.array(w["log_tau"])[1])
        score = self._linear(fused, w["score__weight"], w["score__bias"]).squeeze(-1)
        score = score + self._linear(
            features, w["feat_lin__weight"], w["feat_lin__bias"]
        ).squeeze(-1)
        score = score / tau
        abstain = self._linear(
            gate, w["abstain__weight"], w["abstain__bias"]
        ).squeeze(-1) / tau
        return mx.concatenate([score, mx.reshape(abstain, (1,))])


def _marker_positions(row, opening, closing):
    def find(pattern):
        width = len(pattern)
        return [i for i in range(len(row) - width + 1) if row[i : i + width] == pattern]

    opens, closes = find(opening), find(closing)
    if len(opens) != 1 or len(closes) != 1:
        raise ValueError(f"option marker count mismatch: opens={opens}, closes={closes}")
    option_start, close_start = opens[0], closes[0]
    close_end = close_start + len(closing) - 1
    inner = list(range(option_start + len(opening), close_start))
    return option_start, close_end, inner


def predict_image(adapter, head, image: Image.Image, question: str, options: list[str]):
    if not 2 <= len(options) <= 256 or len(set(options)) != len(options):
        raise ValueError("supply 2–256 distinct options")
    started = time.perf_counter()
    tokenizer = adapter.tokenizer
    opening = tokenizer("<|opt|>", add_special_tokens=False)["input_ids"]
    closing = tokenizer("<|/opt|>", add_special_tokens=False)["input_ids"]
    prompt = adapter.processor.apply_chat_template(
        [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": question}]}],
        tokenize=False,
        add_generation_prompt=True,
    )
    full_prompt = prompt + "".join(f"<|opt|>{option}<|/opt|>" for option in options)
    inputs = adapter.prepare(full_prompt, image)
    full_ids = np.asarray(inputs["input_ids"])[0].tolist()
    open_at = [
        i for i in range(len(full_ids) - len(opening) + 1)
        if full_ids[i : i + len(opening)] == opening
    ]
    if len(open_at) != len(options):
        raise ValueError(f"expanded option markers not found: {open_at}")
    prefix_end = open_at[0]
    prefix = dict(inputs)
    prefix["input_ids"] = inputs["input_ids"][:, :prefix_end]
    prefix["attention_mask"] = mx.ones_like(prefix["input_ids"])
    cache, next_position = adapter.prefill(prefix)

    all_closes = [
        i for i in range(len(full_ids) - len(closing) + 1)
        if full_ids[i : i + len(closing)] == closing
    ]
    rows, picks, option_inner = [], [], []
    for index, option_start in enumerate(open_at):
        close_start = all_closes[index]
        close_end = close_start + len(closing) - 1
        row = full_ids[prefix_end:open_at[0]] + full_ids[option_start : close_end + 1]
        marker_start, close_pos, inner = _marker_positions(row, opening, closing)
        rows.append(row)
        picks.append([close_pos, max(0, marker_start - 1)] + [position - 1 for position in inner])
        option_inner.append(inner)

    hidden_rows = adapter.suffix_hidden(cache, next_position, rows, picks)
    all_hidden = mx.concatenate(hidden_rows, axis=0)
    language_model = adapter.model.language_model
    logits = (
        language_model.model.embed_tokens.as_linear(all_hidden)
        if language_model.args.tie_word_embeddings
        else language_model.lm_head(all_hidden)
    ).astype(mx.float32)
    scored, offset = [], 0
    for row_tokens, row_hidden, inner in zip(rows, hidden_rows, option_inner):
        count = len(inner)
        token_logits = logits[offset + 2 : offset + 2 + count]
        targets = mx.array([row_tokens[position] for position in inner])
        log_probs = token_logits - mx.logsumexp(token_logits, axis=-1, keepdims=True)
        values = log_probs[mx.arange(count), targets] if count else mx.array([])
        scored.append((row_hidden[0], row_hidden[1], values, count))
        offset += len(row_hidden)
    sums = mx.array([float(mx.sum(row[2]).item()) for row in scored])
    counts = mx.array([row[3] for row in scored])
    first = mx.array([float((row[2][0] if row[3] else 0).item()) for row in scored])
    normalized = sums - mx.logsumexp(sums)
    features = mx.stack(
        [sums / 10, sums / mx.maximum(counts, 1), counts / 10, normalized, first, mx.ones_like(sums)],
        axis=1,
    )
    option_hidden = mx.stack([row[0] for row in scored]).astype(mx.float32)
    question_hidden = scored[0][1].astype(mx.float32)
    scores = head.logits(option_hidden, question_hidden, features)
    probabilities = mx.softmax(scores / head.choice_temperature)
    mx.eval(probabilities)
    values = probabilities.tolist()
    choice_index = int(mx.argmax(probabilities[:-1]).item())
    return {
        "prediction": options[choice_index],
        "prediction_index": choice_index,
        "probabilities": values[:-1],
        "abstain_probability": values[-1],
        "elapsed_ms": (time.perf_counter() - started) * 1000,
        "quantization_bits": 5,
    }


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", type=Path, default=Path("."))
    parser.add_argument("--head", type=Path, default=Path("decision_head/head.npz"))
    parser.add_argument("--image", type=Path, required=True)
    parser.add_argument("--state", default="")
    parser.add_argument("--question", required=True)
    parser.add_argument("--options", nargs="+", required=True)
    parser.add_argument("--max-pixels", type=int, default=87808)
    parser.add_argument("--repeat", type=int, default=1)
    args = parser.parse_args()
    if args.repeat < 1:
        parser.error("repeat must be positive")
    model_path = args.model.resolve()
    head_path = args.head if args.head.is_absolute() else model_path / args.head
    adapter = _load_runtime(model_path)
    adapter.processor.image_processor.max_pixels = args.max_pixels
    head = OmniHead(head_path)
    with Image.open(args.image) as opened:
        image = opened.convert("RGB")
    question = args.question if not args.state else f"{args.question}\nContext: {args.state}"
    results = []
    with adapter.execution_context():
        for _ in range(args.repeat):
            results.append(predict_image(adapter, head, image, question, args.options))
    output = {"model": str(model_path), "requests": results}
    if len(results) > 1:
        output["warm_median_ms"] = statistics.median(result["elapsed_ms"] for result in results[1:])
    print(json.dumps(output, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()