"""
TRM-text-ISM: 単一ファイル構成。

なぜ1ファイルにしたか:
  config と modeling を別ファイル + relative import (`from .config import ...`) に
  分けると、`save_pretrained()` → `push_to_hub()` → `from_pretrained(trust_remote_code=True)`
  という往復で `ModuleNotFoundError` を起こす既知の不具合がある
  (huggingface/transformers issue #40496, 2025-08)。
  Falcon/ChatGLM2など実運用のHubモデルの多くも、複数ファイル構成を避けて
  configとmodelingを1ファイルに収めることでこれを回避している。

  このファイルだけを `modeling_trm_text_ism.py` としてHubに置けば、
  Colabでの直importでも、Hubのtrust_remote_code経由でも、同一コードパスで動く。
"""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint
from transformers import PreTrainedModel, PretrainedConfig
from transformers.generation import GenerationMixin
from transformers.modeling_outputs import CausalLMOutputWithPast


# ============================== Config ==============================

class TRMTextISMConfig(PretrainedConfig):
    """
    TRM-text (ISM) config.

    アーキテクチャ: RMSNorm + SwiGLU + RoPE + gated residual の `TRMBlock` を
    `n_layers` 層積み、それを `recurrence_steps` 回ループ（recurrent-depth）。
    n_layers=1 で「1ブロックをrecurrence_steps回」という最初の形と同一挙動。

    制約: n_heads * head_dim == dim （qkvがdimにしか射影しないためMHAのみ）。
    """

    model_type = "trm_text_ism"
    auto_map = {
        "AutoConfig": "modeling_trm_text_ism.TRMTextISMConfig",
        "AutoModelForCausalLM": "modeling_trm_text_ism.TRMTextISMForCausalLM",
    }

    def __init__(
        self,
        vocab_size: int = 151936,
        dim: int = 2048,
        n_layers: int = 1,
        n_heads: int = 16,
        head_dim: int = 128,
        mlp_ratio: float = 2.6875,
        mlp_hidden_size: int | None = 5632,
        recurrence_steps: int = 4,
        max_seq_len: int = 2048,
        residual_scale: float = 1.0,
        tie_word_embeddings: bool = False,
        pad_token_id: int | None = None,
        bos_token_id: int | None = None,
        eos_token_id: int | None = None,
        **kwargs,
    ):
        self.vocab_size = vocab_size
        self.dim = dim
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.head_dim = head_dim
        self.mlp_ratio = mlp_ratio
        self.mlp_hidden_size = mlp_hidden_size
        self.recurrence_steps = recurrence_steps
        self.max_seq_len = max_seq_len
        self.residual_scale = residual_scale
        kwargs["use_cache"] = False  # KVキャッシュ未実装。generate()のcache分岐を踏ませない

        super().__init__(
            tie_word_embeddings=tie_word_embeddings,
            pad_token_id=pad_token_id,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            **kwargs,
        )

    @property
    def hidden_size(self) -> int:
        return self.dim

    @property
    def num_attention_heads(self) -> int:
        return self.n_heads

    @property
    def num_hidden_layers(self) -> int:
        return self.n_layers

    @property
    def num_key_value_heads(self) -> int:
        return self.n_heads

    @property
    def _mlp_hidden(self) -> int:
        return self.mlp_hidden_size or int(self.dim * self.mlp_ratio)

    def param_breakdown(self) -> dict:
        d, h, V = self.dim, self._mlp_hidden, self.vocab_size
        attn = 3 * d * d + d * d
        mlp = 3 * d * h
        norms = 2 * d
        gates = 2 * d
        per_block = attn + mlp + norms + gates
        blocks = self.n_layers * per_block
        final_norm = d
        token_emb = V * d
        lm_head = 0 if self.tie_word_embeddings else V * d
        non_emb = blocks + final_norm
        total = non_emb + token_emb + lm_head
        return {
            "token_emb": token_emb, "lm_head": lm_head, "per_block": per_block,
            "blocks_total": blocks, "final_norm": final_norm,
            "non_embedding": non_emb, "total": total,
            "embedding_share": (token_emb + lm_head) / total,
        }

    def num_parameters(self, include_embeddings: bool = True) -> int:
        b = self.param_breakdown()
        return b["total"] if include_embeddings else b["non_embedding"]

    def __post_init_check__(self):
        assert self.n_heads * self.head_dim == self.dim, (
            f"n_heads*head_dim ({self.n_heads}*{self.head_dim}) != dim ({self.dim})."
        )


TRM_TEXT_PRESETS: dict[str, dict] = {
    "debug": dict(dim=512, n_layers=4, n_heads=8, head_dim=64,
                  mlp_hidden_size=1408, recurrence_steps=4, max_seq_len=1024),
    "950m": dict(dim=2048, n_layers=7, n_heads=16, head_dim=128,
                 mlp_hidden_size=5632, recurrence_steps=4, max_seq_len=2048),
    "1b": dict(dim=2048, n_layers=8, n_heads=16, head_dim=128,
               mlp_hidden_size=5632, recurrence_steps=4, max_seq_len=2048),
    "1b-single": dict(dim=3072, n_layers=1, n_heads=24, head_dim=128,
                      mlp_hidden_size=8192, recurrence_steps=4, max_seq_len=2048),
    "1.3b": dict(dim=2304, n_layers=10, n_heads=18, head_dim=128,
                 mlp_hidden_size=6144, recurrence_steps=4, max_seq_len=2048),
}


def trm_text_config(preset: str = "1b", **overrides) -> TRMTextISMConfig:
    if preset not in TRM_TEXT_PRESETS:
        raise KeyError(f"unknown preset {preset!r}. choices: {list(TRM_TEXT_PRESETS)}")
    cfg_kwargs = {**TRM_TEXT_PRESETS[preset], **overrides}
    cfg = TRMTextISMConfig(**cfg_kwargs)
    cfg.__post_init_check__()
    return cfg


# ============================== Model ==============================

def apply_rope(x, cos, sin):
    S = x.shape[2]
    c, s = cos[:, :, :S, :].to(x.dtype), sin[:, :, :S, :].to(x.dtype)
    x1, x2 = x[..., :x.shape[-1] // 2], x[..., x.shape[-1] // 2:]
    return torch.cat([x1 * c - x2 * s, x2 * c + x1 * s], dim=-1)


class SwiGLUMLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        h = config.mlp_hidden_size or int(config.dim * config.mlp_ratio)
        self.gate_proj = nn.Linear(config.dim, h, bias=False)
        self.up_proj = nn.Linear(config.dim, h, bias=False)
        self.down_proj = nn.Linear(h, config.dim, bias=False)
        self.down_proj._scale_init = True

    def forward(self, x):
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


class TRMAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.n_heads, self.head_dim = config.n_heads, config.head_dim
        assert self.n_heads * self.head_dim == config.dim, \
            "n_heads*head_dim must equal dim (qkv projects to dim only)"
        self.qkv = nn.Linear(config.dim, 3 * config.dim, bias=False)
        self.out = nn.Linear(config.dim, config.dim, bias=False)
        self.out._scale_init = True

    def forward(self, x, attn_mask, cos, sin, is_causal=False):
        B, S, _ = x.shape
        q, k, v = self.qkv(x).chunk(3, dim=-1)
        q, k, v = [t.view(B, S, self.n_heads, self.head_dim).transpose(1, 2) for t in (q, k, v)]
        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
        if attn_mask is None:
            y = F.scaled_dot_product_attention(q, k, v, is_causal=is_causal)
        else:
            y = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask[:, None, :, :])
        return self.out(y.transpose(1, 2).reshape(B, S, -1))


class TRMBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.res = config.residual_scale
        self.norm1 = nn.RMSNorm(config.dim)
        self.attn = TRMAttention(config)
        self.norm2 = nn.RMSNorm(config.dim)
        self.mlp = SwiGLUMLP(config)
        self.attn_gate = nn.Parameter(torch.ones(config.dim))
        self.mlp_gate = nn.Parameter(torch.ones(config.dim))

    def forward(self, x, attn_mask, c, s, is_causal=False):
        x = x + self.res * torch.sigmoid(self.attn_gate).view(1, 1, -1) * self.attn(self.norm1(x), attn_mask, c, s, is_causal)
        return x + self.res * torch.sigmoid(self.mlp_gate).view(1, 1, -1) * self.mlp(self.norm2(x))


class TRMTextISMForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = TRMTextISMConfig
    supports_gradient_checkpointing = True

    def __init__(self, config):
        super().__init__(config)
        self.token_emb = nn.Embedding(config.vocab_size, config.dim)
        self.blocks = nn.ModuleList([TRMBlock(config) for _ in range(config.n_layers)])
        self.norm = nn.RMSNorm(config.dim)
        self.lm_head = nn.Linear(config.dim, config.vocab_size, bias=False)
        self.gradient_checkpointing = False

        pos = torch.arange(config.max_seq_len).float()
        theta = 1.0 / (10000.0 ** (torch.arange(0, config.head_dim // 2).float() / (config.head_dim // 2)))
        f = torch.outer(pos, theta)
        # persistent=True (デフォルト): from_pretrained の low_cpu_mem_usage 経路では
        # モデルが meta device 上に一旦構築され、その後 state_dict から重みがロードされる。
        # persistent=False のバッファは state_dict に乗らないため、このロード経路では
        # meta device 上の未初期化値のまま残ってしまい、cos()/sin() の出力が
        # 1e+34 のような異常値になってNaNが全体に伝播する事故が起きた。
        # config から再計算可能な値であっても、ロード安全性のため persistent のままにする。
        self.register_buffer("rope_cos", f.cos().view(1, 1, config.max_seq_len, -1))
        self.register_buffer("rope_sin", f.sin().view(1, 1, config.max_seq_len, -1))
        self.post_init()

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            std = 0.02
            if getattr(module, "_scale_init", False):
                eff_depth = self.config.n_layers * self.config.recurrence_steps
                std = 0.02 / math.sqrt(2 * max(1, eff_depth))
            nn.init.normal_(module.weight, mean=0.0, std=std)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def get_input_embeddings(self):
        return self.token_emb

    def set_input_embeddings(self, value):
        self.token_emb = value

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, value):
        self.lm_head = value

    def gradient_checkpointing_enable(self, **kwargs):
        self.gradient_checkpointing = True

    def gradient_checkpointing_disable(self):
        self.gradient_checkpointing = False

    def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **kwargs):
        if attention_mask is None:
            attention_mask = torch.ones_like(input_ids)
        return {"input_ids": input_ids, "attention_mask": attention_mask, "use_cache": False}

    def forward(self, input_ids, attention_mask=None, labels=None, **kwargs):
        if input_ids.numel() > 0:
            lo, hi = input_ids.min().item(), input_ids.max().item()
            if lo < 0 or hi >= self.config.vocab_size:
                raise ValueError(
                    f"input_ids out of range for vocab_size={self.config.vocab_size}: "
                    f"min={lo}, max={hi}. tokenizerのvocabとconfig.vocab_sizeが食い違っている、"
                    f"またはpad_token_id/eos_token_idがNoneのまま渡っている可能性が高い。"
                )
        B, S = input_ids.shape
        x = self.token_emb(input_ids)

        if attention_mask is None:
            m, is_causal = None, True
        else:
            m = torch.tril(torch.ones(S, S, device=input_ids.device)).bool().unsqueeze(0).expand(B, -1, -1)
            m = m & attention_mask[:, None, :].bool()
            is_causal = False

        c, s = self.rope_cos, self.rope_sin

        for _ in range(self.config.recurrence_steps):
            for blk in self.blocks:
                if self.gradient_checkpointing and self.training:
                    x = checkpoint(blk, x, m, c, s, is_causal, use_reentrant=False)
                else:
                    x = blk(x, m, c, s, is_causal)

        logits = self.lm_head(self.norm(x))

        loss = None
        if labels is not None:
            loss = F.cross_entropy(
                logits[:, :-1].reshape(-1, logits.size(-1)).float(),
                labels[:, 1:].reshape(-1),
                ignore_index=-100,
            )
        return CausalLMOutputWithPast(loss=loss, logits=logits)