import torch
import torch.nn as nn
from einops import rearrange, repeat
from layers.SelfAttention_Family import TwoStageAttentionLayer

#用于合并时间序列的不同片段
class SegMerging(nn.Module):
    #初始化方法，参数包含模型维度d_model、窗口大小win_size、以及正则化层默认为LayerNorm
    def __init__(self, d_model, win_size, norm_layer=nn.LayerNorm):
        super().__init__() #调用父类的构造方法进行初始化
        self.d_model = d_model
        self.win_size = win_size
        #定义线性变换层，将合并后的窗口重新投影到原始维度空间
        self.linear_trans = nn.Linear(win_size * d_model, d_model)
        #定义正则化层，标准化输入数据
        self.norm = norm_layer(win_size * d_model)
    #
    def forward(self, x):
        #从输入数据的形状中提取批次大小、时间步长、段数和模型维度
        batch_size, ts_d, seg_num, d_model = x.shape
        #计算需要填充的段数以保证窗口大小
        pad_num = seg_num % self.win_size
        #如果需要填充，则在时间序列的末尾复制最后几个段以进行填充
        if pad_num != 0:
            pad_num = self.win_size - pad_num
            x = torch.cat((x, x[:, :, -pad_num:, :]), dim=-2)
        #以窗口为步长，从每个窗口开始选取片段，并将其添加在seg_to_merge
        seg_to_merge = []
        for i in range(self.win_size):
            seg_to_merge.append(x[:, :, i::self.win_size, :])
        #将列表中的所有片段按最后一个维度（特征维度）拼接
        x = torch.cat(seg_to_merge, -1)
        #对合并后的数据进行正则化
        x = self.norm(x)
        #通过线性变换将数据投影回原来的特征维度
        x = self.linear_trans(x)

        return x

#处理不同尺度时间序列的模块
class scale_block(nn.Module):
    #初始化方法，参数包括模型的配置configs，窗口大小win_size，模型维度d_model，注意力头数n_heads，前馈网络的维度d_ff，模块的深度depth，dropout率，段数seg_num以及一个因子factor。
    def __init__(self, configs, win_size, d_model, n_heads, d_ff, depth, dropout, \
                 seg_num=10, factor=10):
        #调用父类的构造方法进行初始化
        super(scale_block, self).__init__()
        #如果窗口大小大于1，则创建一个合并层，否则设置为None。
        if win_size > 1:
            self.merge_layer = SegMerging(d_model, win_size, nn.LayerNorm)
        else:
            self.merge_layer = None
        #创建模块列表以存储编码层。
        self.encode_layers = nn.ModuleList()
        #向模块列表中添加depth数量的TwoStageAttentionLayer，这是一个包含两阶段注意力机制的层。
        for i in range(depth):
            self.encode_layers.append(TwoStageAttentionLayer(configs, seg_num, factor, d_model, n_heads, \
                                                             d_ff, dropout))
    #定义前向传播方法，接收输入x和额外的参数。
    def forward(self, x, attn_mask=None, tau=None, delta=None):
        #从输入数据的形状中提取时间步长维度。
        _, ts_dim, _, _ = x.shape
        #如果存在合并层，则将输入数据通过合并层。
        if self.merge_layer is not None:
            x = self.merge_layer(x)
        #将数据通过所有编码层。
        for layer in self.encode_layers:
            x = layer(x)

        return x, None

#编码器类
class Encoder(nn.Module):
    #初始化方法，参数为注意力层列表
    def __init__(self, attn_layers):
        #调用父类的构造方法进行初始化
        super(Encoder, self).__init__()
        #将传入的注意力层参数转换为模块列表
        self.encode_blocks = nn.ModuleList(attn_layers)
    #定义前向传播方法
    def forward(self, x):
        #创建列表并添加原始输入x
        encode_x = []
        encode_x.append(x)
        #通过所有编码块迭代，并将每个块的输出添加到列表中。
        for block in self.encode_blocks:
            x, attns = block(x)
            encode_x.append(x)

        return encode_x, None

#解码器层类
class DecoderLayer(nn.Module):
    #初始化方法接收自注意力和交叉注意力层，段长度seg_len，模型维度d_model，可选的前馈网络维度d_ff，以及dropout率。
    def __init__(self, self_attention, cross_attention, seg_len, d_model, d_ff=None, dropout=0.1):
        #调用父类的构造方法进行初始化。
        super(DecoderLayer, self).__init__()
        #存储自注意力和交叉注意力机制。
        self.self_attention = self_attention
        self.cross_attention = cross_attention
        #定义两个正则化层，用于标准化数据
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        #定义dropout层
        self.dropout = nn.Dropout(dropout)
        #创建一个简单的多层感知器（MLP）结构，其中包含GELU激活函数
        self.MLP1 = nn.Sequential(nn.Linear(d_model, d_model),
                                  nn.GELU(),
                                  nn.Linear(d_model, d_model))
        #线性预测层，用于将解码器输出映射到段长度。
        self.linear_pred = nn.Linear(d_model, seg_len)
    #定义前向传播方法
    def forward(self, x, cross):
        #获取批次大小
        batch = x.shape[0]
        #将数据通过自注意力层
        x = self.self_attention(x)
        #使用rearrange函数重组数据以匹配交叉注意力层的期望输入格式
        x = rearrange(x, 'b ts_d out_seg_num d_model -> (b ts_d) out_seg_num d_model')

        #将交叉编码数据也重组以匹配期望的格式
        cross = rearrange(cross, 'b ts_d in_seg_num d_model -> (b ts_d) in_seg_num d_model')
        #通过交叉注意力层，输入是重组后的数据和交叉编码数据
        tmp, attn = self.cross_attention(x, cross, cross, None, None, None,)
        #将注意力层的输出加到输入上，并应用dropout
        x = x + self.dropout(tmp)
        #应用第一个正则化层，并将结果保存在y中。
        y = x = self.norm1(x)
        #将y通过多层感知器。
        y = self.MLP1(y)
        #将MLP的输出与x相加，并应用第二个正则化层。
        dec_output = self.norm2(x + y)
        
        #重新解码输出以恢复原始批次维度。
        dec_output = rearrange(dec_output, '(b ts_d) seg_dec_num d_model -> b ts_d seg_dec_num d_model', b=batch)
        #通过线性预测层生成解码层的预测
        layer_predict = self.linear_pred(dec_output)
        #重新预测输出以匹配期望的输出格式
        layer_predict = rearrange(layer_predict, 'b out_d seg_num seg_len -> b (out_d seg_num) seg_len')

        return dec_output, layer_predict

#解码器类
class Decoder(nn.Module):
    #初始化方法，参数为解码器层列表
    def __init__(self, layers):
        #调用父类的构造方法进行初始化
        super(Decoder, self).__init__()
        #将解码器层转换为模块列表
        self.decode_layers = nn.ModuleList(layers)

    #定义前向传播方法，接收输入x和编码器的输出cross
    def forward(self, x, cross):
        #初始化最终预测为None
        final_predict = None
        i = 0
        #从输入数据的形状中提取时间步长维度
        ts_d = x.shape[1]
        #迭代所有解码器层
        for layer in self.decode_layers:
            #获取当前对应的编码器输出
            cross_enc = cross[i]
            #将数据通过解码器层
            x, layer_predict = layer(x, cross_enc)
            #如果是第一次迭代，将最终预测设置为该层预测，否则将该层预测累加到最终预测上
            if final_predict is None:
                final_predict = layer_predict
            else:
                final_predict = final_predict + layer_predict
            i += 1
        #重新安排最终预测的形状以匹配期望的输出格式
        final_predict = rearrange(final_predict, 'b (out_d seg_num) seg_len -> b (seg_num seg_len) out_d', out_d=ts_d)

        return final_predict