import os
import sys

import torch
import torch.nn as nn
from torch.nn import functional as F

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

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

        self.flatten = nn.Flatten()
        self.lin1 = nn.LazyLinear(features).to(device)
        self.act1 = nn.GELU().to(device)
        self.lin2 = nn.Linear(features, features * 15 * 20).to(device)
        self.act2 = nn.Tanh().to(device)
    

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        x: torch.Tensor = self.flatten(input)
        x= self.lin1(x)
        x = self.act1(x)
        x = self.lin2(x)
        x = self.act2(x)
        x = x.view(x.shape[0],self.features,15,20)

        return x

class DepthEstimation(nn.Module):
    def __init__(self, view_features=128, device='cuda', use_skip=False):
        super().__init__()

        
        block_setting = [
            CNBlockConfig(96, 192, 3),
            CNBlockConfig(192, 384, 3),
            CNBlockConfig(384, 768, 9),
            CNBlockConfig(768, None, 3),
        ]
        
        block_setting_reverse = [
            CNBlockConfig(768 + view_features, 384, 3),
            CNBlockConfig(384 * ( 2 if use_skip else 1), 192, 9),
            CNBlockConfig(192 * ( 2 if use_skip else 1), 96, 3),
            CNBlockConfig(96 * ( 2 if use_skip else 1), 27, 3),
            CNBlockConfig(27, 9, 3),
            CNBlockConfig(9, None, 3),
        ]

        self.convnext = ConvNeXt(block_setting, device=device)
        self.flatten = nn.Flatten()
        self.decoder = ConvNeXt(block_setting_reverse, use_skip=use_skip, reversed=True, device=device)

        self.view_encoder = ViewEncoder(view_features)


    def forward(self, input) -> torch.Tensor:
        image, view = input
        x, features = self.convnext(image)
        view_features = self.view_encoder(view)
        # x = self.flatten(x)
        x = self.decoder([torch.concat((x, view_features), axis=1),features])
        
        return x