"""DGTRS-CLIP model — text encoder components.

Adapted from https://github.com/MitsuiChen14/DGTRS (Apache-2.0).
Only text-encoder classes are kept; vision encoder removed.
"""

from __future__ import annotations

import logging
from collections import OrderedDict
from pathlib import Path
from typing import Union

import coloredlogs
import numpy as np
import torch
import torch.nn.functional as F
from torch import nn

from src.models.dgtrs.simple_tokenizer import SimpleTokenizer as _Tokenizer

LOGGER = logging.getLogger("caption_test.dgtrs")
coloredlogs.install(level="INFO", logger=LOGGER, fmt="%(asctime)s %(name)s %(levelname)s %(message)s")

_tokenizer = _Tokenizer()


# ---------------------------------------------------------------------------
# Transformer blocks (DGTRS original — sequence-first LND format)
# ---------------------------------------------------------------------------

class LayerNorm(nn.LayerNorm):
    """LayerNorm that handles fp16 by casting to fp32 internally."""

    def forward(self, x: torch.Tensor):
        orig_type = x.dtype
        ret = super().forward(x.type(torch.float32))
        return ret.type(orig_type)


class QuickGELU(nn.Module):
    def forward(self, x: torch.Tensor):
        return x * torch.sigmoid(1.702 * x)


class ResidualAttentionBlock(nn.Module):
    def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None):
        super().__init__()
        self.attn = nn.MultiheadAttention(d_model, n_head)
        self.ln_1 = LayerNorm(d_model)
        self.mlp = nn.Sequential(OrderedDict([
            ("c_fc", nn.Linear(d_model, d_model * 4)),
            ("gelu", QuickGELU()),
            ("c_proj", nn.Linear(d_model * 4, d_model)),
        ]))
        self.ln_2 = LayerNorm(d_model)
        self.attn_mask = attn_mask

    def attention(self, x: torch.Tensor):
        self.attn_mask = (
            self.attn_mask.to(dtype=x.dtype, device=x.device)
            if self.attn_mask is not None else None
        )
        return self.attn(x, x, x, need_weights=False, attn_mask=self.attn_mask)[0]

    def forward(self, x: torch.Tensor):
        x = x + self.attention(self.ln_1(x))
        x = x + self.mlp(self.ln_2(x))
        return x


class Transformer(nn.Module):
    def __init__(self, width: int, layers: int, heads: int, attn_mask: torch.Tensor = None):
        super().__init__()
        self.width = width
        self.layers = layers
        self.resblocks = nn.Sequential(
            *[ResidualAttentionBlock(width, heads, attn_mask) for _ in range(layers)]
        )
        self.gradient_checkpointing = False

    def forward(self, x: torch.Tensor):
        if self.gradient_checkpointing and self.training:
            for block in self.resblocks:
                x = torch.utils.checkpoint.checkpoint(
                    block, x, use_reentrant=False,
                )
            return x
        return self.resblocks(x)


# ---------------------------------------------------------------------------
# DGTRS-CLIP text encoder
# ---------------------------------------------------------------------------

class DGTRSTextEncoder(nn.Module):
    """DGTRS-CLIP text encoder with KPS dual positional embeddings.

    Context length: 248 tokens.
    Uses mask1 (positions 0-19) for positional_embedding and
    mask2 (positions 20-247) for positional_embedding_res.

    This is the official architecture from
    https://github.com/MitsuiChen14/DGTRS/blob/main/model/model_longclip.py
    """

    CONTEXT_LENGTH = 248

    def __init__(
        self,
        vocab_size: int = 49408,
        transformer_width: int = 768,
        transformer_heads: int = 12,
        transformer_layers: int = 12,
        embed_dim: int = 768,
    ) -> None:
        super().__init__()
        self.context_length = self.CONTEXT_LENGTH
        self.embed_dim = embed_dim

        self.token_embedding = nn.Embedding(vocab_size, transformer_width)
        self.positional_embedding = nn.Parameter(
            torch.empty(self.CONTEXT_LENGTH, transformer_width),
        )
        self.positional_embedding_res = nn.Parameter(
            torch.empty(self.CONTEXT_LENGTH, transformer_width),
        )

        self.transformer = Transformer(
            width=transformer_width,
            layers=transformer_layers,
            heads=transformer_heads,
            attn_mask=self._build_attention_mask(),
        )

        self.ln_final = LayerNorm(transformer_width)
        self.text_projection = nn.Parameter(torch.empty(transformer_width, embed_dim))

        # KPS masks: positions 0-19 use positional_embedding,
        # positions 20-247 use positional_embedding_res.
        mask1 = torch.zeros(self.CONTEXT_LENGTH, 1)
        mask1[:20, :] = 1
        mask2 = torch.zeros(self.CONTEXT_LENGTH, 1)
        mask2[20:, :] = 1
        self.register_buffer("mask1", mask1)
        self.register_buffer("mask2", mask2)

    def _build_attention_mask(self) -> torch.Tensor:
        mask = torch.empty(self.CONTEXT_LENGTH, self.CONTEXT_LENGTH)
        mask.fill_(float("-inf"))
        mask.triu_(1)
        return mask

    @property
    def dtype(self) -> torch.dtype:
        return self.token_embedding.weight.dtype

    def forward(self, text: torch.Tensor) -> torch.Tensor:
        """Encode tokenized text.

        Args:
            text: Token IDs [B, 248] (int/long).

        Returns:
            Text embeddings [B, embed_dim].
        """
        x = self.token_embedding(text).type(self.dtype)  # [B, 248, D]

        # Dual masked positional embeddings (KPS).
        x = x + (
            self.positional_embedding * self.mask1
            + self.positional_embedding_res * self.mask2
        ).type(self.dtype)

        x = x.permute(1, 0, 2)  # NLD -> LND (sequence-first for nn.MultiheadAttention)
        x = self.transformer(x)
        x = x.permute(1, 0, 2)  # LND -> NLD
        x = self.ln_final(x).type(self.dtype)

        # Take features from EOT token (highest token ID in each sequence).
        x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
        return x


# ---------------------------------------------------------------------------
# Loading and tokenization
# ---------------------------------------------------------------------------

def _extract_text_state(full_state: dict) -> dict:
    """Extract text encoder keys from a full DGTRS-CLIP state dict."""
    text_keys = {
        "token_embedding", "positional_embedding", "positional_embedding_res",
        "transformer", "ln_final", "text_projection",
    }
    return {
        k: v for k, v in full_state.items()
        if any(k == prefix or k.startswith(prefix + ".") for prefix in text_keys)
    }


def build_model(state_dict: dict) -> DGTRSTextEncoder:
    """Build DGTRSTextEncoder from a DGTRS-CLIP state dict.

    Auto-detects architecture dimensions from the state dict.
    """
    embed_dim = state_dict["text_projection"].shape[1]
    vocab_size = state_dict["token_embedding.weight"].shape[0]
    transformer_width = state_dict["ln_final.weight"].shape[0]
    transformer_heads = transformer_width // 64
    transformer_layers = len(set(
        k.split(".")[2] for k in state_dict if k.startswith("transformer.resblocks")
    ))

    model = DGTRSTextEncoder(
        vocab_size=vocab_size,
        transformer_width=transformer_width,
        transformer_heads=transformer_heads,
        transformer_layers=transformer_layers,
        embed_dim=embed_dim,
    )
    # mask1/mask2 are buffers created in __init__, not in checkpoint.
    model.load_state_dict(state_dict, strict=False)
    return model.eval()


def load_dgtrs_text_encoder(
    checkpoint_path: str | Path,
    device: str = "cpu",
) -> DGTRSTextEncoder:
    """Load DGTRS-CLIP text encoder from checkpoint.

    Extracts only text encoder weights from the full CLIP checkpoint.
    """
    path = Path(checkpoint_path)
    LOGGER.info("📝 Loading DGTRS-CLIP text encoder from %s", path.name)
    full_state = torch.load(str(path), map_location="cpu", weights_only=False)
    if "state_dict" in full_state:
        full_state = full_state["state_dict"]

    text_state = _extract_text_state(full_state)
    model = build_model(text_state)
    model = model.to(device)

    n_params = sum(p.numel() for p in model.parameters())
    LOGGER.info(
        "📝 DGTRS text encoder loaded: %s params, context=%d tokens",
        f"{n_params:,}", model.context_length,
    )
    return model


def tokenize_dgtrs(
    texts: str | list[str],
    context_length: int = 248,
    truncate: bool = True,
) -> torch.Tensor:
    """Tokenize text for DGTRS-CLIP (248 token context).

    Args:
        texts: Input string or list of strings.
        context_length: Output sequence length (default 248).
        truncate: Whether to truncate long sequences.

    Returns:
        Token IDs [B, context_length] (int tensor).
    """
    if isinstance(texts, str):
        texts = [texts]

    sot_token = _tokenizer.encoder["<|startoftext|>"]
    eot_token = _tokenizer.encoder["<|endoftext|>"]
    all_tokens = [[sot_token] + _tokenizer.encode(text) + [eot_token] for text in texts]

    result = torch.zeros(len(all_tokens), context_length, dtype=torch.int)
    for i, tokens in enumerate(all_tokens):
        if len(tokens) > context_length:
            if truncate:
                tokens = tokens[:context_length]
                tokens[-1] = eot_token
            else:
                raise RuntimeError(f"Input too long ({len(tokens)} > {context_length})")
        result[i, :len(tokens)] = torch.tensor(tokens)

    return result
