# ---------- mkd_core.py ----------
# (주식회사엠케이디 | MKD)
# MKDCore Transformer — Optimized for 10GB GPU
# Author: MKD AI Systems

import warnings
warnings.filterwarnings("ignore", category=FutureWarning)

import torch
from torch import nn
from tokenizers import Tokenizer

# ============================================================
# CONFIGURATION
# ============================================================
SEQ_LEN = 1024       # Context window (tokens)
D_MODEL = 1024       # Hidden dimension
N_HEADS = 16         # Attention heads
N_LAYERS = 12        # Transformer depth
FF_DIM = D_MODEL * 4 # Feed-forward expansion
DROPOUT = 0.1        # Regularization rate


# ============================================================
# 🔤 TOKENIZER WRAPPER
# ============================================================
class MyTokenizer:
    """
    Wrapper for your trained BPE tokenizer (prod_tokenizer.json)
    Adds <|im_start|> and <|im_end|> for chat/sequence training.
    """
    def __init__(self, json_path: str):
        self.tok = Tokenizer.from_file(json_path)

    def encode(self, text: str, add_special: bool = True) -> list[int]:
        if add_special:
            text = "<|im_start|>" + text + "<|im_end|>"
        return self.tok.encode(text).ids

    def decode(self, ids: list[int]) -> str:
        return self.tok.decode(ids, skip_special_tokens=True)

    @property
    def vocab_size(self) -> int:
        return self.tok.get_vocab_size()


# ============================================================
#  MODEL COMPONENTS
# ============================================================
class FeedForward(nn.Module):
    """Feed-Forward Network (FFN) with GELU activation"""
    def __init__(self, dim: int):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim, FF_DIM),
            nn.GELU(),
            nn.Linear(FF_DIM, dim),
            nn.Dropout(DROPOUT)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


class MultiHeadAttention(nn.Module):
    """Scaled Dot-Product Multi-Head Attention"""
    def __init__(self, dim: int, heads: int):
        super().__init__()
        self.heads = heads
        self.scale = (dim // heads) ** -0.5
        self.qkv = nn.Linear(dim, dim * 3)
        self.proj = nn.Linear(dim, dim)
        self.dropout = nn.Dropout(DROPOUT)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, C = x.shape
        qkv = self.qkv(x).chunk(3, dim=-1)
        q, k, v = [t.view(B, T, self.heads, C // self.heads).transpose(1, 2) for t in qkv]

        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)
        attn = self.dropout(attn)

        out = (attn @ v).transpose(1, 2).contiguous().view(B, T, C)
        return self.proj(out)


class Block(nn.Module):
    """Transformer block: Attention + FeedForward"""
    def __init__(self, dim: int, heads: int):
        super().__init__()
        self.norm1 = nn.LayerNorm(dim)
        self.attn = MultiHeadAttention(dim, heads)
        self.norm2 = nn.LayerNorm(dim)
        self.ff = FeedForward(dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.attn(self.norm1(x))
        x = x + self.ff(self.norm2(x))
        return x


# ============================================================
# 🧠 MKDCore TRANSFORMER MODEL
# ============================================================
class MKDCore(nn.Module):
    """MKDCore Transformer Language Model"""
    def __init__(self, vocab_size: int):
        super().__init__()
        self.token_emb = nn.Embedding(vocab_size, D_MODEL)
        self.pos_emb = nn.Embedding(SEQ_LEN, D_MODEL)
        self.drop = nn.Dropout(DROPOUT)

        self.layers = nn.ModuleList([Block(D_MODEL, N_HEADS) for _ in range(N_LAYERS)])
        self.norm = nn.LayerNorm(D_MODEL)
        self.head = nn.Linear(D_MODEL, vocab_size, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T = x.shape
        pos = torch.arange(0, T, device=x.device).unsqueeze(0)
        x = self.token_emb(x) + self.pos_emb(pos)
        x = self.drop(x)

        for layer in self.layers:
            x = layer(x)

        x = self.norm(x)
        logits = self.head(x)
        return logits


# ============================================================
# ✅ TEST (Run only if executed directly)
# ============================================================
if __name__ == "__main__":
    print("🔍 Testing MKDCore model...")
    vocab_size = 32000
    model = MKDCore(vocab_size=vocab_size)

    dummy = torch.randint(0, vocab_size, (2, 64))  # batch=2, seq=64
    out = model(dummy)
    print(f"✅ Output shape: {out.shape}")  # [2, 64, 32000]