"""Export pinned Decision 3.0 text inference; vision export is separate."""

import argparse
import json
from pathlib import Path

import torch
import torch.nn.functional as F
from coreai_torch import TorchConverter, get_decomp_table
from safetensors.torch import load_file
from transformers import Qwen3_5Model
from transformers.models.qwen3_5.modeling_qwen3_5 import apply_rotary_pos_emb_vision
from transformers.vision_utils import (
    get_vision_interpolation_indices_and_weights,
    get_vision_position_ids,
)


class DecisionText(torch.nn.Module):
    def __init__(self, backbone, readout):
        super().__init__()
        self.language_model = backbone.language_model
        self.register_buffer("readout", readout.float())

    def forward(
        self, input_ids, attention_mask, position_ids, visual_embeddings, visual_mask
    ):
        embeddings = self.language_model.embed_tokens(input_ids)
        embeddings = torch.where(
            visual_mask, visual_embeddings.to(embeddings.dtype), embeddings
        )
        hidden = self.language_model(
            inputs_embeds=embeddings,
            attention_mask={
                "full_attention": attention_mask[:, None, None, :].bool(),
                "linear_attention": attention_mask,
            },
            position_ids=position_ids,
            use_cache=False,
        ).last_hidden_state
        return hidden[:, -1].float() @ self.readout.T


class DecisionVision(torch.nn.Module):
    """The unchanged vision weights at a fixed one-image 224px grid."""

    def __init__(self, visual, side=224):
        super().__init__()
        self.visual = visual
        grid = torch.tensor(
            [[1, side // visual.config.patch_size, side // visual.config.patch_size]]
        )
        indices, weights = get_vision_interpolation_indices_and_weights(
            grid,
            num_grid_per_side=visual.num_grid_per_side,
            mode=visual.interpolation_mode,
            align_corners=visual.interpolation_align_corners,
            spatial_merge_size=visual.config.spatial_merge_size,
            kwargs={},
        )
        self.register_buffer("indices", indices)
        self.register_buffer("weights", weights)
        self.register_buffer(
            "positions",
            get_vision_position_ids(grid, visual.spatial_merge_size, kwargs={}),
        )

    def forward(self, pixel_values):
        v = self.visual
        hidden = F.linear(
            pixel_values.to(v.patch_embed.proj.weight.dtype),
            v.patch_embed.proj.weight.flatten(1),
            v.patch_embed.proj.bias,
        )
        hidden = hidden + (v.pos_embed(self.indices) * self.weights[:, :, None]).sum(
            1
        ).to(hidden.dtype)
        cos, sin = v.rotary_pos_emb(hidden, self.positions)
        for block in v.blocks:
            x = block.norm1(hidden)
            attn = block.attn
            q, k, value = (
                attn.qkv(x)
                .reshape(x.shape[0], 3, attn.num_heads, -1)
                .permute(1, 0, 2, 3)
                .unbind(0)
            )
            q, k = apply_rotary_pos_emb_vision(q, k, cos, sin)
            y = F.scaled_dot_product_attention(
                q.transpose(0, 1).unsqueeze(0),
                k.transpose(0, 1).unsqueeze(0),
                value.transpose(0, 1).unsqueeze(0),
                scale=attn.scaling,
            )
            hidden = hidden + attn.proj(
                y.squeeze(0).transpose(0, 1).reshape(x.shape[0], -1)
            )
            hidden = hidden + block.mlp(block.norm2(hidden))
        return v.merger(hidden).float()


def save(model, sample, path):
    with torch.no_grad():
        exported = torch.export.export(model, sample, strict=False)
        print("Decomposing", flush=True)
        exported = exported.run_decompositions(get_decomp_table())
    print("Converting", flush=True)
    program = TorchConverter().add_exported_program(exported).to_coreai()
    program.save_asset(path)
    return exported.graph_signature.user_outputs[0]


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--source", default="upstream/d3-edge")
    p.add_argument("--output", default="bundles/d3-edge-CoreAI")
    p.add_argument("--length", type=int, default=512)
    p.add_argument("--precision", choices=["bfloat16", "float32"], default="bfloat16")
    a = p.parse_args()
    torch.set_num_threads(4)
    source = Path(a.source)
    output = Path(a.output)
    output.mkdir(parents=True, exist_ok=True)
    print("Loading backbone", flush=True)
    backbone = Qwen3_5Model.from_pretrained(
        source, dtype=getattr(torch, a.precision), attn_implementation="sdpa"
    ).eval()
    model = DecisionText(
        backbone, load_file(str(source / "readout.safetensors"))["weight"]
    ).eval()
    hidden = backbone.config.text_config.hidden_size
    sample = (
        torch.ones(1, a.length, dtype=torch.int32),
        torch.ones(1, a.length, dtype=torch.int32),
        torch.arange(a.length, dtype=torch.int32).view(1, 1, -1).expand(4, 1, -1),
        torch.zeros(1, a.length, hidden),
        torch.zeros(1, a.length, 1, dtype=torch.bool),
    )
    print("Exporting", flush=True)
    out = save(model, sample, output / "decision.aimodel")
    vision = DecisionVision(backbone.visual).eval()
    patches = (224 // backbone.config.vision_config.patch_size) ** 2
    width = (
        3
        * backbone.config.vision_config.temporal_patch_size
        * backbone.config.vision_config.patch_size**2
    )
    vision_out = save(vision, (torch.zeros(patches, width),), output / "vision.aimodel")
    (output / "manifest.json").write_text(
        json.dumps(
            {
                "source": source.name,
                "sequence_length": a.length,
                "hidden_size": hidden,
                "inputs": [
                    "input_ids",
                    "attention_mask",
                    "position_ids",
                    "visual_embeddings",
                    "visual_mask",
                ],
                "output": out,
                "vision_output": vision_out,
                "image_side": 224,
                "precision": a.precision,
            },
            indent=2,
        )
        + "\n"
    )
    print("Saved", flush=True)


if __name__ == "__main__":
    main()