import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from layers.STEmbedding import SNIP, AttentionPrompt

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 Model(nn.Module):
    """
    Normalization-Linear
    """
    def __init__(self, args):
        super(Model, self).__init__()
        self.seq_len = args.input_length
        self.pred_len = args.predict_length
        
        # Decompsition Kernel Size
        kernel_size = 3
        self.decompsition = series_decomp(kernel_size)
        self.individual = False
        self.channels = None # not used

        if self.individual:
            self.Linear_Seasonal = nn.ModuleList()
            self.Linear_Trend = nn.ModuleList()
            
            for i in range(self.channels):
                self.Linear_Seasonal.append(nn.Linear(self.seq_len,self.pred_len))
                self.Linear_Trend.append(nn.Linear(self.seq_len,self.pred_len))

                # Use this two lines if you want to visualize the weights
                # self.Linear_Seasonal[i].weight = nn.Parameter((1/self.seq_len)*torch.ones([self.pred_len,self.seq_len]))
                # self.Linear_Trend[i].weight = nn.Parameter((1/self.seq_len)*torch.ones([self.pred_len,self.seq_len]))
        else:
            self.Linear_Seasonal = nn.Linear(self.seq_len,self.pred_len)
            self.Linear_Trend = nn.Linear(self.seq_len,self.pred_len)
            
            # Use this two lines if you want to visualize the weights
            # self.Linear_Seasonal.weight = nn.Parameter((1/self.seq_len)*torch.ones([self.pred_len,self.seq_len]))
            # self.Linear_Trend.weight = nn.Parameter((1/self.seq_len)*torch.ones([self.pred_len,self.seq_len]))


        self.prompt_type = args.prompt_type
        self.static_feats_dim_list = args.static_feats_dim_list
        self.static_func_type = args.static_func_type
        self.dynamic_func_type = args.dynamic_func_type
        self.se_emb_dropout = args.se_emb_dropout
        self.support_len = args.support_len
        self.hid_dim = args.hid_dim
        
        if self.prompt_type=='SNIP':
            self.getSNIP = SNIP(self.static_feats_dim_list, self.seq_len, support_len=self.support_len, 
                                static_func_type = self.static_func_type, dynamic_func_type= self.dynamic_func_type,
                                dropout_rate=self.se_emb_dropout)
        
        elif self.prompt_type=='AttPrompt':
            self.getAttPrompt = AttentionPrompt(args.M, self.seq_len)


    def forward(self, x, mode=None):
        x, x_mark, y_mark, aux_data = x
        x = x.squeeze(-1) #BTN

        # x: [Batch, Input length, Channel]
        seasonal_init, trend_init = self.decompsition(x)
        seasonal_init, trend_init = seasonal_init.permute(0,2,1), trend_init.permute(0,2,1) # B N T


        enc_out = x.permute(0,2,1) # B N T
        if self.prompt_type == 'SNIP':
            adj_mx = aux_data['adj_mx']
            period_feats = aux_data['period_feats']
            stci_feats = aux_data['stci_feats']
            structure_feats = aux_data['structure_feats']
            input_static_feat_list = []
            if self.static_feats_dim_list[0] >0:
                input_static_feat_list.append(period_feats)
            if self.static_feats_dim_list[1] >0:
                input_static_feat_list.append(stci_feats)
            if self.static_feats_dim_list[2] >0:
                input_static_feat_list.append(structure_feats)
                
            prompt_emb = self.getSNIP(enc_out.unsqueeze(1), input_static_feat_list, adj_mx)
            prompt_emb = prompt_emb.squeeze(1) # B N T
            seasonal_init = seasonal_init + prompt_emb
            trend_init = trend_init + prompt_emb
        elif self.prompt_type == 'AttPrompt':
            prompt_emb = self.getAttPrompt(enc_out.unsqueeze(1))
            prompt_emb = prompt_emb.squeeze(1)
            enc_out = enc_out + prompt_emb
            trend_init = trend_init + prompt_emb

        if self.individual:
            seasonal_output = torch.zeros([seasonal_init.size(0),seasonal_init.size(1),self.pred_len],dtype=seasonal_init.dtype).to(seasonal_init.device)
            trend_output = torch.zeros([trend_init.size(0),trend_init.size(1),self.pred_len],dtype=trend_init.dtype).to(trend_init.device)
            for i in range(self.channels):
                seasonal_output[:,i,:] = self.Linear_Seasonal[i](seasonal_init[:,i,:])
                trend_output[:,i,:] = self.Linear_Trend[i](trend_init[:,i,:])
        else:
            seasonal_output = self.Linear_Seasonal(seasonal_init)
            trend_output = self.Linear_Trend(trend_init)

        predy = seasonal_output + trend_output
        predy = predy.permute(0,2,1).unsqueeze(-1)
        
        return predy, None