import torch, torch.nn as nn, math

class M(nn.Module):
    def __init__(self, d=128, L=3, H=2, nc=10):
        super().__init__()
        self.pe = nn.Conv2d(3, d, 16, 16)
        self.cls = nn.Parameter(torch.randn(1,1,d)*.02)
        self.pos = nn.Parameter(torch.randn(1,197,d)*.02)
        self.blks = nn.ModuleList([nn.TransformerEncoderLayer(d, H, d*4, .1, activation='gelu', batch_first=True, norm_first=True) for _ in range(L)])
        self.ln = nn.LayerNorm(d)
        self.te = nn.Embedding(30522, d, padding_idx=0)
        self.tpos = nn.Parameter(torch.randn(1,128,d)*.02)
        self.tblks = nn.ModuleList([nn.TransformerEncoderLayer(d, H, d*4, .1, activation='gelu', batch_first=True, norm_first=True) for _ in range(L)])
        self.tln = nn.LayerNorm(d)
        self.fuse = nn.ModuleList([nn.TransformerEncoderLayer(d, H, d*4, .1, activation='gelu', batch_first=True, norm_first=True) for _ in range(2)])
        self.fln = nn.LayerNorm(d)
        self.head = nn.Sequential(nn.Linear(d,d), nn.GELU(approximate='quick'), nn.Dropout(.1), nn.Linear(d,nc))
        self._init()

    def _init(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.trunc_normal_(m.weight, std=0.02)
                if m.bias is not None: nn.init.zeros_(m.bias)

    def enc_img(self, x):
        x = self.pe(x).flatten(2).transpose(1,2)
        x = torch.cat([self.cls.expand(x.size(0),-1,-1), x], 1)
        x = x + self.pos
        for b in self.blks: x = b(x)
        return self.ln(x)

    def enc_txt(self, ids, mask=None):
        x = self.te(ids) + self.tpos[:, :ids.size(1)]
        m = (mask == 0) if mask is not None else None
        for b in self.tblks: x = b(x, src_key_padding_mask=m)
        return self.tln(x)

    def forward(self, img, ids, mask=None, lbl=None):
        fi = self.enc_img(img)
        ft = self.enc_txt(ids, mask)
        x = ft
        for f in self.fuse: x = f(x)
        x = self.fln(x[:, 0])
        logits = self.head(x)
        loss = nn.functional.cross_entropy(logits, lbl) if lbl is not None else None
        return {'logits': logits, 'loss': loss}

if __name__ == '__main__':
    m = M()
    print(f'Params: {sum(p.numel() for p in m.parameters()):,}')
    o = m(torch.randn(2,3,224,224), torch.randint(0,30522,(2,128)), torch.ones(2,128), torch.tensor([0,1]))
    print(o['logits'].shape, o['loss'].item())