import torch
from torch import nn
from models.convnext_orig import ConvNeXt, CNBlock, CNBlockConfig
from typing import Any, Callable, List, Optional, Sequence


class DirectModel(nn.Module):
    
    def __init__(self, block_setting: List[CNBlockConfig], kernel_size: int = 7, input_features: int = 16, device: str = 'cuda'):
        super(DirectModel, self).__init__()

        self.encoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, input_features=input_features, device=device)
        self.decoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, device=device, reversed=True)

        layers = [CNBlock(block_setting[-1].input_channels, kernel_size=kernel_size) for _ in range(5)]
        self.mean_branch = nn.Sequential(*layers).to(device)
        layers = [CNBlock(block_setting[-1].input_channels, kernel_size=kernel_size) for _ in range(5)]
        self.var_branch = nn.Sequential(*layers).to(device)

        self.flatten_mean = nn.Flatten()
        self.flatten_var = nn.Flatten()


    def forward(self, inputs):
        body, template = inputs

        x = self.encoder(x)

        return mean