import torch
import torch.nn as nn
import torch.nn.functional as F
from utils.masking import TriangularCausalMask, ProbMask
from layers.Transformer_EncDec import Decoder, DecoderLayer, Encoder, EncoderLayer, ConvLayer
from layers.SelfAttention_Family import FullAttention, ProbAttention, AttentionLayer
from layers.Embed import DataEmbedding,DataEmbedding_only_timemodel,DataEmbedding_timemodel,DataEmbedding_wo_pos,DataEmbedding_wo_temp,DataEmbedding_wo_pos_temp
import numpy as np
import torch.fft


class moving_avg(nn.Module):
    """
    Moving average block to highlight the trend of time series
    """
    def __init__(self, kernel_size, stride):
        super(moving_avg, self).__init__()
        self.kernel_size = kernel_size
        self.avg = nn.AvgPool1d(kernel_size=kernel_size, stride=stride, padding=0)

    def forward(self, x):
        # padding on the both ends of time series
        front = x[:, 0:1, :].repeat(1, (self.kernel_size - 1) // 2, 1)
        end = x[:, -1:, :].repeat(1, (self.kernel_size - 1) // 2, 1)
        x = torch.cat([front, x, end], dim=1)
        x = self.avg(x.permute(0, 2, 1))
        x = x.permute(0, 2, 1)
        return x


class series_decomp(nn.Module):
    """
    Series decomposition block
    """
    def __init__(self, kernel_size):
        super(series_decomp, self).__init__()
        self.moving_avg = moving_avg(kernel_size, stride=1)

    def forward(self, x):
        moving_mean = self.moving_avg(x)
        res = x - moving_mean
        return res, moving_mean


class FeedForwardNetwork(nn.Module):
    def __init__(self, hidden_size, filter_size, dropout_rate=0.1):
        super(FeedForwardNetwork, self).__init__()

        self.layer1 = nn.Linear(hidden_size, filter_size)
        self.relu = nn.Sigmoid()

        self.dropout = nn.Dropout(dropout_rate)
        self.layer2 = nn.Linear(filter_size, hidden_size)

        self.initialize_weight(self.layer1)
        self.initialize_weight(self.layer2)

    def forward(self, x):
        x = self.layer1(x)
        x = self.relu(x)
        # x = self.dropout(x)
        x = self.layer2(x)
        return x

    def initialize_weight(self, x):
        nn.init.xavier_uniform_(x.weight)
        if x.bias is not None:
            nn.init.constant_(x.bias, 0)


class Model(nn.Module):
    """
    Informer with Propspare attention in O(LlogL) complexity
    """
    def __init__(self, configs):
        super(Model, self).__init__()
        self.seq_len = configs.seq_len
        self.pred_len = configs.pred_len
        self.enc_in = configs.enc_in
        self.c_out = configs.c_out
        self.d_model = configs.d_model
        self.n_model = configs.n_model
        self.t_model = configs.t_model
        self.conv_kernel = [1, 3, 4, 6]
        self.sample_kernel = [14, 14, 14, 14]
        self.output_attention = configs.output_attention
        self.device = 'cuda:0'
        self.d_layers = configs.d_layers
        self.dropout = 0.05

        # Embedding
        if configs.embed_type == 0:
            self.enc_embedding = DataEmbedding(configs.enc_in, configs.d_model, configs.embed, configs.freq,
                                            configs.dropout)
            self.dec_embedding = DataEmbedding(configs.dec_in, configs.d_model, configs.embed, configs.freq,
                                           configs.dropout)
        elif configs.embed_type == 1:
            # self.enc_embedding = DataEmbedding(configs.enc_in, configs.d_model, configs.embed, configs.freq,
            #                                         configs.dropout)
            # self.dec_embedding = DataEmbedding(configs.dec_in, configs.d_model, configs.embed, configs.freq,
            #                                         configs.dropout)
            self.enc_embedding = DataEmbedding_only_timemodel(configs.enc_in, configs.t_model, configs.d_model, configs.embed, configs.freq,configs.dropout)
            # self.enc_embedding = DataEmbedding_timemodel(configs.enc_in, configs.d_model, configs.embed, configs.freq,
            #                                    configs.dropout)
        elif configs.embed_type == 2:
            self.enc_embedding = DataEmbedding_wo_pos(configs.enc_in, configs.d_model, configs.embed, configs.freq,
                                                    configs.dropout)
            self.dec_embedding = DataEmbedding_wo_pos(configs.dec_in, configs.d_model, configs.embed, configs.freq,
                                                    configs.dropout)

        elif configs.embed_type == 3:
            self.enc_embedding = DataEmbedding_wo_temp(configs.enc_in, configs.d_model, configs.embed, configs.freq,
                                                    configs.dropout)
            self.dec_embedding = DataEmbedding_wo_temp(configs.dec_in, configs.d_model, configs.embed, configs.freq,
                                                    configs.dropout)
        elif configs.embed_type == 4:
            self.enc_embedding = DataEmbedding_wo_pos_temp(configs.enc_in, configs.d_model, configs.embed, configs.freq,
                                                    configs.dropout)
            self.dec_embedding = DataEmbedding_wo_pos_temp(configs.dec_in, configs.d_model, configs.embed, configs.freq,
                                                    configs.dropout)
        # if configs.freq == 't':
        #     self.Linear = nn.Linear(self.t_model * 5, self.pred_len)
        # else:
        #     self.Linear = nn.Linear(self.t_model * 4, self.pred_len)
        self.Linear = nn.Linear(42*self.t_model, self.pred_len)
        # self.NLinear = nn.Linear(self.seq_len, self.pred_len)

        self.nodevec = nn.Parameter(torch.randn(self.enc_in, 42).to(self.device), requires_grad=True).to(
            self.device)  # (num_nodes,10)

        self.conv1 = nn.ModuleList()
        self.conv2 = nn.ModuleList()
        self.sample_conv = nn.ModuleList()
        self.projection_sample = nn.ModuleList()
        self.projection_s = nn.ModuleList()

        # Add Conv
        # downsampling convolution: padding=i//2, stride=i
        for i in range(len(self.conv_kernel)):
            # self.conv1.append(nn.Conv1d(in_channels=1, out_channels=1,
            #                         kernel_size=self.conv_kernel[i], padding=self.conv_kernel[i] // 2, stride=1))
            self.conv2.append(nn.Conv1d(in_channels=1, out_channels=1,
                                    kernel_size=self.conv_kernel[i], padding=self.conv_kernel[i] // 2, stride=self.conv_kernel[i]))
            self.sample_conv.append(nn.Conv1d(in_channels=1, out_channels=self.t_model,
                                    dilation=24 // self.conv_kernel[i], kernel_size=14, padding=0, stride=1))

            # self.projection_s.append(nn.Linear(self.t_model * (24 // self.conv_kernel[i]), self.t_model))
#         self.projection_sample.append(nn.Linear(self.t_model * len(self.conv_kernel), 24))

#         self.tweight = nn.Parameter(torch.randn(self.pred_len).to(self.device), requires_grad=True).to(
#             self.device)


    def forward(self, x, x_mark_dec, y, y_mark_enc, x_fa):
        # x_enc  torch.Size([32, 336, 1])   x_mark_enc  torch.Size([32, 336, 4])
        # print("x_time ", x_time.shape)
        # x_time = self.enc_embedding(x, x_mark_dec)
        # y_time = self.enc_embedding(y, y_mark_enc)
        # x_out = self.Linear(x_time)
        # x = x - x_out

        seq_last = x[:,-1:,:].detach()
        x = x - seq_last
        x_enc = x.permute(0, 2, 1).unsqueeze(dim=-2).reshape(-1, 1, self.seq_len)   # [B*self.enc_in, 1, self.seq_len]

        n_sample = torch.tensor([], device=self.device)
        c = 0
        for i in range(len(self.conv_kernel)):
            c += (24 // self.conv_kernel[i])
            # sample_out = self.conv1[i](x_enc)   # [B*self.enc_in, 1, self.seq_len]
            sample_out = self.conv2[i](x_enc)   # [B*self.enc_in, 1, self.seq_len]
            
            sample_out = self.sample_conv[i](sample_out)  # [B*self.enc_in, self.t_model, 24]
            sample_out = sample_out[:,:,0:(24 // self.conv_kernel[i])]
            n_sample = torch.cat((n_sample, sample_out), dim=-1)
            
        # 拼接Node_vec
        dec_out = n_sample.permute(0, 2, 1).reshape(-1, self.enc_in, c, self.t_model)  # [B, 7, 42, self.t_model]
        # nodevec = torch.sigmoid(self.nodevec).unsqueeze(dim=0).unsqueeze(dim=-2)   # 竖着乘
        nodevec = torch.sigmoid(self.nodevec).unsqueeze(dim=0).unsqueeze(dim=-1)   # 横着乘
        dec_out = dec_out * nodevec 
        # for i in range(self.enc_in):
        #     dec_out = torch.cat((dec_out, self.nodevec[i:i+1,:].unsqueeze(dim=0).unsqueeze(dim=0)), dim=-2)

        # spatial Attention
        # print("mg_sample.shape", mg_sample.shape)
        # mg_att = torch.einsum('bctd,cd->bct', (dec_out, self.nodevec)).contiguous()
        # mg_score = nn.functional.softmax(mg_att, dim=-1)  # [B, 7, 42]
        # mg_s = torch.einsum('bctd,bct->bcd', (dec_out, mg_score))   # [B, 7, self.t_model]

        dec = self.Linear(dec_out.reshape(-1,self.enc_in, (c) * self.t_model)).permute(0, 2, 1)
        # temporal Aggregation
        # mg_pattern = torch.tensor([], device=self.device)
        # mg_node = torch.tensor([], device=self.device)
        # for i in range(self.pred_len // 24 + 1):
        #     if i == self.pred_len // 24:
        #         if self.pred_len % 24 > 0 :
        #             mg_pattern = torch.cat((mg_pattern, mg_s[:, 0:self.pred_len % 24, :]), dim=1)
        #             mg_node = torch.cat((mg_node, self.nodevec[:, 0:self.pred_len % 24]), dim=1)
        #         else:
        #             break
        #     else:
        #         mg_pattern = torch.cat((mg_pattern, mg_s), dim=1)   # [B, self.pred_len, self.enc_in]
        #         mg_node = torch.cat((mg_node, self.nodevec), dim=1)   # [self.enc_in, self.pred_len]
        #
        # # print("mg_pattern.shape", mg_pattern.shape)
        # # gate = torch.sigmoid(self.tweight.unsqueeze(dim=0).unsqueeze(dim=-1))
        # # one_vec = torch.ones_like(gate)
        #
        # gate = torch.sigmoid(self.tweight * mg_node)   # [self.enc_in, self.pred_len]
        # dec = mg_pattern * gate.permute(1,0)
        # dec = mg_pattern * self.tweight.unsqueeze(dim=0).unsqueeze(dim=-1)
        
        # gate = torch.matmul(self.nodevec, self.tweight)   # [node_num, self.pred_len]
        # dec = mg_pattern * gate.permute(1,0)
        # dec = dec + self.NLinear(x.permute(0,2,1)).permute(0,2,1)
        dec = dec + seq_last

        return dec