import json
from pathlib import Path

import torch
import torch.nn.functional as F
from coreai_torch import TorchConverter, get_decomp_table
from transformers import AutoModel, AutoProcessor


def unfold(x, dimension, size, step):
    dimension %= x.ndim
    return torch.stack(
        [
            x.narrow(dimension, i, size)
            for i in range(0, x.shape[dimension] - size + 1, step)
        ],
        dim=dimension,
    ).movedim(dimension + 1, -1)


def save(model, sample, path):
    with torch.no_grad():
        ep = torch.export.export(model, sample, strict=False)
        table = get_decomp_table()
        table[torch.ops.aten.unfold.default] = unfold
        ep = ep.run_decompositions(table)
    print("Converting", path, flush=True)
    TorchConverter().add_exported_program(ep).to_coreai().save_asset(path)
    return list(ep.graph_signature.user_outputs)


class Text(torch.nn.Module):
    def __init__(self, m):
        super().__init__()
        self.m = m.language_model

    def forward(self, input_ids, attention_mask, media_embeddings, media_mask):
        emb = self.m.embed_tokens(input_ids)
        emb = torch.where(media_mask, media_embeddings.to(emb.dtype), emb)
        n = input_ids.shape[1]
        pos = torch.arange(n, device=input_ids.device).unsqueeze(0)
        mask = attention_mask[:, None, None, :].bool().expand(1, 1, n, n)
        sliding = (
            pos[:, :, None] - pos[:, None, :]
        ).abs() <= self.m.config.sliding_window
        hidden = self.m(
            inputs_embeds=emb,
            position_ids=pos,
            attention_mask={
                "full_attention": mask,
                "sliding_attention": mask & sliding[:, None],
            },
        ).last_hidden_state
        weights = attention_mask[:, :, None].float()
        pooled = (hidden.float() * weights).sum(1) / weights.sum(1).clamp(min=1)
        return F.normalize(pooled, dim=-1)


class Vision(torch.nn.Module):
    def __init__(self, m):
        super().__init__()
        self.v = m.vision_tower
        self.project = m.embed_vision

    def forward(self, pixel_values, pixel_position_ids):
        v = self.v
        pad = (pixel_position_ids == -1).all(-1)
        emb = v.patch_embedder(pixel_values, pixel_position_ids, pad)
        hidden = v.encoder(
            inputs_embeds=emb,
            attention_mask=~pad,
            pixel_position_ids=pixel_position_ids,
        ).last_hidden_state
        hidden, mask = v.pooler(
            hidden, pixel_position_ids, pad, pixel_values.shape[1] // 9
        )
        return self.project(hidden.to(emb.dtype)), mask


class Audio(torch.nn.Module):
    def __init__(self, m):
        super().__init__()
        self.a = m.audio_tower
        self.project = m.embed_audio

    def forward(self, input_features, input_features_mask):
        out = self.a(input_features, attention_mask=input_features_mask)
        return self.project(out.last_hidden_state), out.attention_mask


if __name__ == "__main__":
    torch.set_num_threads(4)
    root = Path("bundles/embeddinggemma-2-CoreAI")
    root.mkdir(parents=True, exist_ok=True)
    m = AutoModel.from_pretrained(
        "upstream/embeddinggemma-2",
        dtype=torch.float32,
        attn_implementation={
            "language_model": "sdpa",
            "vision_tower": "sdpa",
            "audio_tower": "eager",
        },
    ).eval()
    print("Loaded", flush=True)
    manifest = {
        "source": "google/embeddinggemma-2",
        "revision": "914f7f89142e33e77833254d9c9b90c3cef7303b",
        "precision": "float32",
        "sequence_length": 1024,
        "transformers_revision": "536ecc007387a50e77603bb5d92100e9b07514cc",
    }
    sample = (
        torch.ones(1, 1024, dtype=torch.int32),
        torch.ones(1, 1024, dtype=torch.int32),
        torch.zeros(1, 1024, 512),
        torch.zeros(1, 1024, 1, dtype=torch.bool),
    )
    manifest["text_outputs"] = save(Text(m).eval(), sample, root / "embedding.aimodel")
    (root / "manifest.json").write_text(json.dumps(manifest, indent=2))
    p = AutoProcessor.from_pretrained("upstream/embeddinggemma-2")
    from PIL import Image

    e = p(images=Image.new("RGB", (224, 224)), return_tensors="pt")
    print("Image shapes", [(k, tuple(v.shape)) for k, v in e.items()], flush=True)
    sample = (e["pixel_values"].float(), e["image_position_ids"].int())
    manifest["vision_outputs"] = save(Vision(m).eval(), sample, root / "vision.aimodel")
    manifest["vision_shape"] = list(sample[0].shape)
    (root / "manifest.json").write_text(json.dumps(manifest, indent=2))
    import numpy as np

    e = p(
        audio=np.sin(np.arange(16000) * 0.1).astype(np.float32),
        sampling_rate=16000,
        return_tensors="pt",
    )
    print("Audio shapes", [(k, tuple(v.shape)) for k, v in e.items()], flush=True)
    sample = (
        F.pad(
            e["input_features"].float(), (0, 0, 0, 1120 - e["input_features"].shape[1])
        ),
        F.pad(
            e["input_features_mask"].bool(),
            (0, 1120 - e["input_features_mask"].shape[1]),
        ),
    )
    manifest["audio_outputs"] = save(Audio(m).eval(), sample, root / "audio.aimodel")
    manifest["audio_shape"] = list(sample[0].shape)
    (root / "manifest.json").write_text(json.dumps(manifest, indent=2))