import os
import sys

# os.environ["CUDA_VISIBLE_DEVICES"] = '3'
# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"

import torch
from torch import nn, Tensor
from torch.nn import functional as F
from functools import partial
from typing import Any, Callable, List, Optional, Sequence

torch.set_float32_matmul_precision('high')

sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
from models.convnext import ConvNeXt, CNBlockConfig, CNBlock, LayerNorm2d


class ConvNeXtReverseMaskPose(nn.Module):
    def __init__(
        self,
        block_setting: List[CNBlockConfig],
        stochastic_depth_prob: float = 0.0,
        layer_scale: float = 1e-6,
        block: Optional[Callable[..., nn.Module]] = None,
        norm_layer: Optional[Callable[..., nn.Module]] = None,
        device: str = 'cuda',
        use_skip=False,
    ) -> None:
        super().__init__()

        self.use_skip = use_skip

        if not block_setting:
            raise ValueError("The block_setting should not be empty")
        elif not (isinstance(block_setting, Sequence) and all([isinstance(s, CNBlockConfig) for s in block_setting])):
            raise TypeError("The block_setting should be List[CNBlockConfig]")

        if block is None:
            block = CNBlock

        if norm_layer is None:
            norm_layer = partial(LayerNorm2d, eps=1e-6)

        stage: List[nn.Module] = []
        for _ in range(3):
            stage.append(block(block_setting[0].input_channels, layer_scale, 0))
        self.stage_smooth = nn.Sequential(*stage).to(device)

        self.first = nn.Sequential(
                        nn.ConvTranspose2d(block_setting[0].input_channels, block_setting[0].out_channels, kernel_size=2, stride=2),
                        norm_layer(block_setting[0].out_channels),
                    ).to(device)


        total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
        stage_block_id = 0
        layers_list: List[nn.Module] = []

        for cnf in block_setting[1:]:
            if cnf.input_channels is not None:
                # Bottlenecks
                stage: List[nn.Module] = []
                for _ in range(cnf.num_layers):
                    # adjust stochastic depth probability based on the depth of the stage block
                    sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
                    stage.append(block(cnf.input_channels, layer_scale, sd_prob))
                    stage_block_id += 1
                layers_list.append(nn.Sequential(*stage).to(device))
            else:
                layers_list.append(nn.Identity().to(device))

            if cnf.out_channels is not None:
                # Upsampling
                layers_list.append(
                    nn.Sequential(
                        norm_layer(cnf.input_channels),
                        nn.ConvTranspose2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
                    ).to(device)
                )

        self.layers = nn.Sequential(*layers_list[:-1]).to(device)
        self.stage_pose = nn.Sequential(*layers_list[-1:]).to(device)

        # for cnf in block_setting[-1:]:
        #     # Bottlenecks
        #     stage: List[nn.Module] = []
        #     for _ in range(cnf.num_layers):
        #         stage.append(block(cnf.input_channels, layer_scale, 0))
        #     self.stage_mask = nn.Sequential(*stage).to(device)

        self.last = nn.Conv2d(
                    block_setting[-1].input_channels,
                    3,
                    kernel_size=3,
                    padding='same',
                    ).to(device)
        # self.last_mask = nn.Conv2d(
        #             block_setting[-1].input_channels,
        #             1,
        #             kernel_size=3,
        #             padding='same',
        #             ).to(device)


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


    def _forward_impl(self, x) -> Tensor:
        if self.use_skip:
            raise NotImplementedError
            input, features = x
        else:
            input = x
        x = self.stage_smooth(input)
        x = self.first(x)
        f = self.layers(x)
        x = self.stage_pose(f)
        pose = torch.sigmoid(self.last(x))

        # x = self.stage_mask(f)
        # mask = torch.sigmoid(self.last_mask(x))
        
        # return pose, mask
        return pose, torch.zeros((x.shape[0], 1, x.shape[2], x.shape[3]), device=x.device)

    def forward(self, x) -> Tensor:
        return self._forward_impl(x)


class PoseEncoder(nn.Module):
    def __init__(self, features=128, device='cuda'):
        super().__init__()

                
        block_setting = [
            CNBlockConfig(96, 192, 3),
            CNBlockConfig(192, 384, 3),
            CNBlockConfig(384, features, 9),
            CNBlockConfig(features, None, 3),
        ]

        block_setting_reverse = [
            CNBlockConfig(features, 384, 3),
            CNBlockConfig(384, 192, 9),
            CNBlockConfig(192, 96, 3),
            CNBlockConfig(96, 27, 3),
            CNBlockConfig(27, 9, 3),
            CNBlockConfig(9, None, 3),
        ]

        rotation_block_setting = [
            CNBlockConfig(64, 128, 3),
            CNBlockConfig(128, 128, 3),
            CNBlockConfig(128, 256, 9),
            CNBlockConfig(256, None, 3),
        ]

        rotation_block_setting_reverse = [
            CNBlockConfig(256, 128, 3),
            CNBlockConfig(128, 64, 9),
            CNBlockConfig(64, 32, 3),
            CNBlockConfig(32, 16, 3),
            CNBlockConfig(16, None, 3),
        ]

        self.encoder = ConvNeXt(block_setting, input_features=7, device=device)
        self.decoder = ConvNeXtReverseMaskPose(block_setting_reverse, device=device)
        self.rotation_encoder = ConvNeXt(rotation_block_setting, input_features=7, device=device)
        self.rotation_decoder = ConvNeXt(rotation_block_setting_reverse, reversed=True, device=device)

        self.flatten = nn.Flatten()


    def forward(self, input) -> torch.Tensor:
        mask, pose_in, pose_in_back, pose_out, mask_out = input
        features, _ = self.encoder(torch.cat([mask, pose_in, pose_in_back], dim=1))
        x = features.view(features.size(0), features.size(1), -1)
        y, _ = self.rotation_encoder(torch.cat([mask, pose_in, pose_out], dim=1))
        rotation = self.rotation_decoder(y)
        rotation = torch.squeeze(rotation, dim=1)
        rotation = torch.softmax(rotation, dim=1)

        # x = torch.einsum('bij,bkj->bik', x, rotation)
        x = torch.einsum('bij,bjk->bik', x, rotation)
        x = x.view(x.size(0), x.size(1), 16, 16)
        pose, mask_out_pred = self.decoder(x)
        pose = pose * mask_out
        # mask_out = torch.round(mask_out)
        # mask_out = (mask_out > 0.5).to(torch.int64)
        # pose = pose * (mask_out > 0.5).to(torch.float32)

        return pose, mask_out_pred, rotation
        # x = pose
        # return x, torch.zeros((x.shape[0], 1, x.shape[2], x.shape[3]), device=x.device), torch.zeros((x.shape[0], 256, 256), device=x.device)

class DumbPoseEncoder(nn.Module):
    def __init__(self, in_channels=7, out_channels=3, device='cuda'):
        super(DumbPoseEncoder, self).__init__()

        # Encoder
        self.encoder_conv1 = nn.Conv2d(in_channels, 64, kernel_size=3, padding=1)
        self.encoder_conv2 = nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1)
        self.encoder_conv3 = nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1)
        self.encoder_conv4 = nn.Conv2d(256, 512, kernel_size=3, padding=1)
        self.encoder_pool = nn.MaxPool2d(kernel_size=2, stride=2)

        # Decoder
        self.decoder_upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        self.decoder_conv1 = nn.Conv2d(512, 256, kernel_size=3, padding=1)
        self.decoder_upsample_2 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        self.decoder_conv2 = nn.Conv2d(256, 128, kernel_size=3, padding=1)
        self.decoder_upsample_3 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        self.decoder_conv3 = nn.Conv2d(128, 64, kernel_size=3, padding=1)
        self.decoder_conv4 = nn.Conv2d(64, out_channels, kernel_size=3, padding=1)

    def forward(self, x):
        # Encoder
        mask, pose_in, pose_out = x
        x = torch.cat([mask, pose_in, pose_out], dim=1)
        x1 = nn.functional.relu(self.encoder_conv1(x))
        x = self.encoder_pool(x1)
        x = nn.functional.relu(self.encoder_conv2(x))
        x = nn.functional.relu(self.encoder_conv3(x))
        x = nn.functional.relu(self.encoder_conv4(x))

        # Decoder
        x = self.decoder_upsample(x)
        x = nn.functional.relu(self.decoder_conv1(x))
        x = self.decoder_upsample_2(x)
        x = nn.functional.relu(self.decoder_conv2(x))
        x = self.decoder_upsample_3(x)
        x = nn.functional.relu(self.decoder_conv3(x))
        x = nn.functional.sigmoid(self.decoder_conv4(x))

        return x, torch.zeros((x.shape[0], 1, x.shape[2], x.shape[3]), device=x.device), torch.zeros((x.shape[0], 256, 256), device=x.device)


if __name__ == '__main__':
    # test the model
    model = PoseEncoder()
    # model = torch.compile(model)
    print(model)
    print(sum(p.numel() for p in model.parameters() if p.requires_grad))
    print(sum(p.numel() for p in model.parameters()))

    x = torch.rand((4, 1, 512, 512)).to('cuda')
    y = torch.rand((4, 3, 512, 512)).to('cuda')
    z = torch.rand((4, 3, 512, 512)).to('cuda')
    output = model((x, y, z))

    from torchviz import make_dot
    graph = make_dot(output, params=dict(model.named_parameters()))
    graph.render("model_graph", format="pdf")

    pass