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 DepthEstimation(nn.Module):
    def __init__(self, device='cuda'):
        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, 384, 3),
            CNBlockConfig(384, 192, 9),
            CNBlockConfig(192, 96, 3),
            CNBlockConfig(96, 27, 3),
            CNBlockConfig(27, 9, 3),
            CNBlockConfig(9, None, 3),
        ]

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


    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x, features = self.convnext(x)
        # x = self.flatten(x)
        x = self.decoder([x, features])
        
        return x