text: from tools.loading import load_models,load_datasets,load_optimizer, ANNEAL_TYPE
  from tools.others import sample_ZCs,make_bin_label,x_R_proccess_free, init_list, PGD_FREE_TYPE
  from tools.metric import get_metrics, set_hinge_bd
  from tools.adv_tools.pgd_free import l2_cliper
  import torch.optim as optim
  from torchvision.utils import save_image
  import argparse
  import torch
  import time
  from tools.metric import Metric_Types
  import math
  import numpy as np
  import os
  import random
  
  parser = argparse.ArgumentParser()
  parser.add_argument('--random_seed', type=int, default=10, help = 'random seed')
  parser.add_argument('--net', type=str, required=True, help='GAN structure: sngan_cifar, sngan_catdog64, sngan_catdog128, sngan_imgnet128, biggan')
  parser.add_argument('--dim_z', type=int, required=True, help = 'dimension of nosie Z')
  parser.add_argument('--img_c', type=int, required=True, help = 'image color channels')
  parser.add_argument('--fm_base_g', type=int, required=True, help = 'conv layer channels multiplier for g')
  parser.add_argument('--fm_base_d', type=int, required=True, help = 'conv layer channels multiplier for d')
  parser.add_argument('--n_classes', type=int, required=True, help = 'number of classes')
  parser.add_argument('--r_batch_size', type=int, required=True, help = 'real batch size')
  parser.add_argument('--g_batch_size', type=int, required=True, help = 'fake batch size')
  parser.add_argument('--num_bp', type=int, required=True, help = 'num of backprop to accumulate')
  parser.add_argument('--data_path', type=str, required=True, help = 'dataset path')
  parser.add_argument('--out_path', type=str, required=True, help = 'path of result to save, gen_*.pth and dis_*.pth')
  parser.add_argument('--bottom_width', type=int, required=True, help = 'the width of feature maps at the first layer of G')
  parser.add_argument('--img_width', type=int, required=True, help = 'the width of imgs')
  parser.add_argument('--itr_d', type=int, required=True, help = 'MAX_D_STEP')
  parser.add_argument('--lr_d', type=float, required=True, help = 'learning rate for D')
  parser.add_argument('--beta1_d', type=float, required=True, help = 'adam b1 for D')
  parser.add_argument('--beta2_d', type=float, required=True, help = 'adam b2 for D') 
  parser.add_argument('--itr_g', type=int, required=True, help = 'efault to 1')
  parser.add_argument('--lr_g', type=float, required=True, help = 'learning rate for G')
  parser.add_argument('--beta1_g', type=float, required=True, help = 'adam b1 for G')
  parser.add_argument('--beta2_g', type=float, required=True, help = 'adam b2 for G')
  parser.add_argument('--anneal_lr', type=str, default='nll', help = 'lr anneall type: exp, linear')
  parser.add_argument('--anneal_lr_p1', type=float, default=1, help = 'lr anneal p1 respect to total_itrs')
  parser.add_argument('--anneal_lr_p2', type=int, default=1, help = 'lr anneal p2 iters respect to total_itrs')
  parser.add_argument('--n_gpu', type=int, required=True, help = 'num of gpus')
  parser.add_argument('--dali', action='store_true', default=False, help = 'if use nvidia dali to load image')
  parser.add_argument('--da', action='store_true', default=False, help = 'data augmentation?')
  parser.add_argument('--n_workers', type=int, required=True, help = 'num of dataloader works')
  parser.add_argument('--G_total_itrs', type=int, required=True, help = 'total updates of gen(total itrs)')
  parser.add_argument('--wide', action='store_true', default=False, help = 'if use wide D')
  parser.add_argument('--bn_d', action='store_true', default=False, help = 'if use batch_norm in D')
  parser.add_argument('--bn_g', action='store_true', default=False, help = 'if use batch_norm in G')
  parser.add_argument('--sn_d', action='store_true', default=False, help = 'if use sepcture_norm in D')
  parser.add_argument('--sn_g', action='store_true', default=False, help = 'if use sepcture_norm in G')
  parser.add_argument('--sa_d', action='store_true', default=False, help = 'if use self-attetion in D')
  parser.add_argument('--sa_g', action='store_true', default=False, help = 'if use self-attetion in G')
  parser.add_argument('--Z_dist', type=str, default='normal', help = 'Z noise distribution type: normal')
  parser.add_argument('--Z_p1', type=float, default=0., help = 'Z noise distribution parameter 1')
  parser.add_argument('--Z_p2', type=float, default=1., help = 'Z noise distribution parameter 2')
  parser.add_argument('--save_bias', type=int, required=True, help = 'save bias respect to iterations')
  parser.add_argument('--start_itr', type=int, default=0, help = 'start itrs, 0 means new training')
  parser.add_argument('--shuffle', action='store_true', default=False, help = 'if shuffle data')
  parser.add_argument('--b_metric', type=str, default='ce',help='bin loss for training, ce for cross_entropy, or hinge')
  parser.add_argument('--c_metric', type=str, default='ce',help='class loss for training, ce for cross_entropy, or hinge projection, ce_kl for ce with KL in L^f_D')
  parser.add_argument('--c_alph_f', type=float, default=1.0,help = 'class loss coefficient of fake data for D')
  parser.add_argument('--c_alph_r', type=float, default=1.0,help = 'class loss coefficient of real data for D')
  parser.add_argument('--c_alph_g', type=float, default=1.0,help = 'class loss coefficient of generated data for G')
  parser.add_argument('--hinge_bd', type=float, default=1.0,help = 'margin boudary of hinge loss')
  parser.add_argument('--pgd_free_steps', type=int, required=True, help = 'number of steps for pgd free')
  parser.add_argument('--pgd_type', type=str, required=True, help = 'Linf or L2')
  parser.add_argument('--pgd_eps', type=float, required=True, help = 'radius of Linf-norm ball for pgd attack')
  parser.add_argument('--pgd_tau', type=float, required=True, help = 'lr of sgd for pgd attack')
  parser.add_argument('--pgd_rs', action='store_true', default=True, help = 'flag to random start pgd attack')
  opt = parser.parse_args()
  
  print(opt)
  torch.manual_seed(opt.random_seed)
  np.random.seed(opt.random_seed)
  set_hinge_bd(opt.hinge_bd)
  
  def train():
      lr_str = 'lr'
  
      Z_dist = opt.Z_dist
      Z_params = [opt.Z_p1,opt.Z_p2]
      sigma = 1e-8
      # models 
      netG,netD = load_models(net=opt.net,n_gpu=opt.n_gpu,dim_z=opt.dim_z,img_c=opt.img_c,img_width=opt.img_width,\
                  fm_base_d=opt.fm_base_d,fm_base_g=opt.fm_base_g,bn_d=opt.bn_d,bn_g=opt.bn_g,sn_d=opt.sn_d,sn_g=opt.sn_g,\
                  sa_d=opt.sa_d,sa_g=opt.sa_g,bottom_width=opt.bottom_width,n_classes=opt.n_classes,c_metric=opt.c_metric,wide=opt.wide)
  
      if opt.start_itr>0:
          netG.load_state_dict(torch.load(f'{opt.out_path}check_pts/gen_{opt.start_itr}.pth'))
          netD.load_state_dict(torch.load(f'{opt.out_path}check_pts/dis_{opt.start_itr}.pth'))
      G_itrs = int(opt.start_itr*opt.save_bias-1)
      D_itrs = int(G_itrs*(opt.itr_d/opt.itr_g))
      
      # optimizers
      optimizerG = load_optimizer(model=netG,lr=opt.lr_g,betas=(opt.beta1_g,opt.beta2_g))
      optimizerD = load_optimizer(model=netD,lr=opt.lr_d,betas=(opt.beta1_d,opt.beta2_d))
  
      # lr anneal
      if opt.anneal_lr in ANNEAL_TYPE:
          if opt.anneal_lr==ANNEAL_TYPE[0]:
              lr_lambda_D = lambda epoch: max(int(G_itrs<=opt.anneal_lr_p1),int(G_itrs>opt.anneal_lr_p1)*(opt.G_total_itrs-G_itrs)/(opt.G_total_itrs-opt.anneal_lr_p1))
              lr_lambda_G = lambda epoch: max(int(G_itrs<=opt.anneal_lr_p1),int(G_itrs>opt.anneal_lr_p1)*(opt.G_total_itrs-G_itrs)/(opt.G_total_itrs-opt.anneal_lr_p1))
          elif opt.anneal_lr==ANNEAL_TYPE[1]:
              lr_lambda_D = lambda epoch: opt.anneal_lr_p1 ** (G_itrs/opt.anneal_lr_p2)
              lr_lambda_G = lambda epoch: opt.anneal_lr_p1 ** (G_itrs/opt.anneal_lr_p2)
          lr_adjuster_D = torch.optim.lr_scheduler.LambdaLR(optimizerD, lr_lambda_D, last_epoch=-1)
          lr_adjuster_G = torch.optim.lr_scheduler.LambdaLR(optimizerG, lr_lambda_G, last_epoch=-1)
  
      if opt.r_batch_size<=0:
          raise ValueError(f"batch size must >0")
  
      if opt.r_batch_size>0:
          loader,class2idx = load_datasets(data_path=opt.data_path,img_width=opt.img_width, \
                          batch_size=opt.r_batch_size,n_workers=opt.n_workers,n_gpu=opt.n_gpu,shuffle=opt.shuffle,dali=opt.dali,seed=opt.random_seed,da=opt.da)
          loader_iter = iter(loader)
      
      fixed_Z,fixed_C_int,fixed_C_vec = sample_ZCs(64,opt.dim_z,opt.n_classes,Z_dist,Z_params,opt.n_gpu)
      
      print("Starting GAN Training Loop...")
      Metric_G, Metric_R, Metric_F = get_metrics()
      D_err_R,D_err_F,G_err,D_cnt,G_cnt=0,0,0,0,0
      tpts = time.time()
      itrs = 0
      delta = torch.zeros([opt.r_batch_size, opt.img_c, opt.img_width, opt.img_width]) # free accumulated perturbation
      if opt.n_gpu>0:
          delta = delta.cuda()
      if opt.pgd_rs: # perturbation init
          delta.normal_()
          if opt.pgd_type==PGD_FREE_TYPE[0]:
              delta.clamp_(-opt.pgd_eps,opt.pgd_eps)
          elif opt.pgd_type==PGD_FREE_TYPE[1]:
              delta = l2_cliper(delta,opt.pgd_eps,opt.n_gpu)
      delta = list(torch.chunk(delta,opt.num_bp))
  
      while G_itrs<opt.G_total_itrs:
          # load data
          if opt.r_batch_size>0:
              try:
                  data = next(loader_iter)
              except StopIteration:
                  loader_iter = iter(loader)
                  data = next(loader_iter)
              x_R = data[0]["data"] if opt.dali else data[0]
              y_R_c = data[0]["label"].view(-1).long() if opt.dali else data[1]
          else:
              x_R = torch.zeros(0, opt.img_c, opt.img_width,opt.img_width)
              y_R_c = torch.full((0,), 0).long()
          # to cuda ?
          if opt.n_gpu>0:
              x_R, y_R_c = x_R.cuda(), y_R_c.cuda()       
          
          x_R, y_R_c = list(torch.chunk(x_R,opt.num_bp)),list(torch.chunk(y_R_c,opt.num_bp))
  
          r_BS, g_BS = init_list(opt.num_bp),init_list(opt.num_bp)
  
          for bp_i in range(opt.num_bp):
              r_BS[bp_i] = x_R[bp_i].size(0)
              g_BS[bp_i] = int(opt.g_batch_size/opt.num_bp)
          if opt.g_batch_size%opt.num_bp !=0:
              g_BS[-1] = opt.g_batch_size%opt.num_bp
  
          y_R_b = init_list(opt.num_bp)
          y_G_b,y_F_b = init_list(opt.num_bp),init_list(opt.num_bp)
          for bp_i in range(opt.num_bp):
              y_R_b[bp_i],y_G_b[bp_i],y_F_b[bp_i] = make_bin_label(r_BS[bp_i],g_BS[bp_i],opt.n_gpu) # prepare inputs and label
  
          ############################
          # (1) Update D network:
          ###########################
          if itrs%opt.itr_g==0:
              D_itrs += 1
              # Train with all-real batch 
              x_F, y_F_c = init_list(opt.num_bp),init_list(opt.num_bp)
              for bp_i in range(opt.num_bp):
                  # produce fake imgs from G
                  Z, C_int, C_vec = sample_ZCs(g_BS[bp_i],opt.dim_z,opt.n_classes,Z_dist,Z_params,opt.n_gpu) # z <- Normal(0,1)
                  x_F[bp_i] = netG(Z,C_vec).detach() # x_F <- G(z)
                  y_F_c[bp_i] = C_int.view(-1)
                  
              for free_step in range(opt.pgd_free_steps):
                  D_cnt += 1
                  netD.zero_grad()
                  for bp_i in range(opt.num_bp):
                      # g_xR, g_wDr <- D(x_R)
                      delta[bp_i], loss_dis_R = x_R_proccess_free(netD=netD,x_R=x_R[bp_i],y_R_b=y_R_b[bp_i],y_R_c=y_R_c[bp_i],metric = Metric_R,\
                                              delta=delta[bp_i], pgd_type=opt.pgd_type, pgd_eps=opt.pgd_eps, pgd_tau=opt.pgd_tau,\
                                              b_metric=opt.b_metric,c_metric=opt.c_metric, c_alph=opt.c_alph_r,n_gpu=opt.n_gpu,num_bp=opt.num_bp)
                      D_err_R += (loss_dis_R.item())
                      o_F_b,o_F_c = netD(x_F[bp_i],y_F_c[bp_i]) 
                      loss_dis_F = Metric_F(o_F_b,o_F_c,y_F_b[bp_i],y_F_c[bp_i],opt.b_metric,opt.c_metric,1.0,opt.c_alph_f) 
                      loss_dis_F = (loss_dis_F/opt.num_bp) 
                      loss_dis_F.backward() # g_wDf <- D(x_F)
                      D_err_F += (loss_dis_F.item())
                  optimizerD.step()  # wD <- wD - (g_wDf + g_wDr) * lr
  
          if itrs % opt.itr_d == 0:
              G_cnt += 1
              G_itrs += 1
              ############################
              # (2) Update G network:
              ###########################
              netD.zero_grad()
              netG.zero_grad()
              for bp_i in range(opt.num_bp):
                  Z, C_int, C_vec = sample_ZCs(g_BS[bp_i],opt.dim_z,opt.n_classes,Z_dist,Z_params,opt.n_gpu)
                  y_G_c = C_int.view(-1)
                  x_G = netG(Z,C_vec)
                  o_G_b,o_G_c = netD(x_G,y_G_c)
                  loss_gen = Metric_G(o_G_b,o_G_c,y_G_b[bp_i],y_G_c,opt.b_metric,opt.c_metric,opt.c_alph_g)
                  loss_gen = (loss_gen/opt.num_bp)
                  G_err += loss_gen.item()
                  loss_gen.backward() # g_wG <- D(G(z))
              optimizerG.step() # wG <- wG - g_wG  * lr
          if G_itrs%opt.save_bias == 0 and (itrs % opt.itr_d==0):
              print(f'Training[{G_itrs}/{opt.G_total_itrs}] G_loss:{round(G_err/(G_cnt+sigma),3)},D_real:{round(D_err_R/(D_cnt+sigma),3)},D_fake:{round(D_err_F/(D_cnt+sigma),3)},secs:{round(time.time()-tpts,1)},lrD:{round(optimizerD.param_groups[0][lr_str],8)},lrG:{round(optimizerG.param_groups[0][lr_str],8)},r:{opt.r_batch_size},g:{opt.g_batch_size}')
              D_err_R,D_err_F,G_err,D_cnt,G_cnt=0,0,0,0,0
              tpts = time.time()
              curr_pts = int(G_itrs/opt.save_bias)
              with torch.no_grad():
                  fixed_x_fake = netG(fixed_Z,fixed_C_vec)
                  fixed_x_fake.data.mul_(0.5).add_(0.5)
                  fixed_x_fake = fixed_x_fake.view(64,opt.img_c,opt.img_width,opt.img_width)
                  save_image(fixed_x_fake.data, f'{opt.out_path}imgs/{curr_pts}.png', nrow=8)    
              torch.save(netG.state_dict(), f'{opt.out_path}check_pts/gen_{curr_pts}.pth')
              torch.save(netD.state_dict(), f'{opt.out_path}check_pts/dis_{curr_pts}.pth')
          
          # anneal part
          if opt.anneal_lr in ANNEAL_TYPE:
              lr_adjuster_G.step()
              lr_adjuster_D.step()
  
          # finally add itrs for keeping count continuing
          itrs += 1 
  
  
  
  if __name__ == "__main__":
      if not os.path.exists(opt.out_path):
          raise ValueError(f"No output path {opt.out_path}")
      if not os.path.exists(opt.out_path+'/check_pts'):
          os.makedirs(opt.out_path+'/check_pts')
      if not os.path.exists(opt.out_path+'/imgs'):
          os.makedirs(opt.out_path+'/imgs')
      if not os.path.exists(opt.out_path+'/log'):
          os.makedirs(opt.out_path+'/log')
      train()
  