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
        
        # Encoder
        self.encoder = nn.Sequential(
            nn.Linear(3*3, 64),
            nn.ReLU(),
            nn.Linear(64, 128),
            nn.ReLU(),
            nn.Linear(128, 256),
            nn.ReLU(),
            nn.Linear(256, 512),
            nn.ReLU(),
            nn.Linear(512, bottleneck_size),
            nn.ReLU()
        )

        # Decoder
        self.decoder = nn.Sequential(
            nn.Linear(bottleneck_size, 512),
            nn.ReLU(),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, 3*3),
            nn.Tanh(),
        )
    
    def forward(self, x):
        # Encode
        matrix = self.encoder(x.view(-1, 3*3))
        
        # 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