"""
UNCHA: Uncertainty-guided Compositional Hyperbolic Alignment
Zero-shot image classification demo using hyperbolic (Lorentz) embeddings.

Paper: https://arxiv.org/abs/2603.22042
Code:  https://github.com/jeeit17/UNCHA
"""

import os
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # MUST come before torch / any CUDA-touching import

import gzip
import html
import math
import re
import regex
from collections import OrderedDict
from pathlib import Path

import numpy as np
import torch
from torch import nn
from torch.nn import functional as F
import timm
import gradio as gr
from PIL import Image

import torchvision.transforms as T

# ---------------------------------------------------------------------------
# Tokenizer (adapted from UNCHA/CLIP BPE tokenizer)
# ---------------------------------------------------------------------------

class Tokenizer:
    """Byte-Pair Encoding tokenizer compatible with CLIP / UNCHA checkpoints."""

    def __init__(self, bpe_path: str | Path | None = None):
        bs = (
            list(range(ord("!"), ord("~") + 1))
            + list(range(ord("\xa1"), ord("\xac") + 1))
            + list(range(ord("\xae"), ord("\xff") + 1))
        )
        self.byte_encoder = {b: chr(b) for b in bs}
        n = 0
        for b in range(2**8):
            if b not in self.byte_encoder:
                self.byte_encoder[b] = chr(2**8 + n)
                n += 1

        if bpe_path is None:
            bpe_path = Path(__file__).resolve().parent / "bpe_simple_vocab_16e6.txt.gz"
        merges = gzip.open(bpe_path).read().decode("utf-8").split("\n")
        merges = merges[1 : 49152 - 256 - 2 + 1]
        merges = [tuple(merge.split()) for merge in merges]
        vocab = list(self.byte_encoder.values())
        vocab = vocab + [v + "</w>" for v in vocab]
        for merge in merges:
            vocab.append("".join(merge))
        vocab.extend(["<|startoftext|>", ""])

        self.encoder = dict(zip(vocab, range(len(vocab))))
        self.bpe_ranks = dict(zip(merges, range(len(merges))))
        self.cache = {"<|startoftext|>": "<|startoftext|>", "": ""}
        self.pat = regex.compile(
            r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""",
            regex.IGNORECASE,
        )

    def __call__(self, text):
        import ftfy

        text_list = [text] if isinstance(text, str) else text
        token_tensors = []
        for text in text_list:
            bpe_tokens = []
            text = ftfy.fix_text(text)
            text = html.unescape(html.unescape(text))
            text = re.sub(r"\s+", " ", text)
            text = text.strip().lower()
            for token in regex.findall(self.pat, text):
                token = "".join(self.byte_encoder[b] for b in token.encode("utf-8"))
                bpe_tokens.extend(
                    self.encoder[bpe_token]
                    for bpe_token in self.bpe(token).split(" ")
                )
            sot = self.encoder["<|startoftext|>"]
            eot = self.encoder[""]
            bpe_tokens = [sot, *bpe_tokens, eot]
            token_tensors.append(torch.IntTensor(bpe_tokens))
        return token_tensors

    @staticmethod
    def get_pairs(word):
        pairs = set()
        prev_char = word[0]
        for char in word[1:]:
            pairs.add((prev_char, char))
            prev_char = char
        return pairs

    def bpe(self, token):
        if token in self.cache:
            return self.cache[token]
        word = tuple(token[:-1]) + (token[-1] + "</w>",)
        pairs = self.get_pairs(word)
        if not pairs:
            return token + "</w>"
        while True:
            bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf")))
            if bigram not in self.bpe_ranks:
                break
            first, second = bigram
            new_word = []
            i = 0
            while i < len(word):
                try:
                    j = word.index(first, i)
                    new_word.extend(word[i:j])
                    i = j
                except ValueError:
                    new_word.extend(word[i:])
                    break
                if word[i] == first and i < len(word) - 1 and word[i + 1] == second:
                    new_word.append(first + second)
                    i += 2
                else:
                    new_word.append(word[i])
                    i += 1
            new_word = tuple(new_word)
            word = new_word
            if len(word) == 1:
                break
            else:
                pairs = self.get_pairs(word)
        word = " ".join(word)
        self.cache[token] = word
        return word


# ---------------------------------------------------------------------------
# Lorentz model hyperbolic operations (adapted from UNCHA/meru)
# ---------------------------------------------------------------------------

def pairwise_inner(x, y, curv=1.0):
    x_time = torch.sqrt(1 / curv + torch.sum(x**2, dim=-1, keepdim=True))
    y_time = torch.sqrt(1 / curv + torch.sum(y**2, dim=-1, keepdim=True))
    return x @ y.T - x_time @ y_time.T


def exp_map0(x, curv=1.0, eps=1e-8):
    rc_xnorm = curv**0.5 * torch.norm(x, dim=-1, keepdim=True)
    sinh_input = torch.clamp(rc_xnorm, min=eps, max=math.asinh(2**15))
    return torch.sinh(sinh_input) * x / torch.clamp(rc_xnorm, min=eps)


# ---------------------------------------------------------------------------
# Text encoder (adapted from UNCHA TransformerTextEncoder)
# ---------------------------------------------------------------------------

class _TransformerBlock(nn.Module):
    def __init__(self, d_model, n_head):
        super().__init__()
        self.attn = nn.MultiheadAttention(d_model, n_head, batch_first=True)
        self.ln_1 = nn.LayerNorm(d_model)
        self.mlp = nn.Sequential(
            OrderedDict([
                ("c_fc", nn.Linear(d_model, d_model * 4)),
                ("gelu", nn.GELU()),
                ("c_proj", nn.Linear(d_model * 4, d_model)),
            ])
        )
        self.ln_2 = nn.LayerNorm(d_model)

    def forward(self, x, attn_mask=None):
        lx = self.ln_1(x)
        ax = self.attn(lx, lx, lx, need_weights=False, attn_mask=attn_mask)[0]
        x = x + ax
        x = x + self.mlp(self.ln_2(x))
        return x


class TransformerTextEncoder(nn.Module):
    def __init__(self, arch="L12_W512", vocab_size=49408, context_length=77):
        super().__init__()
        self.vocab_size = vocab_size
        self.context_length = context_length
        self.layers = int(re.search(r"L(\d+)", arch).group(1))
        self.width = int(re.search(r"W(\d+)", arch).group(1))
        _attn = re.search(r"A(\d+)", arch)
        self.heads = int(_attn.group(1)) if _attn else self.width // 64

        self.token_embed = nn.Embedding(vocab_size, self.width)
        self.posit_embed = nn.Parameter(torch.empty(context_length, self.width))
        _resblocks = [_TransformerBlock(self.width, self.heads) for _ in range(self.layers)]
        self.resblocks = nn.ModuleList(_resblocks)
        self.ln_final = nn.LayerNorm(self.width)

        attn_mask = torch.triu(
            torch.full((context_length, context_length), float("-inf")), diagonal=1
        )
        self.register_buffer("attn_mask", attn_mask.bool())

        nn.init.normal_(self.token_embed.weight, std=0.02)
        nn.init.normal_(self.posit_embed.data, std=0.01)
        out_proj_std = (2 * self.width * self.layers) ** -0.5
        for block in self.resblocks:
            nn.init.normal_(block.attn.in_proj_weight, std=self.width**-0.5)
            nn.init.normal_(block.attn.out_proj.weight, std=out_proj_std)
            nn.init.normal_(block.mlp[0].weight, std=(2 * self.width) ** -0.5)
            nn.init.normal_(block.mlp[2].weight, std=out_proj_std)

    def forward(self, text_tokens: torch.Tensor) -> torch.Tensor:
        """Forward pass: tokenize -> transformer blocks -> ln_final."""
        max_len = text_tokens.shape[-1]
        _posit_embed = self.posit_embed[:max_len, :]
        _attn_mask = self.attn_mask[:max_len, :max_len]

        token_embeddings = self.token_embed(text_tokens) + _posit_embed

        textual_features = token_embeddings
        for block in self.resblocks:
            textual_features = block(textual_features, _attn_mask)

        textual_features = self.ln_final(textual_features)
        return textual_features


def build_timm_vit(arch="vit_base_patch16_224", global_pool="token",
                   use_sincos2d_pos=True):
    model = timm.create_model(
        arch, num_classes=0, global_pool=global_pool,
        class_token=global_pool == "token", norm_layer=nn.LayerNorm,
    )
    model.width = model.embed_dim
    if use_sincos2d_pos:
        h, w = model.patch_embed.grid_size
        grid_w = torch.arange(w, dtype=torch.float32)
        grid_h = torch.arange(h, dtype=torch.float32)
        grid_w, grid_h = torch.meshgrid(grid_w, grid_h)
        pos_dim = model.embed_dim // 4
        omega = torch.arange(pos_dim, dtype=torch.float32) / pos_dim
        omega = 1.0 / (10000.0**omega)
        out_w = torch.einsum("m,d->md", [grid_w.flatten(), omega])
        out_h = torch.einsum("m,d->md", [grid_h.flatten(), omega])
        pos_emb = torch.cat(
            [torch.sin(out_w), torch.cos(out_w), torch.sin(out_h), torch.cos(out_h)],
            dim=1,
        )[None, :, :]
        if global_pool == "token":
            pe_token = torch.zeros([1, 1, model.embed_dim], dtype=torch.float32)
            pos_emb = torch.cat([pe_token, pos_emb], dim=1)
        model.pos_embed.data.copy_(pos_emb)
        model.pos_embed.requires_grad = False
    return model


# ---------------------------------------------------------------------------
# UNCHA model (inference-only)
# ---------------------------------------------------------------------------

class UNCHAModel(nn.Module):
    """Inference-only UNCHA model: hyperbolic image-text alignment."""

    def __init__(self, embed_dim=512, visual_arch="vit_base_patch16_224",
                 text_arch="L12_W512", vocab_size=49408, context_length=77):
        super().__init__()
        self.visual = build_timm_vit(arch=visual_arch)
        self.textual = TransformerTextEncoder(
            arch=text_arch, vocab_size=vocab_size, context_length=context_length
        )
        self.embed_dim = embed_dim
        self.visual_proj = nn.Linear(self.visual.width, embed_dim, bias=False)
        self.textual_proj = nn.Linear(self.textual.width, embed_dim, bias=False)
        self.logit_scale = nn.Parameter(torch.tensor(1 / 0.07).log())
        self.curv = nn.Parameter(torch.tensor(1.0).log())
        self.visual_alpha = nn.Parameter(torch.tensor(embed_dim**-0.5).log())
        self.textual_alpha = nn.Parameter(torch.tensor(embed_dim**-0.5).log())
        self.tokenizer = Tokenizer()
        self.register_buffer("pixel_mean", torch.tensor((0.485, 0.456, 0.406)).view(-1, 1, 1))
        self.register_buffer("pixel_std", torch.tensor((0.229, 0.224, 0.225)).view(-1, 1, 1))

    @property
    def device(self):
        return self.logit_scale.device

    def encode_image(self, images, project=True):
        images = (images - self.pixel_mean) / self.pixel_std
        feats = self.visual(images)
        feats = self.visual_proj(feats)
        if project:
            feats = feats * self.visual_alpha.exp()
            with torch.autocast(self.device.type, dtype=torch.float32):
                feats = exp_map0(feats, self.curv.exp())
        return feats

    def encode_text(self, tokens, project=True):
        context_len = self.textual.context_length
        batch_size = len(tokens)
        padded = torch.zeros((batch_size, context_len), dtype=torch.long)
        for idx, inst in enumerate(tokens):
            L_ = min(inst.shape[0], context_len)
            if inst.shape[0] > context_len:
                inst = inst[:context_len]
            padded[idx, :L_] = inst[:L_]
        padded = padded.to(self.device)
        feats = self.textual(padded)
        eos = padded.argmax(dim=-1)
        batch_idx = torch.arange(batch_size, device=self.device)
        feats = feats[batch_idx, eos]
        feats = self.textual_proj(feats)
        if project:
            feats = feats * self.textual_alpha.exp()
            with torch.autocast(self.device.type, dtype=torch.float32):
                feats = exp_map0(feats, self.curv.exp())
        return feats


# ---------------------------------------------------------------------------
# Image preprocessing (matching the evaluation pipeline)
# ---------------------------------------------------------------------------

IMAGE_TRANSFORM = T.Compose([
    T.Resize(224, T.InterpolationMode.BICUBIC),
    T.CenterCrop(224),
    T.ToTensor(),
])

# ---------------------------------------------------------------------------
# Load model at module scope
# ---------------------------------------------------------------------------

CHECKPOINT_REPO = "hayeonkim/uncha"
CHECKPOINT_FILE = "uncha_vit_b.pth"

print("Loading UNCHA model...")
model = UNCHAModel(
    embed_dim=512,
    visual_arch="vit_base_patch16_224",
    text_arch="L12_W512",
    vocab_size=49408,
    context_length=77,
)

# Download and load checkpoint
from huggingface_hub import hf_hub_download
_ckpt_path = hf_hub_download(CHECKPOINT_REPO, CHECKPOINT_FILE)
_ckpt = torch.load(_ckpt_path, map_location="cpu", weights_only=False)
_sd = _ckpt["model"]

# The checkpoint may contain min_radius_head keys from the text encoder that
# our inference model doesn't have — filter them out.
_model_sd = model.state_dict()
_filtered_sd = {}
for k, v in _sd.items():
    if k in _model_sd:
        _filtered_sd[k] = v
    else:
        print(f"  Skipping checkpoint key not in model: {k}")

_missing, _unexpected = model.load_state_dict(_filtered_sd, strict=False)
if _missing:
    print(f"  Missing keys: {_missing}")
if _unexpected:
    print(f"  Unexpected keys: {_unexpected}")

model = model.eval().to("cuda")
print(f"Model loaded. curv={model.curv.exp().item():.4f}, "
      f"visual_alpha={model.visual_alpha.exp().item():.4f}, "
      f"textual_alpha={model.textual_alpha.exp().item():.4f}")

# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------

PROMPT_TEMPLATES = [
    "a photo of a {}.",
    "a blurry photo of a {}.",
    "a black and white photo of a {}.",
    "a low contrast photo of a {}.",
    "a high contrast photo of a {}.",
    "a bad photo of a {}.",
    "a good photo of a {}.",
    "a photo of a small {}.",
    "a photo of a big {}.",
    "a photo of the {}.",
]

@spaces.GPU(duration=60)
def classify(image: Image.Image, candidate_labels: str) -> dict:
    """Zero-shot image classification using UNCHA hyperbolic vision-language model.

    Args:
        image: Input image to classify.
        candidate_labels: Comma-separated list of candidate class labels.

    Returns:
        Dictionary mapping each label to its probability score.
    """
    if image is None:
        return {}
    if image.mode != "RGB":
        image = image.convert("RGB")

    # Parse labels
    labels = [l.strip() for l in candidate_labels.split(",") if l.strip()]
    if not labels:
        return {}

    # Preprocess image
    img_tensor = IMAGE_TRANSFORM(image).unsqueeze(0).to(model.device)

    # Encode image into hyperbolic space
    with torch.inference_mode():
        img_feats = model.encode_image(img_tensor, project=True)  # (1, D)

    # Encode text prompts for each label (prompt ensemble in tangent space)
    with torch.inference_mode():
        all_class_feats = []
        for label in labels:
            prompts = [pt.format(label) for pt in PROMPT_TEMPLATES]
            tokens = model.tokenizer(prompts)
            text_feats = model.encode_text(tokens, project=False)  # (N_prompts, D)
            # Ensemble in tangent space, then project to hyperboloid
            text_feats = text_feats.mean(dim=0)  # (D,)
            text_feats = text_feats * model.textual_alpha.exp()
            text_feats = exp_map0(text_feats.unsqueeze(0), model.curv.exp())  # (1, D)
            all_class_feats.append(text_feats.squeeze(0))

        classifier = torch.stack(all_class_feats, dim=0)  # (num_classes, D)

        # Lorentzian pairwise inner product as classification scores
        scores = pairwise_inner(img_feats, classifier, model.curv.exp())  # (1, num_classes)
        scores = scores.squeeze(0)  # (num_classes,)

        # Convert to probabilities via softmax
        probs = F.softmax(scores * model.logit_scale.exp(), dim=-1)

    # Build label->prob dict, sorted by probability
    result = {label: float(probs[i]) for i, label in enumerate(labels)}
    result = dict(sorted(result.items(), key=lambda x: x[1], reverse=True))
    return result


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------

CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(elem_id="col-container", css=CSS) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # UNCHA: Uncertainty-guided Compositional Hyperbolic Alignment

            Zero-shot image classification using a hyperbolic vision-language model.
            Upload an image and provide candidate labels — the model will compute
            similarity scores using Lorentzian (hyperbolic) geometry.

            [Paper](https://arxiv.org/abs/2603.22042) | [Code](https://github.com/jeeit17/UNCHA) | [Model](https://huggingface.co/hayeonkim/uncha)
            """
        )

        with gr.Row():
            with gr.Column(scale=1):
                image_input = gr.Image(
                    type="pil", label="Input Image",
                    sources=["upload", "clipboard"],
                )
                labels_input = gr.Textbox(
                    label="Candidate Labels (comma-separated)",
                    placeholder="cat, dog, bird, car, person",
                    value="cat, dog, bird, car, person",
                )
                run_btn = gr.Button("Classify", variant="primary")

            with gr.Column(scale=1):
                output_labels = gr.Label(
                    label="Classification Scores",
                    num_top_classes=10,
                )

    run_btn.click(
        fn=classify,
        inputs=[image_input, labels_input],
        outputs=output_labels,
        api_name="classify",
    )

    gr.Examples(
        examples=[
            ["examples/sample1.jpg", "dog, cat, animal, pet, mammal"],
            ["examples/sample2.jpg", "building, landscape, nature, city, water"],
            ["examples/sample3.jpg", "person, food, vehicle, animal, plant"],
        ],
        inputs=[image_input, labels_input],
        outputs=output_labels,
        fn=classify,
        cache_examples=True,
        cache_mode="lazy",
    )

demo.launch(mcp_server=True, theme=gr.themes.Citrus())