text: """
  extra the last layer embedding of inception3
  
  refer: https://github.com/pytorch/examples/blob/master/imagenet/main.py
  """
  
  import argparse
  from tqdm import tqdm
  import numpy as np
  import json 
  import datetime
  import sys
  import os
  from scipy.io import loadmat
  import pickle
  
  import torch.nn.functional as F
  import torch.nn as nn
  import torch
  import yaml
  import torchvision.transforms as transforms
  import networkx as nx
  #import torchvision.datasets as datasets
  from torch.autograd import Variable
  
  from tensorboardX import SummaryWriter
  
  import utils
  from models.S2ILinearModel import S2ILayer2Model_IAS
  TMP=10
  
  class train(object):
      def __init__(self,device,cfg_file,IAS=False):
          self.IAS = IAS
          self.device = device
          with open(cfg_file,'r') as f:
             self.cfg = yaml.safe_load(f) 
          utils.init_seeds(self.cfg['model_hyp']['random_seed'])
  
          attSplit = loadmat('./dataset/xlsa/'+self.cfg['dataset']['name']+'/att_splits.mat') 
          res101 = loadmat('./dataset/xlsa/'+self.cfg['dataset']['name']+'/res101.mat')
          labels = res101['labels'].astype(int).squeeze() - 1
          seen_dataLoc = attSplit['test_seen_loc'].squeeze() - 1
          unseen_dataLoc = attSplit['test_unseen_loc'].squeeze() -1
          
          seen_labels = labels[seen_dataLoc]
          unseen_labels = labels[unseen_dataLoc]
          self.seen_labels = np.unique(seen_labels)
          self.unseen_labels = np.unique(unseen_labels)
  
          self.clsname = [ attSplit['allclasses_names'][i][0][0] for i in range(len(attSplit['allclasses_names']))]
          att_matrix = np.transpose(attSplit['att'])
          self.cfg['model_hyp']['att_feats'] = att_matrix.shape[1]
  
          self.attMatrix = att_matrix.copy()
          self.attMatrix[:len(self.seen_labels)] = att_matrix[self.seen_labels]
          self.attMatrix[len(self.seen_labels):] = att_matrix[self.unseen_labels]
  
          self.attMatrix = torch.FloatTensor(self.attMatrix).to(self.device)
  
      def createLinearModel(self,):
  
          self.model1 = S2ILayer2Model_IAS(att_dims=self.cfg['model_hyp']['att_feats'],img_dims=self.cfg['model_hyp']['img_feats'],tmp=5).to(self.device) 
  
      def apply_classification_weights(self,features, cls_weights,norm=False):
      
          features = F.normalize(features,dim=-1)
          cls_weights = F.normalize(cls_weights, p=2, dim=-1, eps=1e-12) 
      
          cls_scores = self.scale_cls * (torch.matmul(cls_weights,features.t()))
          cls_scores = cls_scores.permute(0,2,1)
          cls_scores = torch.diagonal(cls_scores,offset=0,dim1=0,dim2=1)
          cls_scores = cls_scores.t()
          if norm:
              return cls_scores,cls_scores/scale_cls
          else:
                  return cls_scores
  
      def IASatt(self,features,AttM,tmp):
          cls_num = len(AttM)
          attdims = AttM.shape[1]
          atten = torch.mm(features,self.w_IAS)+self.b_IAS
          atten = F.relu(atten)
          atten = F.softmax(atten/tmp,dim=1).reshape(-1,1,attdims)
          atten = atten + torch.ones_like(atten)
          atten = atten.repeat(1,cls_num,1)
          AttM = AttM.unsqueeze(0)
          AttM = AttM.repeat(len(features),1,1)
          AttM = atten*AttM
          return AttM
      
      def forward(self,att,features,tmp):
              
          features = features.squeeze()
          att = self.IASatt(features,att,tmp)
  
          a1 = F.relu(torch.matmul(att, self.w1) + self.b1)
          a2 = F.relu(torch.matmul(a1, self.w2) + self.b2)
          return a2
  
      def loadDataFromCustom(self,dataType):
          """
              Load Data from Xlsa:
              ------
  
          """
          print('load from customed')
  
          with open('./dataset/finetune/'+self.cfg['dataset']['name']+'/dvbeExtracted.pkl','rb') as f:
              ExtractFeature = pickle.load(f)
          imgfs = np.array(ExtractFeature[dataType]['features'])
          labels = np.array(ExtractFeature[dataType]['labels'])
  
          print(imgfs.shape)
          print(labels.shape)
          select_id, idx = np.unique(labels,return_inverse=True)
          if dataType.lower() == "test_unseen":
              labels = idx + len(self.seen_labels)
  
          else:
              labels = idx
  
          imgfs = torch.FloatTensor(imgfs.reshape(-1,2048))
          labels = torch.LongTensor(labels)
          dataset = torch.utils.data.TensorDataset(imgfs,labels)
          dataloader = torch.utils.data.DataLoader(dataset,batch_size=self.cfg['dataset_hyp']['batch_size'],shuffle=False)
          return dataloader
  
      def loadChk(self,chkFile1,chkFile2):
          """
              if resume training
              args:
                  chkFile
              return:
                  bestAcc in chkFile
          """
  
          print("=> load checkpoint '{}'".format(chkFile1))
          chk = torch.load(chkFile1)
          self.cfg['model_hyp']['start_epoch'] = chk['epoch']
          print(chk['model'].keys())
          self.model1.load_state_dict(chk['model'])
  
          print("=> load checkpoint '{}'".format(chkFile2))
          
          model = torch.load(chkFile2)
          self.w1 = Variable(model['w1'], requires_grad=False)
          self.b1 = Variable(model['b1'], requires_grad=False)
          self.w2 = Variable(model['w2'], requires_grad=False)
          self.b2 = Variable(model['b2'], requires_grad=False)
          self.w_IAS = Variable(model['w_IAS'],requires_grad=False)
          self.b_IAS = Variable(model['b_IAS'],requires_grad=False)
          self.scale_cls = Variable(model['scale_cls'], requires_grad=False)
          self.bias = Variable(model['bias'], requires_grad=False)
          return 
  
      def evaluationSplit(self,dataloader,dataType='seen',accType='allCls',imgfType='customed',recordFile=None,seenOrunseenCls=None):
          """
          evaluation model on input dataloader
          args:
              dataloader
              dataType: 'seen'   : count acc, CEloss, MSEloss 
                        'unseen' : count acc, CEloss(1), MSEloss 
              accType:  'allCls' : return acc mean of all cls 
                        'allImg' : return acc mean of all images
              recordFile: None   : do not record
                          'print'： print eval
                         filepath: record in this file
          return:
              accs: (int) mean of all cls or all images 
                          decide by accType
              accs_cls: (dict) mean acc according to cls
              generalAccs_cls: (dict) mean generalAccs 2 cls,
              splitacc: datasplit bw seen and unseen
              cos_cls:(dict){cls:[sample_cos]} sample cosimilarity 
          """ 
  
          total_spp_loss = []
          generalAccs = []
          accs = []
          generalAccs_cls = {}
          accs_cls = {}
          cos_cls = {}
          spp_loss_cls = {}
          splitaccs = []
          TF_l = []
          cos = []
  
          with torch.no_grad():
              for images,targets in tqdm(dataloader):
                  images = images.to(self.device)
                  targets = targets.to(self.device)
  
                  if self.cfg['dataset_hyp']['imgfType'] == 'customed':
                      images = self.caffeRes101(images).squeeze()
  
                  batch_visual_norm = F.normalize(images, p=2, dim=images.dim()-1, eps=1e-12)                
                  generalpreds1,generalpreds1_n = self.model1(self.attMatrix,images,TrainOrTest="Test")
                  batch_weights = self.forward(self.attMatrix[len(self.seen_labels):],batch_visual_norm,tmp=TMP)       
                  all_cls_weights = batch_weights
                  generalpreds2 = self.apply_classification_weights(batch_visual_norm, all_cls_weights)
                  generalpreds1 = torch.argmax(generalpreds1,dim=1)
                  generalpreds2 = torch.argmax(generalpreds2,dim=1)+len(self.seen_labels)
                  generalpreds = generalpreds1
                  for i in range(len(generalpreds1)):
                      if generalpreds1[i] > len(self.seen_labels)-1: # judge by model as unseen class
                          generalpreds[i] = generalpreds2[i]
                 
                  if dataType.lower()=='seen':
                      preds,_ = self.model1(self.attMatrix[:len(self.seen_labels)],images,TrainOrTest="Test")
                      preds = torch.argmax(preds,dim=1)
                  elif dataType.lower()=='unseen':
                      preds = generalpreds2
  
                  # seen or unseen split accuracy
                  splitacc = [ i in seenOrunseenCls for i in generalpreds1.cpu().detach().tolist()]
                  splitaccs.extend(splitacc)
  
                  # calculate accs
                  acc = preds == targets
                  TF_l.extend(acc.cpu().detach().tolist())
                  max_cos,_ = torch.max(generalpreds1_n,dim=-1)
                  cos.extend(max_cos.cpu().detach().tolist())
  
                  # calculate accs_cls
                  for c in torch.unique(targets):
                      loc = targets == c
                      acc_c = preds[loc] == targets[loc]
                      generalacc_c = generalpreds[loc] == targets[loc]
                      dist_c = generalpreds1_n[loc,c].cpu().detach().tolist()
                      acc_c = acc_c.cpu().detach().tolist()
                      generalacc_c = generalacc_c.cpu().detach().tolist()
                      try:
                          accs_cls[c.cpu().detach().tolist()].extend(acc_c)
                          generalAccs_cls[c.cpu().detach().tolist()].extend(generalacc_c)
                          cos_cls[c.cpu().detach().tolist()].extend(dist_c)
                      except(KeyError):
                          accs_cls[c.cpu().detach().tolist()] = acc_c
                          generalAccs_cls[c.cpu().detach().tolist()] = generalacc_c
                          cos_cls[c.cpu().detach().tolist()] = dist_c
  
          # average class accuracy
          if accType == 'allCls':
              for c in accs_cls:
                  accs.append(np.mean(accs_cls[c]))
                  generalAccs.append(np.mean(generalAccs_cls[c]))
          for c in accs_cls:
              accs_cls[c] = np.mean(accs_cls[c])*100
              generalAccs_cls[c] = np.mean(generalAccs_cls[c])*100
  
          return  {'accs':np.mean(accs)*100,
                   'generalAccs':np.mean(generalAccs)*100,
                   'generalAccs_cls':generalAccs_cls,
                   'accs_cls':accs_cls,
                   'splitacc':np.mean(splitaccs)*100,
                   'cos_cls':cos_cls,
                   'TrueFalseList':TF_l,
                   'cos':cos
                  }
  
      def prepareModel(self,):
          self.createLinearModel()
  
      def evalModel(self,chkFile1,chkFile2,resume=False):
          """
              train Model #epochs defined in cfg
  
              args:
                  chkFile: save model chkpoint
          """
  
          testSeenData = self.loadDataFromCustom(dataType='test_seen')
          testUnseenData = self.loadDataFromCustom(dataType='test_unseen')
  
          self.loadChk(chkFile1,chkFile2)
          
          testSeenEval = self.evaluationSplit(testSeenData,dataType='seen',accType=self.cfg['model_hyp']['accType'],imgfType=self.cfg['dataset_hyp']['imgfType'],seenOrunseenCls=np.arange(0,len(self.seen_labels)))
          testUnseenEval = self.evaluationSplit(testUnseenData,dataType='unseen',accType=self.cfg['model_hyp']['accType'],imgfType=self.cfg['dataset_hyp']['imgfType'],seenOrunseenCls=np.arange(len(self.seen_labels),len(self.seen_labels)+len(self.unseen_labels)))
  
          accu = testUnseenEval['generalAccs']
          accs = testSeenEval['generalAccs']
          H = 2*(accu*accs)/(accu+accs)
          print('ZSL(T1):',testUnseenEval['accs'])   # ZSL(T1)
  #         print('test/SeenAcc',testSeenEval['accs'])
          print('GZSL H:',H)                                # H
          print('GZSL Seen:',accs)                   # GZSL Seen
          print('GZSL Unseen:',accu)                   # GZSL Unseen
  #         print('test/fu',testUnseenEval['splitacc']) 
  #         print('test/fs',testSeenEval['splitacc'])
          return testSeenEval['TrueFalseList'],testSeenEval['cos']
  
  if __name__ =="__main__":
  
      parser = argparse.ArgumentParser() 
      parser.add_argument('--modelChk',type=str,help='model chkpoint file')
      parser.add_argument('--yaml',type=str,help='checkpoint file saved in yaml')
      parser.add_argument('--device',default='cuda:0',type=str,help='cuda:0,cuda:1,cuda:2,cpu'
  )
      opt = parser.parse_args()
  
      #chkFile1 ='./chk/201115CUBS2IFinetune_ResNet101_fullv_r520_IAStmp10CustomedExtractedCos520LR5e-4WD5e-7/chk_best_Hs.pt' #class weight + IAS
      chkFile1 ='./chk/CUBLR5e-4WD9e-7/chk_best_Hs.pt' #class weight + IAS
      chkFile2 = './chk/cublr5e-4_opt5e-7_w20_s4_51FinetunebestunseenAcc.pt'#mixup_IAS
      device = torch.device(opt.device) 
  
      t = train(device,opt.yaml,IAS=True)
      t.prepareModel()
  
      truefalselistBase,cos_Base = t.evalModel(chkFile1,chkFile2)
      
  
  
  