import torch.nn as nn
from .resnet_backbone import build_backbone
from .modules import SpatialSoftmax
import numpy as np
import torch

class ResNetEncoder(nn.Module):
    def __init__(self, len_cameras=3, use_film=False):
        super().__init__()
        backbones = []
        pools = []
        linears = []
        img_fea_dim = stsm_num_kp = 512
        self.len_cameras = len_cameras
        self.use_film = use_film
        self.backbone_name = 'resnet50'
        for _ in range(len_cameras):
            backbone = build_backbone({"backbone": "resnet50"})
            backbones.append(backbone)

            input_shape = [2048, 8, 10]

            pools.append(
                nn.Sequential(
                    SpatialSoftmax(**{'input_shape': input_shape, 'num_kp': stsm_num_kp, 'temperature': 1.0,
                                      'learnable_temperature': False, 'noise_std': 0.0}),
                    nn.Flatten(start_dim=1, end_dim=-1)
                )
            )
            linears.append(
                nn.Sequential(
                    nn.Linear(int(np.prod([stsm_num_kp, 2])), stsm_num_kp),
                    nn.ReLU(),
                    nn.Linear(stsm_num_kp, img_fea_dim)
                )
            )

        self.backbones = nn.ModuleList(backbones)
        self.pools = nn.ModuleList(pools)
        self.linears = nn.ModuleList(linears)
        self.projection = nn.Sequential(
            nn.Linear(len_cameras * 512, 768),
            nn.ReLU(),
            nn.Linear(768, 768),
        )

    def forward(self, images, lang_embed=None):
        all_cam_features = []
        images = (images / 255.0).to(torch.bfloat16)
        for cam_id in range(self.len_cameras):
            if self.use_film and lang_embed is not None:
                cur_img = images[:, cam_id]

                # if self.color_randomizer is not None:
                #     cur_img = self.color_randomizer._forward_in(cur_img)


                features = self.backbones[cam_id](cur_img, lang_embed)

            else:
                cur_img = images[:, cam_id]
                # if self.color_randomizer is not None:
                #     cur_img = self.color_randomizer._forward_in(cur_img)
                features = self.backbones[cam_id](cur_img)

            pool_features = self.pools[cam_id](
                features).to(torch.bfloat16)
            out_features = self.linears[cam_id](pool_features)

            all_cam_features.append(out_features)
        obs_cond = torch.cat(all_cam_features, dim=1)
        obs_cond = self.projection(obs_cond)
        return obs_cond
