from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import onnxruntime as ort
import torch
import torch.nn.functional as functional

from export_decoder_fp32 import (
    HEAD_DIM,
    HIDDEN_SIZE,
    NUM_KV_HEADS,
    NUM_LAYERS,
    OPSET,
    OUTPUT_DIR,
    REPORT_DIR,
    VISION_TOKENS,
    VOCAB_SIZE,
    DecoderBase,
    compare_outputs,
    flatten_cache,
    inspect_graph,
    load_model,
    make_past,
    output_names,
    past_names,
)
from export_decoder_qdq_int8 import rewrite_graph
from export_decoder_unified import empty_past, empty_vision


class DecoderUnifiedGather(DecoderBase):
    def forward(
        self,
        vision_embeds: torch.Tensor,
        token_ids: torch.Tensor,
        position_ids: torch.Tensor,
        past_k0: torch.Tensor,
        past_k1: torch.Tensor,
        past_k2: torch.Tensor,
        past_k3: torch.Tensor,
        past_k4: torch.Tensor,
        past_k5: torch.Tensor,
        past_v0: torch.Tensor,
        past_v1: torch.Tensor,
        past_v2: torch.Tensor,
        past_v3: torch.Tensor,
        past_v4: torch.Tensor,
        past_v5: torch.Tensor,
    ):
        keys = (past_k0, past_k1, past_k2, past_k3, past_k4, past_k5)
        values = (past_v0, past_v1, past_v2, past_v3, past_v4, past_v5)
        token_embed = functional.embedding(token_ids, self.lm_head.weight)
        inputs_embeds = torch.cat((vision_embeds, token_embed), dim=1)
        outputs = self.decoder(
            inputs_embeds=inputs_embeds,
            position_ids=position_ids,
            past_key_values=tuple(zip(keys, values)),
            use_cache=True,
            return_dict=True,
        )
        final_hidden_state = outputs.last_hidden_state[:, -1:, :]
        return (self.project_logits(final_hidden_state),) + flatten_cache(
            outputs.past_key_values
        )


def token(token_id: int) -> torch.Tensor:
    return torch.tensor([[token_id]], dtype=torch.int32)


def export_graph(
    wrapper: DecoderUnifiedGather,
    destination: Path,
    *,
    custom_opsets: dict[str, int] | None = None,
) -> None:
    torch.manual_seed(41)
    vision = torch.randn(1, VISION_TOKENS, HIDDEN_SIZE, dtype=torch.float32)
    past = make_past(1)
    positions = torch.arange(1, VISION_TOKENS + 2, dtype=torch.int32).unsqueeze(0)
    dynamic_axes = {
        "vision_embeds": {1: "vision_len"},
        "position_ids": {1: "input_len"},
        **{name: {2: "past_len"} for name in past_names()},
        **{name: {2: "total_len"} for name in output_names()[1:]},
    }
    torch.onnx.export(
        wrapper,
        (vision, token(1), positions, *past),
        destination,
        input_names=["vision_embeds", "token_ids", "position_ids", *past_names()],
        output_names=output_names(),
        dynamic_axes=dynamic_axes,
        opset_version=OPSET,
        do_constant_folding=True,
        dynamo=False,
        custom_opsets=custom_opsets,
    )


def feeds(values: tuple[torch.Tensor, ...]) -> dict[str, np.ndarray]:
    names = ["vision_embeds", "token_ids", "position_ids", *past_names()]
    return {name: value.numpy() for name, value in zip(names, values)}


def validate_mode(
    wrapper: DecoderUnifiedGather,
    session: ort.InferenceSession,
    label: str,
    values: tuple[torch.Tensor, ...],
) -> dict:
    with torch.inference_mode():
        expected = wrapper(*values)
    actual = session.run(None, feeds(values))
    return compare_outputs(label, expected, actual)


def validate_graph(wrapper: DecoderUnifiedGather, path: Path) -> dict:
    session = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"])
    torch.manual_seed(43)
    prefill = (
        torch.randn(1, VISION_TOKENS, HIDDEN_SIZE, dtype=torch.float32),
        token(1),
        torch.arange(VISION_TOKENS + 1, dtype=torch.int32).unsqueeze(0),
        *empty_past(),
    )
    step = (
        empty_vision(),
        token(4),
        torch.tensor([[VISION_TOKENS + 1]], dtype=torch.int32),
        *make_past(VISION_TOKENS + 1),
    )
    return {
        "prefill_zero_past": validate_mode(wrapper, session, "gather-prefill", prefill),
        "step_empty_vision": validate_mode(wrapper, session, "gather-step", step),
    }


def main() -> None:
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    REPORT_DIR.mkdir(parents=True, exist_ok=True)
    torch.manual_seed(0)
    torch.set_grad_enabled(False)
    wrapper = DecoderUnifiedGather(load_model()).eval()
    fp32_path = OUTPUT_DIR / "decoder_unified_gather_fp32.onnx"
    qdq_path = OUTPUT_DIR / "decoder_unified_gather_qdq_int8.onnx"
    print(f"Exporting {fp32_path}", flush=True)
    export_graph(wrapper, fp32_path)
    fp32_parity = validate_graph(wrapper, fp32_path)
    print(f"Quantizing {qdq_path}", flush=True)
    qdq_metadata = rewrite_graph(
        fp32_path,
        qdq_path,
        quantize_gather_shapes=frozenset({(VOCAB_SIZE, HIDDEN_SIZE)}),
    )
    qdq_parity = validate_graph(wrapper, qdq_path)
    if qdq_metadata["quantized_gathers"] != 1:
        raise RuntimeError(
            f"Expected one quantized token embedding Gather, got {qdq_metadata['quantized_gathers']}"
        )
    report = {
        "description": (
            "Complete six-layer Baberu decoder using an INT32 token ID and a quantized "
            "embedding Gather instead of a full-vocabulary one-hot MatMul"
        ),
        "fp32": {"graph": inspect_graph(fp32_path), "parity": fp32_parity},
        "qdq_int8": {
            "graph": inspect_graph(qdq_path),
            "quantization": qdq_metadata,
            "parity": qdq_parity,
        },
    }
    report_path = REPORT_DIR / "decoder-gather-report.json"
    report_path.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()