import os
import sys

import torch.nn as nn

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

from infer.lib.predictors.DJCM.utils import init_bn
from infer.lib.predictors.DJCM.decoder import PE_Decoder, SVS_Decoder
from infer.lib.predictors.DJCM.encoder import ResEncoderBlock, Encoder

class LatentBlocks(nn.Module):
    def __init__(
        self, 
        n_blocks, 
        latent_layers
    ):
        super(LatentBlocks, self).__init__()
        self.latent_blocks = nn.ModuleList([
            ResEncoderBlock(
                384, 
                384, 
                n_blocks, 
                None
            ) 
            for _ in range(latent_layers)
        ])

    def forward(self, x):
        for layer in self.latent_blocks:
            x = layer(x)

        return x

class DJCMM(nn.Module):
    def __init__(
        self, 
        in_channels, 
        n_blocks, 
        latent_layers, 
        svs=False, 
        window_length=1024, 
        n_class=360
    ):
        super(DJCMM, self).__init__()
        self.bn = nn.BatchNorm2d(
            window_length // 2 + 1, 
            momentum=0.01
        )
        self.pe_encoder = Encoder(
            in_channels, 
            n_blocks
        )
        self.pe_latent = LatentBlocks(
            n_blocks, 
            latent_layers
        )
        self.pe_decoder = PE_Decoder(
            n_blocks, 
            window_length=window_length, 
            n_class=n_class
        )

        self.svs = svs

        if svs:
            self.svs_encoder = Encoder(
                in_channels, 
                n_blocks
            )
            self.svs_latent = LatentBlocks(
                n_blocks, 
                latent_layers
            )
            self.svs_decoder = SVS_Decoder(
                in_channels, 
                n_blocks
            )

        init_bn(self.bn)

    def spec(self, x, spec_m):
        bs, c, time_steps, freqs_steps = x.shape
        x = x.reshape(bs, c // 4, 4, time_steps, freqs_steps)

        mask_spec = x[:, :, 0, :, :].sigmoid()
        linear_spec = x[:, :, 3, :, :]

        out_spec = (
            spec_m.detach() * mask_spec + linear_spec
        ).relu()

        return out_spec

    def forward(self, spec):
        x = self.bn(
            spec.transpose(1, 3)
        ).transpose(1, 3)[..., :-1]

        if self.svs:
            x, concat_tensors = self.svs_encoder(x)

            x = self.svs_decoder(
                self.svs_latent(x), 
                concat_tensors
            )

            x = self.spec(
                nn.functional.pad(x, pad=(0, 1)), 
                spec
            )[..., :-1]

        x, concat_tensors = self.pe_encoder(x)

        pe_out = self.pe_decoder(
            self.pe_latent(x), 
            concat_tensors
        )

        return pe_out