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

D = 128
D2 = 96

class DepthwiseSeparableConv(nn.Module):
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.dw = nn.Conv2d(in_ch, in_ch, 3, padding=1, groups=in_ch)
        self.pw = nn.Conv2d(in_ch, out_ch, 1)
    def forward(self, x):
        return self.pw(self.dw(x))

class AtomEncoder(nn.Module):
    def __init__(self, embed_dim=D):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(1, 16, 3, stride=2, padding=1), nn.GELU(),
            nn.Conv2d(16, 32, 3, stride=2, padding=1), nn.GELU(),
            DepthwiseSeparableConv(32, 64), nn.GELU(),
            nn.AdaptiveAvgPool2d(1),
        )
        self.proj = nn.Linear(64, embed_dim)
    def forward(self, x):
        return self.proj(self.net(x).flatten(1))

class GridRelationalAttention(nn.Module):
    def __init__(self, dim=D, heads=4, n_relations=4):
        super().__init__()
        self.missing_token = nn.Parameter(torch.randn(dim) * 0.02)
        self.row_attn = nn.MultiheadAttention(dim, heads, batch_first=True)
        self.col_attn = nn.MultiheadAttention(dim, heads, batch_first=True)
        self.relation_queries = nn.Parameter(torch.randn(n_relations, dim) * 0.02)
        self.cross_attn = nn.MultiheadAttention(dim, heads, batch_first=True)
        self.rule_proj = nn.Linear(n_relations * dim, dim)
        self.n_relations, self.dim = n_relations, dim

    def forward(self, context_panels):
        B = context_panels.shape[0]
        missing = self.missing_token.expand(B, 1, self.dim)
        grid = torch.cat([context_panels, missing], dim=1)
        rows_idx = [[0,1,2],[3,4,5],[6,7,8]]
        cols_idx = [[0,3,6],[1,4,7],[2,5,8]]
        row_vecs = torch.cat([
            self.row_attn(grid[:, idx, :], grid[:, idx, :], grid[:, idx, :])[0].mean(1, keepdim=True)
            for idx in rows_idx
        ], dim=1)
        col_vecs = torch.cat([
            self.col_attn(grid[:, idx, :], grid[:, idx, :], grid[:, idx, :])[0].mean(1, keepdim=True)
            for idx in cols_idx
        ], dim=1)
        combined_ctx = torch.cat([row_vecs, col_vecs], dim=1)
        rq = self.relation_queries.unsqueeze(0).expand(B, -1, -1)
        rel_out, _ = self.cross_attn(rq, combined_ctx, combined_ctx)
        return self.rule_proj(rel_out.flatten(1))

class CandidateScorer(nn.Module):
    def __init__(self, dim=D):
        super().__init__()
        self.score_proj = nn.Linear(dim, dim)
    def forward(self, rule_embedding, candidate_embeds):
        proj_rule = self.score_proj(rule_embedding)
        return torch.einsum("bd,bkd->bk", proj_rule, candidate_embeds)

class GRAFT(nn.Module):
    def __init__(self, vocab_size, narrator_cls):
        super().__init__()
        self.encoder = AtomEncoder()
        self.gra = GridRelationalAttention()
        self.scorer = CandidateScorer()
        self.narrator = narrator_cls(vocab_size=vocab_size)