import torch
import torch.nn as nn
from model.config import ModelConfig
from model.transformer import GPT, Block, RMSNorm

class DraftModel(nn.Module):
    """
    A very small version of the main GPT model used for Speculative Decoding.
    Usually 2-4 layers with smaller embedding dimension.
    """
    def __init__(self, config: ModelConfig, draft_layers=2, draft_embd=128):
        super().__init__()
        # Create a modified config for the draft model
        self.config = config
        
        self.token_embedding_table = nn.Embedding(config.vocab_size, draft_embd)
        
        # We reuse the Block architecture but with smaller dims
        # Note: In a real scenario, DraftModel would have its own specific smaller Block
        # For simplicity, we implement a simple one here
        self.layers = nn.ModuleList([
            nn.TransformerEncoderLayer(
                d_model=draft_embd,
                nhead=4,
                dim_feedforward=draft_embd * 4,
                dropout=0.0,
                activation='gelu',
                batch_first=True,
                norm_first=True
            ) for _ in range(draft_layers)
        ])
        
        self.ln_f = nn.LayerNorm(draft_embd)
        self.lm_head = nn.Linear(draft_embd, config.vocab_size, bias=False)

    def forward(self, idx):
        B, T = idx.shape
        x = self.token_embedding_table(idx)
        
        for layer in self.layers:
            x = layer(x)
            
        x = self.ln_f(x)
        logits = self.lm_head(x)
        return logits, None # Matches GPT interface for sampling