import os
import sys

import torch.nn as nn

sys.path.append(os.getcwd())

from main.library.predictors.DJCM.utils import ResConvBlock

class ResEncoderBlock(nn.Module):
    def __init__(
        self, 
        in_channels, 
        out_channels, 
        n_blocks, 
        kernel_size
    ):
        super(ResEncoderBlock, self).__init__()
        self.conv = nn.ModuleList([
            ResConvBlock(
                in_channels, 
                out_channels
            )
        ])

        for _ in range(n_blocks - 1):
            self.conv.append(
                ResConvBlock(
                    out_channels, 
                    out_channels
                )
            )

        self.pool = nn.MaxPool2d(kernel_size) if kernel_size is not None else None

    def forward(self, x):
        for each_layer in self.conv:
            x = each_layer(x)

        if self.pool is not None: return x, self.pool(x)
        return x

class Encoder(nn.Module):
    def __init__(
        self, 
        in_channels, 
        n_blocks
    ):
        super(Encoder, self).__init__()
        self.en_blocks = nn.ModuleList([
            ResEncoderBlock(
                in_channels, 
                32, 
                n_blocks, 
                (1, 2)
            ), 
            ResEncoderBlock(
                32, 
                64, 
                n_blocks, 
                (1, 2)
            ), 
            ResEncoderBlock(
                64, 
                128, 
                n_blocks, 
                (1, 2)
            ), 
            ResEncoderBlock(
                128, 
                256, 
                n_blocks, 
                (1, 2)
            ), 
            ResEncoderBlock(
                256, 
                384, 
                n_blocks, 
                (1, 2)
            ), 
            ResEncoderBlock(
                384, 
                384, 
                n_blocks, 
                (1, 2)
            )
        ])

    def forward(self, x):
        concat_tensors = []

        for layer in self.en_blocks:
            _, x = layer(x)
            concat_tensors.append(_)

        return x, concat_tensors