text: # Copyright 2021 Huawei Technologies Co., Ltd
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  # http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ============================================================================
  """model"""
  
  import mindspore.nn as nn
  
  from encoder import Encoder, EncoderLayer, ConvLayer, EncoderStack
  from decoder import Decoder, DecoderLayer
  from attn import FullAttention, ProbAttention, AttentionLayer
  from embed import DataEmbedding
  
  
  class Informer(nn.Cell):
      """Informer"""
      def __init__(self, enc_in, dec_in, c_out, out_len,
                   factor=5, d_model=512, n_heads=8, e_layers=3, d_layers=2, d_ff=512,
                   dropout=0.0, attn='prob', embed='fixed', freq='h', activation='gelu',
                   output_attention=False, distil=True, mix=True):
          super(Informer, self).__init__()
          self.pred_len = out_len
          self.attn = attn
          self.output_attention = output_attention
  
          # Encoding
          self.enc_embedding = DataEmbedding(enc_in, d_model, embed, freq, dropout)
          self.dec_embedding = DataEmbedding(dec_in, d_model, embed, freq, dropout)
          # Attention
          attn_f = ProbAttention if attn == 'prob' else FullAttention
          # Encoder
          self.encoder = Encoder(
              [
                  EncoderLayer(
                      AttentionLayer(attn_f(False, factor, attention_dropout=dropout, output_attention=output_attention),
                                     d_model, n_heads, mix=False),
                      d_model,
                      d_ff,
                      dropout=dropout,
                      activation=activation
                  ) for l in range(e_layers)
              ],
              [
                  ConvLayer(
                      d_model
                  ) for l in range(e_layers - 1)
              ] if distil else None,
              norm_layer=nn.LayerNorm(normalized_shape=d_model, epsilon=1e-05)
          )
          # Decoder
          self.decoder = Decoder(
              [
                  DecoderLayer(
                      AttentionLayer(attn_f(True, factor, attention_dropout=dropout, output_attention=False),
                                     d_model, n_heads, mix=mix),
                      AttentionLayer(FullAttention(False, factor, attention_dropout=dropout, output_attention=False),
                                     d_model, n_heads, mix=False),
                      d_model,
                      d_ff,
                      dropout=dropout,
                      activation=activation,
                  )
                  for l in range(d_layers)
              ],
              norm_layer=nn.LayerNorm(normalized_shape=d_model, epsilon=1e-05)
          )
          # self.end_conv1 = nn.Conv1d(in_channels=label_len+out_len, out_channels=out_len, kernel_size=1, bias=True)
          # self.end_conv2 = nn.Conv1d(in_channels=d_model, out_channels=c_out, kernel_size=1, bias=True)
          self.projection = nn.Dense(in_channels=d_model, out_channels=c_out, has_bias=True)
  
      def construct(self, x_enc, x_mark_enc, x_dec, x_mark_dec,
                    enc_self_mask=None, dec_self_mask=None, dec_enc_mask=None):
          """build"""
          enc_out = self.enc_embedding(x_enc, x_mark_enc)
          enc_out, attns = self.encoder(enc_out, attn_mask=enc_self_mask)
  
          dec_out = self.dec_embedding(x_dec, x_mark_dec)
          dec_out = self.decoder(dec_out, enc_out, x_mask=dec_self_mask, cross_mask=dec_enc_mask)
          dec_out = self.projection(dec_out)
  
          # dec_out = self.end_conv1(dec_out)
          # dec_out = self.end_conv2(dec_out.transpose(2,1)).transpose(1,2)
          if self.output_attention:
              return dec_out[:, -self.pred_len:, :], attns
          return dec_out[:, -self.pred_len:, :]  # [B, L, D]
  
  
  class InformerStack(nn.Cell):
      """InformerStack"""
      def __init__(self, enc_in, dec_in, c_out, out_len,
                   factor=5, d_model=512, n_heads=8, e_layers=(3, 2, 1), d_layers=2, d_ff=512,
                   dropout=0.0, attn='prob', embed='fixed', freq='h', activation='gelu',
                   output_attention=False, distil=True, mix=True):
          super(InformerStack, self).__init__()
          self.pred_len = out_len
          self.attn = attn
          self.output_attention = output_attention
  
          # Encoding
          self.enc_embedding = DataEmbedding(enc_in, d_model, embed, freq, dropout)
          self.dec_embedding = DataEmbedding(dec_in, d_model, embed, freq, dropout)
          # Attention
          attn_f = ProbAttention if attn == 'prob' else FullAttention
          # Encoder
  
          inp_lens = list(range(len(e_layers)))  # [0,1,2,...] you can customize here
          encoders = [
              Encoder(
                  [
                      EncoderLayer(
                          AttentionLayer(
                              attn_f(False, factor, attention_dropout=dropout, output_attention=output_attention),
                              d_model, n_heads, mix=False),
                          d_model,
                          d_ff,
                          dropout=dropout,
                          activation=activation
                      ) for l in range(el)
                  ],
                  [
                      ConvLayer(
                          d_model
                      ) for l in range(el - 1)
                  ] if distil else None,
                  norm_layer=nn.LayerNorm(normalized_shape=d_model, epsilon=1e-05)
              ) for el in e_layers]
          self.encoder = EncoderStack(encoders, inp_lens)
          # Decoder
          self.decoder = Decoder(
              [
                  DecoderLayer(
                      AttentionLayer(attn_f(True, factor, attention_dropout=dropout, output_attention=False),
                                     d_model, n_heads, mix=mix),
                      AttentionLayer(FullAttention(False, factor, attention_dropout=dropout, output_attention=False),
                                     d_model, n_heads, mix=False),
                      d_model,
                      d_ff,
                      dropout=dropout,
                      activation=activation,
                  )
                  for l in range(d_layers)
              ],
              norm_layer=nn.LayerNorm(normalized_shape=d_model, epsilon=1e-05)
          )
          # self.end_conv1 = nn.Conv1d(in_channels=label_len+out_len, out_channels=out_len, kernel_size=1, bias=True)
          # self.end_conv2 = nn.Conv1d(in_channels=d_model, out_channels=c_out, kernel_size=1, bias=True)
          self.projection = nn.Dense(in_channels=d_model, out_channels=c_out, has_bias=True)
  
      def construct(self, x_enc, x_mark_enc, x_dec, x_mark_dec,
                    dec_self_mask=None, dec_enc_mask=None):
          """build"""
          enc_out = self.enc_embedding(x_enc, x_mark_enc)
          enc_out, attns = self.encoder(enc_out)
  
          dec_out = self.dec_embedding(x_dec, x_mark_dec)
          dec_out = self.decoder(dec_out, enc_out, x_mask=dec_self_mask, cross_mask=dec_enc_mask)
          dec_out = self.projection(dec_out)
  
          # dec_out = self.end_conv1(dec_out)
          # dec_out = self.end_conv2(dec_out.transpose(2,1)).transpose(1,2)
          if self.output_attention:
              return dec_out[:, -self.pred_len:, :], attns
          return dec_out[:, -self.pred_len:, :]  # [B, L, D]
  