import os
import sys

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

import torch
import numpy as np
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__)) + '/..')



class ViewEncoder(nn.Module):
    def __init__(self, bottleneck_size):
        super(ViewEncoder, self).__init__()
        self.bottleneck_size = bottleneck_size
        
        self.encoder = nn.ModuleList([
            nn.Sequential(
                nn.Linear(3*3, 64),
                nn.BatchNorm1d(64),
                nn.ReLU(),),
            nn.Sequential(
                nn.Linear(64, 128),
                nn.BatchNorm1d(128),
                nn.ReLU(),),
            nn.Sequential(
                nn.Linear(128, 256),
                nn.BatchNorm1d(256),
                nn.ReLU(),),
            nn.Sequential(
                nn.Linear(256, 512),
                nn.BatchNorm1d(512),
                nn.ReLU(),),
            nn.Sequential(
                nn.Linear(512, bottleneck_size),
                nn.BatchNorm1d(bottleneck_size),)
        ])

        # Decoder
        self.decoder = nn.ModuleList([
            nn.Sequential(
                nn.Linear(bottleneck_size, 512),
                nn.BatchNorm1d(512),
                nn.ReLU(),),
            nn.Sequential(
                nn.Linear(512*2, 256),
                nn.BatchNorm1d(256),
                nn.ReLU(),),
            nn.Sequential(
                nn.Linear(256*2, 128),
                nn.BatchNorm1d(128),
                nn.ReLU(),),
            nn.Sequential(
                nn.Linear(128*2, 64),
                nn.BatchNorm1d(64),
                nn.ReLU(),),
            nn.Sequential(
                nn.Linear(64*2, 3*3),
                nn.Tanh(),),
        ])
    
    def forward(self, x):
        # Encode
        f1 = self.encoder[0](x.view(-1, 3*3))
        f2 = self.encoder[1](f1)
        f3 = self.encoder[2](f2)
        f4 = self.encoder[3](f3)
        matrix = self.encoder[4](f4)

        d1 = self.decoder[0](matrix)
        d2 = self.decoder[1](torch.cat((d1, f4), dim=1))
        d3 = self.decoder[2](torch.cat((d2, f3), dim=1))
        d4 = self.decoder[3](torch.cat((d3, f2), dim=1))
        x = self.decoder[4](torch.cat((d4, f1), dim=1))
        
        x = x * np.pi
        return x.view(-1, 3, 3), matrix
        
        # # Decode
        # x = self.decoder(matrix)
        # x = x * np.pi
        # return x.view(-1, 3, 3), matrix
    

if __name__ == '__main__':
    # test the model
    model = ViewEncoder(bottleneck_size=128).to('cuda')
    # 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, 3, 3)).to('cuda')
    output = model(x)

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

    pass