import torch
import torch.nn as nn
from huggingface_hub import PyTorchModelHubMixin

class DiffusionTextModel(nn.Module, PyTorchModelHubMixin):
    def __init__(self, vocab_size, max_seq_len, max_time_steps,
                 embed_dim=128, n_layers=4, n_heads=4):
        super().__init__()
        self.config = {
            "vocab_size": vocab_size,
            "max_seq_len": max_seq_len,
            "max_time_steps": max_time_steps,
            "embed_dim": embed_dim,
            "n_layers": n_layers,
            "n_heads": n_heads
        }

        self.token_emb = nn.Embedding(vocab_size, embed_dim)
        self.pos_emb   = nn.Embedding(max_seq_len, embed_dim)
        self.time_emb  = nn.Embedding(max_time_steps+1, embed_dim)

        enc_layer = nn.TransformerEncoderLayer(
            d_model=embed_dim, nhead=n_heads,
            dim_feedforward=4*embed_dim, activation="gelu"
        )
        self.transformer = nn.TransformerEncoder(enc_layer, num_layers=n_layers)
        self.out = nn.Linear(embed_dim, vocab_size)

    def forward(self, x, t):
        B, L = x.shape
        tok = self.token_emb(x)
        pos = self.pos_emb(torch.arange(L, device=x.device).unsqueeze(0).expand(B, L))
        tim = self.time_emb(t).unsqueeze(1).expand(B, L, -1)
        h = tok + pos + tim
        h = self.transformer(h.transpose(0,1)).transpose(0,1)
        return self.out(h)