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 CNBlock, LayerNorm2d
from models.pose_ae import PoseAE
from utils.checkpoints import load_model


class FeatureManipulator(nn.Module):
    def __init__(
        self,
        features: int = 256,
        device: str = 'cuda',
    ) -> None:
        super().__init__()

        norm_layer = partial(LayerNorm2d, eps=1e-6)
        self.last = nn.Sequential(
                        norm_layer(features + 1),
                        nn.Conv2d(features + 1, features, kernel_size=3, padding='same'),
                    ).to(device)
        
        stage: List[nn.Module] = []
        for _ in range(10):
            stage.append(CNBlock(features + 1, 1.0, 0))
        self.stage = nn.Sequential(*stage).to(device)

    def _forward_impl(self, x) -> Tensor:
        x = self.stage(x)
        x = self.last(x)
        return x

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


class PoseManipulator(nn.Module):
    def __init__(self, features=256, device='cuda', ae_checpoint=None):
        super().__init__()

        self.pose_ae: PoseAE = PoseAE(features=features, device=device)
        iteration = load_model(self.pose_ae.encoder, ae_checpoint + '_enc')
        load_model(self.pose_ae.decoder, ae_checpoint + '_dec')
        for param in self.pose_ae.encoder.parameters():
            param.requires_grad = False
        for param in self.pose_ae.decoder.parameters():
            param.requires_grad = False

        print(f'Loaded pose_ae checkpoint at iteration: {iteration}')

        self.manipulator = FeatureManipulator(features=features, device=device)


    def forward(self, input) -> torch.Tensor:
        mask, pose_in, mask_out, pose_out, angle = input
        original_features, _ = self.pose_ae.encoder(torch.cat([mask, pose_in], dim=1))
        expected_features, _ = self.pose_ae.encoder(torch.cat([mask_out, pose_out], dim=1))

        tiled_angle = torch.ones_like(original_features[:, :1, :, :]) * angle[:,:,None,None]
        modified_features = self.manipulator(torch.cat((original_features, tiled_angle), dim=1))
        pose_pred, mask_pred = self.pose_ae.decoder(modified_features)

        return pose_pred, mask_pred, expected_features, modified_features


if __name__ == '__main__':
    pass