import os
import sys

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

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 ConvNeXtReversePose(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)

        stage: List[nn.Module] = []
        for _ in range(cnf.num_layers):
            stage.append(block(block_setting[-1].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))
        pose = pose * (mask > 0.5).to(torch.float32)

        return pose, mask

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


class PoseAE(nn.Module):
    def __init__(self, features=256, 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),
        ]

        self.encoder = ConvNeXt(block_setting, input_features=4, device=device)
        self.decoder = ConvNeXtReversePose(block_setting_reverse, device=device)

        self.flatten = nn.Flatten()


    def forward(self, input) -> torch.Tensor:
        mask, pose_in = input
        x, _ = self.encoder(torch.cat([mask, pose_in], dim=1))

        pose_out, mask_out = self.decoder(x)

        return pose_out, mask_out


if __name__ == '__main__':
    # test the model
    model = PoseAE()
    # 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')
    output = model((x, y))



    pass