text: import os
  import time
  import warnings
  
  import numpy as np
  import pandas as pd
  import torch
  import torch.nn as nn
  from torch import optim
  from torch.optim.lr_scheduler import CosineAnnealingLR
  
  from tsformer.datasets.data_factory import data_provider
  from tsformer.exp_basic import Exp_Basic
  from tsformer.models.autoformer import AutoFormer
  from tsformer.models.custom_informer import Informer
  from tsformer.models.transformer import Transformer
  from tsformer.utils.metrics import metric
  from tsformer.utils.tools import EarlyStopping, adjust_learning_rate, visual
  
  warnings.filterwarnings('ignore')
  
  
  class Exp_Main(Exp_Basic):
  
      def __init__(self, args):
          super(Exp_Main, self).__init__(args)
  
      def _build_model(self):
  
          enc_in = self.args.enc_in
          dec_in = self.args.dec_in
          c_out = self.args.c_out
          seq_len = self.args.seq_len
          label_len = self.args.label_len
          pred_len = self.args.pred_len
          factor = self.args.factor
          d_model = self.args.d_model
          n_heads = self.args.n_heads
          e_layers = self.args.e_layers
          d_layers = self.args.d_layers
          d_ffn = self.args.d_ffn
          dropout = self.args.dropout
          embed = self.args.embed
          moving_avg = self.args.moving_avg
          # freq = self.args.freq
          # activation = self.args.activation
          # distil = self.args.distil
  
          if self.args.model == 'transformer':
              model = Transformer(
                  enc_in=enc_in,
                  dec_in=dec_in,
                  c_out=c_out,
                  seq_len=seq_len,
                  label_len=label_len,
                  pred_len=pred_len,
                  factor=factor,
                  d_model=d_model,
                  n_heads=n_heads,
                  e_layers=e_layers,
                  d_layers=d_layers,
                  d_ffn=d_ffn,
                  dropout=dropout,
                  embed=embed)
  
          elif self.args.model == 'informer':
              model = Informer(
                  enc_in=enc_in,
                  dec_in=dec_in,
                  c_out=c_out,
                  seq_len=seq_len,
                  label_len=label_len,
                  pred_len=pred_len,
                  factor=factor,
                  d_model=d_model,
                  n_heads=n_heads,
                  e_layers=e_layers,
                  d_layers=d_layers,
                  d_ffn=d_ffn,
                  dropout=dropout,
                  embed=embed)
  
          elif self.args.model == 'autoformer':
              model = AutoFormer(
                  enc_in=enc_in,
                  dec_in=dec_in,
                  c_out=c_out,
                  seq_len=seq_len,
                  label_len=label_len,
                  pred_len=pred_len,
                  factor=factor,
                  d_model=d_model,
                  n_heads=n_heads,
                  e_layers=e_layers,
                  d_layers=d_layers,
                  d_ffn=d_ffn,
                  dropout=dropout,
                  embed=embed,
                  moving_avg=moving_avg)
  
          if self.args.use_multi_gpu and self.args.use_gpu:
              model = nn.DataParallel(model, device_ids=self.args.device_ids)
          return model
  
      def _get_lr_scheduler(self, epochs):
          optimizer = self._select_optimizer()
          scheduler = CosineAnnealingLR(
              optimizer, T_max=epochs, eta_min=0, last_epoch=-1)
          return scheduler
  
      def _get_data(self, flag):
          data_set, data_loader = data_provider(self.args, flag)
          return data_set, data_loader
  
      def _select_optimizer(self):
          model_optim = optim.Adam(
              self.model.parameters(), lr=self.args.learning_rate)
          return model_optim
  
      def _select_criterion(self):
          criterion = nn.MSELoss()
          #  criterion = nn.SmoothL1Loss()
          return criterion
  
      def vali(self, vali_data, vali_loader, criterion):
          total_loss = []
          self.model.eval()
          with torch.no_grad():
              for i, (batch_x, batch_y, batch_x_mark,
                      batch_y_mark) in enumerate(vali_loader):
                  batch_x = batch_x.float().to(self.device)
                  batch_y = batch_y.float()
  
                  batch_x_mark = batch_x_mark.float().to(self.device)
                  batch_y_mark = batch_y_mark.float().to(self.device)
  
                  # decoder input
                  dec_inp = torch.zeros_like(
                      batch_y[:, -self.args.pred_len:, :]).float()
                  dec_inp = torch.cat(
                      [batch_y[:, :self.args.label_len, :], dec_inp],
                      dim=1).float().to(self.device)
                  # encoder - decoder
                  if self.args.use_amp:
                      with torch.cuda.amp.autocast():
                          if self.args.output_attention:
                              outputs = self.model(batch_x, batch_x_mark,
                                                   dec_inp, batch_y_mark)[0]
                          else:
                              outputs = self.model(batch_x, batch_x_mark,
                                                   dec_inp, batch_y_mark)
                  else:
                      if self.args.output_attention:
                          outputs = self.model(batch_x, batch_x_mark, dec_inp,
                                               batch_y_mark)[0]
                      else:
                          outputs = self.model(batch_x, batch_x_mark, dec_inp,
                                               batch_y_mark)
                  f_dim = -1 if self.args.features == 'MS' else 0
                  batch_y = batch_y[:, -self.args.pred_len:,
                                    f_dim:].to(self.device)
  
                  pred = outputs.detach().cpu()
                  true = batch_y.detach().cpu()
  
                  loss = criterion(pred, true)
  
                  total_loss.append(loss)
          total_loss = np.average(total_loss)
          self.model.train()
          return total_loss
  
      def train(self, setting):
          train_data, train_loader = self._get_data(flag='train')
          vali_data, vali_loader = self._get_data(flag='val')
          test_data, test_loader = self._get_data(flag='test')
  
          path = os.path.join(self.args.results_dir, 'checkpoints', setting)
          if not os.path.exists(path):
              os.makedirs(path)
  
          time_now = time.time()
  
          train_steps = len(train_loader)
          early_stopping = EarlyStopping(
              patience=self.args.patience, verbose=True)
  
          model_optim = self._select_optimizer()
          criterion = self._select_criterion()
  
          if self.args.use_amp:
              scaler = torch.cuda.amp.GradScaler()
  
          for epoch in range(self.args.train_epochs):
              iter_count = 0
              train_loss = []
  
              self.model.train()
              epoch_time = time.time()
              for i, (batch_x, batch_y, batch_x_mark,
                      batch_y_mark) in enumerate(train_loader):
                  iter_count += 1
                  model_optim.zero_grad()
                  batch_x = batch_x.float().to(self.device)
  
                  batch_y = batch_y.float().to(self.device)
                  batch_x_mark = batch_x_mark.float().to(self.device)
                  batch_y_mark = batch_y_mark.float().to(self.device)
  
                  # decoder input
                  dec_inp = torch.zeros_like(
                      batch_y[:, -self.args.pred_len:, :]).float()
                  dec_inp = torch.cat(
                      [batch_y[:, :self.args.label_len, :], dec_inp],
                      dim=1).float().to(self.device)
  
                  # encoder - decoder
                  if self.args.use_amp:
                      with torch.cuda.amp.autocast():
                          if self.args.output_attention:
                              outputs = self.model(batch_x, batch_x_mark,
                                                   dec_inp, batch_y_mark)[0]
                          else:
                              outputs = self.model(batch_x, batch_x_mark,
                                                   dec_inp, batch_y_mark)
  
                          f_dim = -1 if self.args.features == 'MS' else 0
                          batch_y = batch_y[:, -self.args.pred_len:,
                                            f_dim:].to(self.device)
                          loss = criterion(outputs, batch_y)
                          train_loss.append(loss.item())
                  else:
                      if self.args.output_attention:
                          outputs = self.model(batch_x, batch_x_mark, dec_inp,
                                               batch_y_mark)[0]
                      else:
                          outputs = self.model(batch_x, batch_x_mark, dec_inp,
                                               batch_y_mark, batch_y)
  
                      f_dim = -1 if self.args.features == 'MS' else 0
                      batch_y = batch_y[:, -self.args.pred_len:,
                                        f_dim:].to(self.device)
                      loss = criterion(outputs, batch_y)
                      train_loss.append(loss.item())
  
                  if (i + 1) % 100 == 0:
                      print('\titers: {0}, epoch: {1} | loss: {2:.7f}'.format(
                          i + 1, epoch + 1, loss.item()))
                      speed = (time.time() - time_now) / iter_count
                      left_time = speed * (
                          (self.args.train_epochs - epoch) * train_steps - i)
                      print('\tspeed: {:.4f}s/iter; left time: {:.4f}s'.format(
                          speed, left_time))
                      iter_count = 0
                      time_now = time.time()
  
                  if self.args.use_amp:
                      scaler.scale(loss).backward()
                      scaler.step(model_optim)
                      scaler.update()
                  else:
                      loss.backward()
                      model_optim.step()
  
              print('Epoch: {} cost time: {}'.format(epoch + 1,
                                                     time.time() - epoch_time))
              train_loss = np.average(train_loss)
              vali_loss = self.vali(vali_data, vali_loader, criterion)
              test_loss = self.vali(test_data, test_loader, criterion)
  
              print(
                  'Epoch: {0}, Steps: {1} | Train Loss: {2:.7f} Vali Loss: {3:.7f} Test Loss: {4:.7f}'
                  .format(epoch + 1, train_steps, train_loss, vali_loss,
                          test_loss))
              early_stopping(vali_loss, self.model, path)
              if early_stopping.early_stop:
                  print('Early stopping')
                  break
  
              adjust_learning_rate(model_optim, epoch + 1, self.args)
  
          best_model_path = os.path.join(path, 'checkpoint.pth')
          self.model.load_state_dict(torch.load(best_model_path))
  
          return self.model
  
      def test(self, setting, test=0):
          test_data, test_loader = self._get_data(flag='test')
          if test:
              print('loading model')
              self.model.load_state_dict(
                  torch.load(
                      os.path.join(self.args.results_dir, 'checkpoints', setting,
                                   'checkpoint.pth')))
  
          preds = []
          trues = []
          folder_path = os.path.join(self.args.results_dir, 'test_results',
                                     setting)
          if not os.path.exists(folder_path):
              os.makedirs(folder_path)
  
          self.model.eval()
          with torch.no_grad():
              for i, (batch_x, batch_y, batch_x_mark,
                      batch_y_mark) in enumerate(test_loader):
                  batch_x = batch_x.float().to(self.device)
                  batch_y = batch_y.float().to(self.device)
  
                  batch_x_mark = batch_x_mark.float().to(self.device)
                  batch_y_mark = batch_y_mark.float().to(self.device)
  
                  # decoder input
                  dec_inp = torch.zeros_like(
                      batch_y[:, -self.args.pred_len:, :]).float()
                  dec_inp = torch.cat(
                      [batch_y[:, :self.args.label_len, :], dec_inp],
                      dim=1).float().to(self.device)
                  # encoder - decoder
                  if self.args.use_amp:
                      with torch.cuda.amp.autocast():
                          if self.args.output_attention:
                              outputs = self.model(batch_x, batch_x_mark,
                                                   dec_inp, batch_y_mark)[0]
                          else:
                              outputs = self.model(batch_x, batch_x_mark,
                                                   dec_inp, batch_y_mark)
                  else:
                      if self.args.output_attention:
                          outputs = self.model(batch_x, batch_x_mark, dec_inp,
                                               batch_y_mark)[0]
  
                      else:
                          outputs = self.model(batch_x, batch_x_mark, dec_inp,
                                               batch_y_mark)
  
                  f_dim = -1 if self.args.features == 'MS' else 0
  
                  batch_y = batch_y[:, -self.args.pred_len:,
                                    f_dim:].to(self.device)
                  outputs = outputs.detach().cpu().numpy()
                  batch_y = batch_y.detach().cpu().numpy()
  
                  pred = outputs  # outputs.detach().cpu().numpy()  # .squeeze()
                  true = batch_y  # batch_y.detach().cpu().numpy()  # .squeeze()
  
                  preds.append(pred)
                  trues.append(true)
                  if i % 20 == 0:
                      input = batch_x.detach().cpu().numpy()
                      gt = np.concatenate((input[0, :, -1], true[0, :, -1]),
                                          axis=0)
                      pd = np.concatenate((input[0, :, -1], pred[0, :, -1]),
                                          axis=0)
                      visual(gt, pd, os.path.join(folder_path, str(i) + '.png'))
  
          preds = np.array(preds)
          trues = np.array(trues)
          print('test shape:', preds.shape, trues.shape)
          preds = preds.reshape(-1, preds.shape[-2], preds.shape[-1])
          trues = trues.reshape(-1, trues.shape[-2], trues.shape[-1])
          print('test shape:', preds.shape, trues.shape)
  
          # result save
          folder_path = os.path.join(self.args.results_dir, 'results', setting)
          if not os.path.exists(folder_path):
              os.makedirs(folder_path)
  
          mae, mse, rmse, mape, mspe = metric(preds, trues)
          print('mse:{}, mae:{},rmse:{}, mape:{}, mspe:{}'.format(
              mse, mae, rmse, mape, mspe))
          test_result_file = os.path.join(folder_path, 'result.txt')
          with open(test_result_file, 'w+') as f:
              f.write(setting + '  \n')
              f.write('mse:{}, mae:{}, rmse:{}, mape:{}, mspe{}'.format(
                  mse, mae, rmse, mape, mspe))
              f.write('\n')
              f.write('\n')
              f.close()
  
      def predict(self, setting, load=False):
          pred_data, pred_loader = self._get_data(flag='pred')
  
          if load:
              path = os.path.join(self.args.results_dir, 'checkpoints', setting)
              best_model_path = os.path.join(path, 'checkpoint.pth')
              self.model.load_state_dict(torch.load(best_model_path))
  
          preds = []
  
          self.model.eval()
          with torch.no_grad():
              for i, (batch_x, batch_y, batch_x_mark,
                      batch_y_mark) in enumerate(pred_loader):
                  batch_x = batch_x.float().to(self.device)
                  batch_y = batch_y.float()
                  batch_x_mark = batch_x_mark.float().to(self.device)
                  batch_y_mark = batch_y_mark.float().to(self.device)
  
                  # decoder input
                  dec_inp = torch.zeros_like(
                      batch_y[:, -self.args.pred_len:, :]).float()
                  dec_inp = torch.cat(
                      [batch_y[:, :self.args.label_len, :], dec_inp],
                      dim=1).float().to(self.device)
                  # encoder - decoder
                  if self.args.use_amp:
                      with torch.cuda.amp.autocast():
                          if self.args.output_attention:
                              outputs = self.model(batch_x, batch_x_mark,
                                                   dec_inp, batch_y_mark)[0]
                          else:
                              outputs = self.model(batch_x, batch_x_mark,
                                                   dec_inp, batch_y_mark)
                  else:
                      if self.args.output_attention:
                          outputs = self.model(batch_x, batch_x_mark, dec_inp,
                                               batch_y_mark)[0]
                      else:
                          outputs = self.model(batch_x, batch_x_mark, dec_inp,
                                               batch_y_mark)
                  pred = outputs.detach().cpu().numpy()  # .squeeze()
                  preds.append(pred)
  
          preds = np.array(preds)
          preds = preds.reshape(-1, preds.shape[-2], preds.shape[-1])
          df = pd.DataFrame(preds)
          # result save
          folder_path = os.path.join(self.args.results_dir + setting)
          if not os.path.exists(folder_path):
              os.makedirs(folder_path)
  
          df.to_csv(folder_path + 'real_prediction.csv', index=False)
  
          return
  