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

class TransformerBlock(nn.Module):
    def __init__(self, d, nhead, mlp=4):
        super().__init__()
        self.ln1 = nn.LayerNorm(d)
        self.attn = nn.MultiheadAttention(d, nhead, batch_first=True)
        self.ln2 = nn.LayerNorm(d)
        self.ff = nn.Sequential(
            nn.Linear(d, mlp*d), nn.GELU(), nn.Linear(mlp*d, d)
        )

    def forward(self, x, attn_mask=None):
        h = self.ln1(x)
        h, _ = self.attn(h, h, h, attn_mask=attn_mask, need_weights=False)
        x = x + h
        h = self.ln2(x)
        h = self.ff(h)
        return x + h

class ActionCondTransformer(nn.Module):
    def __init__(self, num_codes, code_dim, n_tokens, action_dim, d_model=512, nlayer=6, nhead=8):
        super().__init__()
        self.token_emb = nn.Embedding(num_codes, d_model)
        self.pos = nn.Parameter(torch.zeros(1, n_tokens, d_model))
        self.act_proj = nn.Linear(action_dim, d_model)
        self.in_proj = nn.Linear(d_model, d_model)
        self.blocks = nn.ModuleList([TransformerBlock(d_model, nhead) for _ in range(nlayer)])
        self.ln = nn.LayerNorm(d_model)
        self.logits_head = nn.Linear(d_model, num_codes)
        self.reward_head = nn.Linear(d_model, 1)
        self.done_head = nn.Linear(d_model, 1)

    def forward(self, idx_t, act_t):
        B, N = idx_t.shape
        x = self.token_emb(idx_t) + self.pos[:, :N]
        a = self.act_proj(act_t).unsqueeze(1)
        x = x + a
        x = self.in_proj(x)
        for blk in self.blocks:
            x = blk(x, attn_mask=None)
        h = self.ln(x)
        logits = self.logits_head(h)
        pooled = h.mean(1)
        rew = self.reward_head(pooled).squeeze(1)
        done = self.done_head(pooled).squeeze(1)
        return logits, rew, done
