# Copyright (c) 2025
# G-Transformer: Energy-Efficient Transformer based on GIT
# Author: Syamsuddin B. Ideris, S.Pd.MM

import math
from typing import Optional, Tuple, List, Dict, Any

import torch
import torch.nn as nn
import torch.nn.functional as F

try:
    from transformers import PreTrainedModel, PretrainedConfig
    from transformers.modeling_outputs import CausalLMOutputWithPast
except Exception as e:
    raise ImportError(
        "Harap instal transformers >= 4.40.0. "
        "pip install transformers"
    ) from e


# ----------------------------
# Konfigurasi
# ----------------------------
class GTransformerConfig(PretrainedConfig):
    model_type = "gtransformer"

    def __init__(
        self,
        vocab_size: int = 65536,
        hidden_size: int = 8192,
        intermediate_size: int = 22016,
        num_hidden_layers: int = 48,
        num_attention_heads: int = 64,
        max_position_embeddings: int = 65536,
        hidden_act: str = "swiglu",
        layer_norm_epsilon: float = 1e-5,
        attention_dropout: float = 0.05,
        hidden_dropout_prob: float = 0.05,
        rotary_emb_base: int = 10000,
        use_flash_attention: bool = True,
        use_low_rank_ffn: bool = True,
        use_entropy_gate: bool = True,
        use_moe: bool = False,
        num_experts: int = 0,
        top_k_experts: int = 0,
        fp8_precision: bool = False,
        dvfs_enabled: bool = False,
        informational_constant_kI: float = 2.612e-20,
        energy_per_token_target_J: float = 0.07,
        delta_I_gate: float = 0.75,
        local_window: int = 512,
        global_rank: int = 64,
        kv_compression_rank: int = 64,
        bos_token_id: int = 1,
        eos_token_id: int = 2,
        pad_token_id: int = 0,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.max_position_embeddings = max_position_embeddings
        self.hidden_act = hidden_act
        self.layer_norm_epsilon = layer_norm_epsilon
        self.attention_dropout = attention_dropout
        self.hidden_dropout_prob = hidden_dropout_prob
        self.rotary_emb_base = rotary_emb_base

        self.use_flash_attention = use_flash_attention
        self.use_low_rank_ffn = use_low_rank_ffn
        self.use_entropy_gate = use_entropy_gate

        self.use_moe = use_moe
        self.num_experts = num_experts
        self.top_k_experts = top_k_experts

        self.fp8_precision = fp8_precision
        self.dvfs_enabled = dvfs_enabled

        self.informational_constant_kI = informational_constant_kI
        self.energy_per_token_target_J = energy_per_token_target_J

        self.delta_I_gate = delta_I_gate
        self.local_window = local_window
        self.global_rank = global_rank
        self.kv_compression_rank = kv_compression_rank

        self.bos_token_id = bos_token_id
        self.eos_token_id = eos_token_id
        self.pad_token_id = pad_token_id


# ----------------------------
# Utilitas
# ----------------------------
def swiglu(x: torch.Tensor) -> torch.Tensor:
    x1, x2 = x.chunk(2, dim=-1)
    return F.silu(x1) * x2


def build_activation(name: str):
    if name.lower() == "swiglu":
        return swiglu
    return getattr(F, name)


# Rotary posisi sederhana
class RotaryEmbedding(nn.Module):
    def __init__(self, dim: int, base: int = 10000):
        super().__init__()
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)

    def forward(self, x: torch.Tensor, seq_len: int):
        t = torch.arange(seq_len, device=x.device, dtype=self.inv_freq.dtype)
        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)
        cos = emb.cos()[None, None, :, :]
        sin = emb.sin()[None, None, :, :]
        return cos, sin


def apply_rotary(q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
    # q,k: [B, H, T, D]
    def rotate(x):
        x1, x2 = x[..., ::2], x[..., 1::2]
        x_rot = torch.stack((-x2, x1), dim=-1).flatten(-2)
        return x_rot
    q_rot = (q * cos) + (rotate(q) * sin)
    k_rot = (k * cos) + (rotate(k) * sin)
    return q_rot, k_rot


# ----------------------------
# IA-Attention
# ----------------------------
class InformationalAttention(nn.Module):
    """
    Atensi hemat energi.
    1. Atensi lokal dengan jendela w.
    2. Seleksi token global berbasis skor informasi.
    3. Proyeksi low-rank untuk jalur global.
    """

    def __init__(self, config: GTransformerConfig):
        super().__init__()
        self.config = config
        self.d_model = config.hidden_size
        self.n_heads = config.num_attention_heads
        self.head_dim = self.d_model // self.n_heads
        assert self.d_model % self.n_heads == 0

        self.w_qkv = nn.Linear(self.d_model, 3 * self.d_model, bias=False)
        self.w_o = nn.Linear(self.d_model, self.d_model, bias=False)

        self.rotary = RotaryEmbedding(self.head_dim)

        # Proyeksi low rank global
        self.rank = config.global_rank
        self.Pk = nn.Linear(self.head_dim, self.rank, bias=False)
        self.Pv = nn.Linear(self.head_dim, self.rank, bias=False)
        self.Uo = nn.Linear(self.rank, self.head_dim, bias=False)

        # Skorer informasi
        self.info_scorer = nn.Sequential(
            nn.Linear(self.d_model, self.d_model // 4, bias=False),
            nn.GELU(),
            nn.Linear(self.d_model // 4, 1, bias=False),
        )

        self.attn_drop = nn.Dropout(config.attention_dropout)
        self.proj_drop = nn.Dropout(config.hidden_dropout_prob)

        self.local_window = config.local_window
        self.delta_I_gate = config.delta_I_gate
        self.use_entropy_gate = config.use_entropy_gate

    def _causal_local_mask(self, T: int, w: int, device) -> torch.Tensor:
        idxs = torch.arange(T, device=device)
        mask = idxs[None, :] - idxs[:, None]
        # izinkan hanya masa lalu dalam jendela lokal
        mask = (mask > 0) | (mask < -(w - 1))
        return mask  # True berarti masked

    def forward(
        self,
        x: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
        use_cache: bool = False,
    ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:

        B, T, C = x.shape
        H, D = self.n_heads, self.head_dim

        qkv = self.w_qkv(x)  # [B, T, 3C]
        q, k, v = qkv.split(C, dim=-1)
        q = q.view(B, T, H, D).transpose(1, 2)  # [B, H, T, D]
        k = k.view(B, T, H, D).transpose(1, 2)
        v = v.view(B, T, H, D).transpose(1, 2)

        cos, sin = self.rotary(q, T)
        q, k = apply_rotary(q, k, cos, sin)

        # Tambah cache jika ada
        if past_key_value is not None:
            pk, pv = past_key_value  # [B, H, T_past, D]
            k = torch.cat([pk, k], dim=2)
            v = torch.cat([pv, v], dim=2)
            T_total = k.size(2)
        else:
            T_total = T

        # Atensi lokal
        w = min(self.local_window, T_total)
        scale = 1.0 / math.sqrt(D)
        attn_scores = torch.einsum("bhtd,bhSd->bhtS", q, k) * scale  # S = T_total
        # Mask kausal lokal
        local_mask = self._causal_local_mask(T_total, w, x.device)  # [T_total, T_total]
        local_mask = local_mask[-T:]  # baris untuk query saat ini
        attn_scores = attn_scores.masked_fill(local_mask[None, None, :, :], float("-inf"))
        if attention_mask is not None:
            attn_scores = attn_scores + attention_mask  # bentuk harus broadcastable

        attn_w_local = F.softmax(attn_scores, dim=-1)
        attn_w_local = self.attn_drop(attn_w_local)
        ctx_local = torch.einsum("bhtS,bhSd->bhtd", attn_w_local, v)

        # Seleksi global berbasis informasi
        # Skor informasi dari representasi x
        with torch.no_grad():
            info_score = self.info_scorer(x).squeeze(-1)  # [B, T]
            # skala ke 0..1 via sigmoid
            info_score = torch.sigmoid(info_score)
        if self.use_entropy_gate:
            gate = (info_score > self.delta_I_gate).float()  # [B, T]
        else:
            gate = torch.ones_like(info_score)

        # Proyeksi low rank untuk jalur global hanya pada token bergated
        # Bentuk sederhana: kompres k,v ke rank kecil lalu atensi penuh pada subset
        # Buat mask indeks global per batch
        ctx_global = torch.zeros_like(ctx_local)
        if gate.sum() > 0:
            # kompres k,v
            k_r = self.Pk(k)  # [B,H,T_total,R]
            v_r = self.Pv(v)  # [B,H,T_total,R]
            q_r = self.Pk(q)  # reuse Pk untuk q

            # gunakan atensi penuh pada subset dengan gate
            # bentuk sederhana, gunakan semua posisi, tapi bobot query di-skala gate query
            gate_q = gate[:, -T:].unsqueeze(1).unsqueeze(-1)  # [B,1,T,1]
            attn_scores_g = torch.einsum("bhtr,bhsr->bhts", q_r, k_r) * (scale * D / self.rank)
            attn_w_g = F.softmax(attn_scores_g, dim=-1)
            attn_w_g = self.attn_drop(attn_w_g)
            ctx_g_r = torch.einsum("bhts,bhsr->bhtr", attn_w_g, v_r)
            ctx_g = self.Uo(ctx_g_r)  # [B,H,T,D]
            ctx_global = ctx_g * gate_q

        ctx = ctx_local + ctx_global
        ctx = ctx.transpose(1, 2).contiguous().view(B, T, C)
        out = self.w_o(ctx)
        out = self.proj_drop(out)

        present = (k, v) if use_cache else None
        return out, present


# ----------------------------
# Low-Rank FFN
# ----------------------------
class LowRankFFN(nn.Module):
    def __init__(self, config: GTransformerConfig):
        super().__init__()
        d = config.hidden_size
        i = config.intermediate_size
        act = build_activation(config.hidden_act)
        self.act = act
        # Faktorisasi: d -> i -> d, dengan bottleneck rank r_ffn
        r_ffn = max(128, i // 8)
        self.w1a = nn.Linear(d, r_ffn, bias=False)
        self.w1b = nn.Linear(d, r_ffn, bias=False)
        self.w2 = nn.Linear(r_ffn, d, bias=False)
        self.drop = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # SWiGLU low-rank
        u = self.w1a(x)
        v = self.w1b(x)
        h = swiglu(torch.cat([u, v], dim=-1))
        out = self.w2(h)
        return self.drop(out)


# ----------------------------
# MoE Router opsional
# ----------------------------
class EntropyMoE(nn.Module):
    def __init__(self, config: GTransformerConfig):
        super().__init__()
        assert config.num_experts > 0
        self.num_experts = config.num_experts
        self.top_k = max(1, config.top_k_experts)
        d = config.hidden_size
        i = config.intermediate_size

        self.router = nn.Sequential(
            nn.Linear(d, d // 2, bias=False),
            nn.GELU(),
            nn.Linear(d // 2, self.num_experts, bias=False),
        )
        self.experts = nn.ModuleList(
            [nn.Sequential(nn.Linear(d, i), nn.GELU(), nn.Linear(i, d)) for _ in range(self.num_experts)]
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, D = x.shape
        logits = self.router(x)  # [B,T,E]
        probs = F.softmax(logits, dim=-1)
        topk = torch.topk(probs, k=self.top_k, dim=-1)
        idx = topk.indices  # [B,T,K]
        wgt = topk.values   # [B,T,K]

        out = torch.zeros_like(x)
        for k in range(self.top_k):
            sel = idx[..., k]  # [B,T]
            # kumpulkan untuk tiap expert
            for e in range(self.num_experts):
                mask = (sel == e).float().unsqueeze(-1)  # [B,T,1]
                if mask.sum() == 0:
                    continue
                xe = x * mask
                ye = self.experts[e](xe)
                out = out + ye * (wgt[..., k].unsqueeze(-1))
        return out


# ----------------------------
# Blok Transformer
# ----------------------------
class GTransformerBlock(nn.Module):
    def __init__(self, config: GTransformerConfig):
        super().__init__()
        self.ln1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon)
        self.attn = InformationalAttention(config)
        self.ln2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon)
        if config.use_moe and config.num_experts > 0:
            self.ff = EntropyMoE(config)
        else:
            self.ff = LowRankFFN(config) if config.use_low_rank_ffn else nn.Sequential(
                nn.Linear(config.hidden_size, config.intermediate_size),
                nn.GELU(),
                nn.Linear(config.intermediate_size, config.hidden_size),
            )

    def forward(
        self,
        x: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
        use_cache: bool = False,
    ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
        h, present = self.attn(self.ln1(x), attention_mask=attention_mask, past_key_value=past_key_value, use_cache=use_cache)
        x = x + h
        x = x + self.ff(self.ln2(x))
        return x, present


# ----------------------------
# Model dasar
# ----------------------------
class GTransformerModel(PreTrainedModel):
    config_class = GTransformerConfig

    def __init__(self, config: GTransformerConfig):
        super().__init__(config)
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = nn.ModuleList([GTransformerBlock(config) for _ in range(config.num_hidden_layers)])
        self.ln_f = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon)

        self.gradient_checkpointing = False

        self.post_init()

    def forward(
        self,
        input_ids: torch.LongTensor,
        attention_mask: Optional[torch.Tensor] = None,
        past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
        use_cache: Optional[bool] = None,
        **kwargs,
    ) -> Tuple[torch.Tensor, Optional[List[Tuple[torch.Tensor, torch.Tensor]]]]:

        B, T = input_ids.shape
        x = self.embed_tokens(input_ids)

        new_past = [] if use_cache else None
        for i, layer in enumerate(self.layers):
            pkv = None if past_key_values is None else past_key_values[i]
            x, present = layer(x, attention_mask=attention_mask, past_key_value=pkv, use_cache=use_cache)
            if use_cache:
                new_past.append(present)

        x = self.ln_f(x)
        return x, new_past


# ----------------------------
# Causal LM
# ----------------------------
class GTransformerForCausalLM(PreTrainedModel):
    config_class = GTransformerConfig

    def __init__(self, config: GTransformerConfig):
        super().__init__(config)
        self.transformer = GTransformerModel(config)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.post_init()

    def get_input_embeddings(self):
        return self.transformer.embed_tokens

    def set_input_embeddings(self, new_embeddings):
        self.transformer.embed_tokens = new_embeddings

    def tie_weights(self):
        # opsional tidak diikat agar stabil FP8
        pass

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        labels: Optional[torch.LongTensor] = None,
        past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
        use_cache: Optional[bool] = None,
        **kwargs,
    ) -> CausalLMOutputWithPast:

        hidden_states, new_past = self.transformer(
            input_ids=input_ids,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            use_cache=use_cache,
        )
        logits = self.lm_head(hidden_states)

        loss = None
        if labels is not None:
            shift_logits = logits[:, :-1, :].contiguous()
            shift_labels = labels[:, 1:].contiguous()
            loss = F.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)),
                shift_labels.view(-1),
                ignore_index=-100,
            )

            # Regularisasi informasi sederhana
            if self.config.use_entropy_gate:
                with torch.no_grad():
                    probs = F.softmax(shift_logits, dim=-1)
                    logp = torch.log(probs + 1e-9)
                    H = -(probs * logp).sum(dim=-1).mean()
                # target penurunan entropi moderat
                loss = loss + 1e-4 * H

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=new_past,
            hidden_states=None,
            attentions=None,
        )

    @torch.no_grad()
    def generate_simple(
        self,
        input_ids: torch.LongTensor,
        max_new_tokens: int = 64,
        temperature: float = 1.0,
    ) -> torch.LongTensor:
        self.eval()
        past = None
        out = input_ids
        for _ in range(max_new_tokens):
            logits = self(out[:, -1:].contiguous(), use_cache=True, past_key_values=past).logits
            past = self(out[:, -1:].contiguous(), use_cache=True, past_key_values=past).past_key_values
            next_token = torch.distributions.Categorical(logits=logits[:, -1, :] / max(1e-6, temperature)).sample()
            out = torch.cat([out, next_token.unsqueeze(-1)], dim=1)
            if int(next_token[0].item()) == self.config.eos_token_id:
                break
        return out