text: from re import X
  import torch
  import numpy as np
  import os.path
  import random
  from federated_dca.models import ZINBAutoEncoder, NBAutoEncoder
  from federated_dca.loss import ZINBLoss, NBLoss
  import torch
  from torch.utils.data import DataLoader
  from federated_dca.datasets import GeneCountData, write_text_matrix
  
  def save_and_load_init_model(model, mname, base='data/checkpoints/'):
      if os.path.exists(os.path.abspath('.') + base + '/init_' + mname + '.npy'):
          saved_params = np.load('init_' + mname + '.npy', allow_pickle=True)
          with torch.no_grad():
              for name, params in model.named_parameters():
                  sh = params.shape
                  if name == 'encoder.0.weight':
                      params.data = torch.from_numpy(np.transpose(saved_params[0]))
                  elif name == 'encoder.0.bias':
                      params.data = torch.from_numpy(saved_params[1])
                  elif name == 'bottleneck.0.weight':
                      params.data = torch.from_numpy(np.transpose(saved_params[2]))
                  elif name == 'bottleneck.0.bias':
                      params.data = torch.from_numpy(saved_params[3])
                  elif name == 'decoder.0.weight':
                      params.data = torch.from_numpy(np.transpose(saved_params[4]))
                  elif name == 'decoder.0.bias':
                      params.data = torch.from_numpy(saved_params[5])
                  elif name == 'mean.0.weight':
                      params.data = torch.from_numpy(np.transpose(saved_params[6]))
                  elif name == 'mean.0.bias':
                      params.data = torch.from_numpy(saved_params[7])
                  elif name == 'disp.0.weight':
                      params.data = torch.from_numpy(np.transpose(saved_params[8]))
                  elif name == 'disp.0.bias':
                      params.data = torch.from_numpy(saved_params[9])
                  elif name == 'drop.0.weight':
                      params.data = torch.from_numpy(np.transpose(saved_params[10]))
                  elif name == 'drop.0.bias':
                      params.data = torch.from_numpy(saved_params[11])
  
          return model
      else:
          print('No original dca params to load!!!')
          params_to_save = []
          for name, params in model.named_parameters():
              sh = params.shape
              if name == 'encoder.0.weight':
                  params_to_save.append(np.transpose(params.detach().numpy()))
              elif name == 'encoder.0.bias':
                  params_to_save.append(params.detach().numpy())
              elif name == 'bottleneck.0.weight':
                  params_to_save.append(np.transpose(params.detach().numpy()))
              elif name == 'bottleneck.0.bias':
                 params_to_save.append(params.detach().numpy())
              elif name == 'decoder.0.weight':
                  params_to_save.append(np.transpose(params.detach().numpy()))
              elif name == 'decoder.0.bias':
                  params_to_save.append(params.detach().numpy())
              elif name == 'mean.0.weight':
                  params_to_save.append(np.transpose(params.detach().numpy()))
              elif name == 'mean.0.bias':
                  params_to_save.append(params.detach().numpy())
              elif name == 'disp.0.weight':
                  params_to_save.append(np.transpose(params.detach().numpy()))
              elif name == 'disp.0.bias':
                  params_to_save.append(params.detach().numpy())
              elif name == 'drop.0.weight':
                  params_to_save.append(np.transpose(params.detach().numpy()))
              elif name == 'drop.0.bias':
                  params_to_save.append(params.detach().numpy())
          np.save(base + '1e_' + mname, params_to_save)
          return model
  
  def load_params(path):
      pass
  
  class trainInstince():
      def __init__(self, params) -> None:
          self.config = params
          self.train_loss = []
          self.val_loss = []
          self.denoise = params['result']['denoise']
          self.result = params['result']['data']
          self.epoch = params['model_parameters']['epoch']
          self.epoch_count = 0
          self.lr = params['model_parameters']['lr']
          self.batch = params['model_parameters']['batch']
          self.encoder_size = params['model_parameters']['encoder_size']
          self.bottleneck_size = params['model_parameters']['bottleneck_size']
          self.ridge = params['model_parameters']['ridge']
          self.reduce_lr = params['model_parameters']['reduce_lr']
          self.early_stopping = params['model_parameters']['early_stopping']
          self.es_count = 0
          self.name = params['model_parameters']['name']
          self.dataset_path = params['local_dataset']['data']
          self.loginput = params['local_dataset']['loginput']
          self.norminput = params['local_dataset']['norminput']
          self.transpose = params['local_dataset']['transpose']
          self.test_split = params['local_dataset']['test_split']
          self.seed = params['model_parameters']['seed']
          self.param_factor = params['model_parameters']['param_factor']
          self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
          self.dataset = GeneCountData('/mnt/input/'+self.dataset_path, self.device, transpose=self.transpose,
                          loginput=self.loginput, norminput=self.norminput, test_split=self.test_split)
          self.input_size = self.dataset.gene_num
          self.dataset.set_mode('train')
          self.trainDataLoader = DataLoader(self.dataset, batch_size=self.batch, shuffle=True)
          self.dataset.set_mode('val')
          self.valDataLoader = DataLoader(self.dataset, batch_size=self.batch)
          self.model_type = params['model_parameters']['model_type']
          if self.model_type == 'zinb':
              self.model = ZINBAutoEncoder(input_size=self.input_size, encoder_size=self.encoder_size,
                          bottleneck_size=self.bottleneck_size).to(self.device)
              self.loss = ZINBLoss(ridge_lambda=self.ridge)
          else:
              self.model = NBAutoEncoder(input_size=self.input_size, encoder_size=self.encoder_size,
                          bottleneck_size=self.bottleneck_size).to(self.device)
              self.loss = NBLoss()
          self.optimizer = torch.optim.RMSprop(self.model.parameters(), lr=self.lr)
          self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(self.optimizer, mode='min', patience=self.reduce_lr)
          self.finished_training = False
          self.best_val_loss = float('inf')
          seed = self.seed
          torch.manual_seed(seed)
          torch.cuda.manual_seed(seed)
          torch.cuda.manual_seed_all(seed)
          np.random.seed(seed)
          random.seed(seed)
          torch.backends.cudnn.benchmark = False
          torch.backends.cudnn.deterministic = True
          os.environ['PYTHONHASHSEED'] = f'{seed}'
  
      def train(self, update, log, id):
          epoch = self.epoch_count
          train_loss = 0
          self.model.train()
          self.dataset.set_mode(self.dataset.train)
          for data, target, size_factor in self.trainDataLoader:
  
              if self.model_type == 'zinb':
                  mean, disp, drop = self.model(data, size_factor)
                  loss = self.loss(target, mean, disp, drop)
              else:
                  mean, disp = self.model(data)
                  loss = self.loss(data, mean, disp)
  
              self.optimizer.zero_grad()
              loss.backward()
              self.optimizer.step()
  
              train_loss += loss.item()
  
          avg_loss = train_loss / len(self.trainDataLoader)
          log(f'Epoch: {epoch}, Avg train loss: '+"%.6f"%avg_loss)
  
          val_loss = 0
          with torch.no_grad():
              self.model.eval()
              self.dataset.set_mode(self.dataset.val)
              for data, target, size_factor in self.valDataLoader:
                  if self.model_type == 'zinb':
                      mean, disp, drop = self.model(data, size_factor)
                      loss = self.loss(target, mean, disp, drop)
                  else:
                      mean, disp = self.model(data)
                      loss = self.loss(data, mean, disp)
  
                  val_loss += loss.item()
              
              avg_loss = val_loss / len(self.valDataLoader)
              self.scheduler.step(avg_loss)
              log(f'Epoch: {epoch}, Avg val loss: '+"%.6f"%avg_loss)
              if avg_loss < self.best_val_loss:
                  self.best_val_loss = avg_loss
                  torch.save(self.model.state_dict(), '/mnt/output/'+self.name+'.pt')
                  self.es_count = 0
              else:
                  self.es_count += 1
  
          self.epoch_count += 1
          if self.epoch - self.epoch_count <= 0:
              self.finished_training = True
      
      def get_weights(self):
          model = self.model
          weights_list = []
          for name, params in model.named_parameters():
              weights_list.append(params.data)
          return weights_list
      
      def set_weights(self, weights):
          model = self.model
          index = 0
          with torch.no_grad():
              for name, params in model.named_parameters():
                  params.data = params.data + self.param_factor * (weights[index] - params.data)
                  index += 1
      
      def finish(self):
          np.save('/mnt/output/train_loss', self.train_loss)
          np.save('/mnt/output/val_loss', self.val_loss)
          if self.denoise:
              self.model.load_state_dict(torch.load('/mnt/output/'+self.name+'.pt'))
              self.model.eval()
              self.dataset.set_mode(self.dataset.test)
              eval_dataloader = DataLoader(self.dataset, batch_size=self.dataset.__len__())
              for data, target, size_factor in eval_dataloader:
                  if self.model_type == 'zinb':
                      mean, disp, drop = self.model(data, size_factor)
                      loss = self.loss(target, mean, disp, drop)
                  else:
                      mean, disp = self.model(data, size_factor)
                      loss = self.loss(target, mean, disp)
              adata = self.dataset.adata.copy()
              adata.X = mean.detach().numpy()
              colnames = adata.var_names.values
              rownames = adata.obs_names.values
              write_text_matrix(adata.X,
                                '/mnt/output/'+self.result,
                                rownames=rownames, colnames=colnames, transpose=True)
  
  
  
  def average_model_params(model_params):
      params = []
      for i in list(range(len(model_params[0]))):
          weight = 0
          for model in model_params:
              if type(weight) is int:
                  weight = model[i]
              else:
                  weight += model[i]
          
          weight = weight / len(model_params)
          params.append(weight)
          
      return params
  