import torch
import torch.nn as nn
from transformers import PreTrainedModel
from .configuration_diffusion import DiffusionConfig

class DiffusionTransformer(nn.Module):
    def __init__(self, vocab_size, d_model=256, n_heads=8, n_layers=6, dropout=0.1, max_len=140, num_steps=1000):
        super().__init__()
        self.vocab_size = vocab_size
        self.max_len = max_len
        self.num_steps = num_steps

        self.token_embed = nn.Embedding(vocab_size, d_model)
        self.pos_embed = nn.Embedding(max_len, d_model)
        self.time_embed = nn.Embedding(num_steps, d_model)

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=n_heads,
            dim_feedforward=4 * d_model,
            dropout=dropout,
            batch_first=True
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=n_layers)
        self.output_head = nn.Linear(d_model, vocab_size)

    def forward(self, input_ids, t=None):
        B, T = input_ids.shape
        device = input_ids.device
        positions = torch.arange(T, device=device).unsqueeze(0).expand(B, T)
        tok_emb = self.token_embed(input_ids)
        pos_emb = self.pos_embed(positions)

        if t is not None:
            t = t.to(device)
            t_ids = (t * (self.num_steps - 1)).long()
            t_emb = self.time_embed(t_ids).unsqueeze(1)
        else:
            t_emb = torch.zeros_like(tok_emb[:, :1])

        x = tok_emb + pos_emb + t_emb
        x = self.transformer(x)
        logits = self.output_head(x)
        return logits

class HuggingFaceDiffusionModel(PreTrainedModel):
    config_class = DiffusionConfig

    def __init__(self, config):
        super().__init__(config)
        self.model = DiffusionTransformer(
            vocab_size=config.vocab_size,
            d_model=config.d_model,
            n_heads=config.n_heads,
            n_layers=config.n_layers,
            dropout=config.dropout,
            max_len=config.max_len,
            num_steps=config.num_steps
        )

    def forward(self, input_ids, t=None):
        return self.model(input_ids, t)