import torch
from torch import nn
from torch import Tensor

from fedot_ind.core.models.nn.network_modules.layers.linear_layers import FlattenHead
from fedot_ind.core.models.nn.network_modules.layers.special import _TSTiEncoder, RevIN


class _PatchTST_backbone(nn.Module):
    def __init__(self, input_dim, seq_len, pred_dim, patch_len, stride,
                 n_layers=3, d_model=128, n_heads=16, d_k=None, d_v=None,
                 d_ff=256, norm='BatchNorm', attn_dropout=0., dropout=0.,
                 act="GELU", res_attention=True, pre_norm=False, store_attn=False,
                 padding_patch=True, individual=False,
                 revin=True, affine=True, subtract_last=False, preprocess_to_lagged=False):

        super().__init__()

        # RevIn
        self.revin = revin
        self.revin_layer = RevIN(
            input_dim, affine=affine, subtract_last=subtract_last)

        # Patching
        self.patch_len = patch_len
        self.stride = stride
        self.padding_patch = padding_patch
        patch_num = int((seq_len - patch_len) / stride + 1) + 1
        self.patch_num = patch_num
        self.padding_patch_layer = nn.ReplicationPad1d(
            (stride, 0))  # original padding at the end
        self.preprocess_to_lagged = preprocess_to_lagged

        # Unfold
        self.unfold = nn.Unfold(kernel_size=(1, patch_len), stride=stride)
        self.patch_len = patch_len

        # Backbone
        self.backbone = _TSTiEncoder(input_dim=input_dim, patch_num=patch_num, patch_len=patch_len,
                                     n_layers=n_layers, d_model=d_model, n_heads=n_heads, d_k=d_k, d_v=d_v, d_ff=d_ff,
                                     attn_dropout=attn_dropout, dropout=dropout, act=act,
                                     res_attention=res_attention, pre_norm=pre_norm, store_attn=store_attn)

        # Head
        self.head_nf = d_model * patch_num
        self.n_vars = input_dim
        self.individual = individual
        self.head = FlattenHead(
            self.individual, self.n_vars, self.head_nf, pred_dim)

    def forward(self, z: Tensor):
        """
        Args:
            z: [batch_size x input_dim x seq_len]
        """
        # norm
        if self.revin:
            z = self.revin_layer(z, torch.tensor(True, dtype=torch.bool))

        b, c, s = z.size()
        z = z.reshape(-1, 1, 1, s).permute(0, 1, 3, 2)

        # model
        z, scores = self.backbone(z)  # z: [bs x nvars x d_model x patch_num]
        z = self.head(z)  # z: [bs x nvars x pred_dim]

        # denorm
        if self.revin:
            z = self.revin_layer(z, torch.tensor(False, dtype=torch.bool))
        return z
