import torch
from torch import nn
from custom_transformer_layer import SelfTransformerEncoderBlock, CrossTransformerEncoderBlock



class VisionPredictor(nn.Module):
    def __init__(
        self,
        input_dim=4096,
        d_model=512,
        nhead=8,
        dim_feedforward=2048,
        num_layers=2,
    ):
        super().__init__()
        self.input_dim = input_dim
        self.d_model = d_model
        self.nhead = nhead
        self.dim_feedforward = dim_feedforward
        self.num_layers = num_layers

        self.down_mlp = nn.Sequential(
            nn.LayerNorm(self.input_dim),
            nn.Linear(self.input_dim, self.d_model),
            nn.GELU(),
        )
        self.transformer = nn.Sequential(
            *[
                SelfTransformerEncoderBlock(
                    dim=self.d_model,
                    num_heads=self.nhead,
                    mlp_ratio=self.dim_feedforward / self.d_model,
                )
                for _ in range(self.num_layers)
            ]
        )
        self.output_mlp = nn.Sequential(
            nn.Linear(self.d_model, self.d_model // 2),
            nn.GELU(),
            nn.Linear(self.d_model // 2, self.d_model // 4),
            nn.GELU(),
            nn.Linear(self.d_model // 4, 2),
            # nn.LogSoftmax(dim=-1),
        )

    def forward(self, x, image_policy) -> torch.FloatTensor:
        predict = []
        new_image_x = self.down_mlp(x)
        new_x = self.transformer(new_image_x * image_policy)
        B, N, C = new_x.size()
        local_x = new_x[:, :, : C // 2]
        global_x = (new_x[:, :, C // 2 :] * image_policy).sum(
            dim=1, keepdim=True
        ) / torch.sum(image_policy, dim=1, keepdim=True)
        new_x = torch.cat([local_x, global_x.expand(B, N, C // 2)], dim=-1)
        predict = self.output_mlp(new_x)
        return predict
    
