#!/usr/bin/python3.6
# -*- coding: utf-8 -*-
# @Time    : 2020/9/3 上午11:03
# @Author  : Joselynzhao
# @Email   : zhaojing17@forxmail.com
# @File    : atmpro1_vsm2.py
# @Software: PyCharm
# @Desc    :

#!/usr/bin/python3.6
# -*- coding: utf-8 -*-
# @Time    : 2020/9/1 下午7:07
# @Author  : Joselynzhao
# @Email   : zhaojing17@forxmail.com
# @File    : atmpro1_vsm.py
# @Software: PyCharm
# @Desc    :

#!/usr/bin/python3.6
# -*- coding: utf-8 -*-
# @Time    : 2020/8/26 下午8:26
# @Author  : Joselynzhao
# @Email   : zhaojing17@forxmail.com
# @File    : atmpro1.py
# @Software: PyCharm
# @Desc    :


from my_reid.eug import *
from my_reid import datasets
from my_reid import models
import numpy as np
import torch
import argparse
import os

import warnings

warnings.filterwarnings("ignore")

from my_reid.utils.logging import Logger
import os.path as osp
import sys
from torch.backends import cudnn
from my_reid.utils.serialization import load_checkpoint
from torch import nn
import time
import pickle
import torch.distributed as dist
from torch.nn.parallel import  DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler
from  pathlib import  Path



def resume(savepath):
    import re
    pattern = re.compile(r'step_(\d+)\.ckpt')
    start_step = -1
    ckpt_file = ""

    # find start step
    files = os.listdir(savepath)
    files.sort()
    for filename in files:
        try:
            iter_ = int(pattern.search(filename).groups()[0])
            print(iter_)
            if iter_ > start_step:
                start_step = iter_
                ckpt_file = osp.join(savepath, filename)
        except:
            continue

    # if need resume
    if start_step >= 0:
        print("continued from iter step", start_step)
    else:
        print("resume failed", start_step, files)
    return start_step, ckpt_file


def main(args):
    father = Path('/mnt/')
    if father.exists(): # 是在服务器上
        data_dir = Path('/mnt/share/datasets/RE-ID/data')  # 服务器
        logs_dir = Path('/mnt/home/{}'.format(args.log_name))  # 服务器
    else: #本地
        data_dir = Path('/home/joselyn/workspace/ATM_SERIES/data')  # 本地跑用这个
        logs_dir = Path('/home/joselyn/workspace/ATM_SERIES/{}'.format(args.log_name))  # 本地跑用这个


    cudnn.benchmark = True
    cudnn.enabled = True
    save_path = os.path.join(logs_dir, args.dataset, args.exp_name, args.exp_order)  # 到编号位置.
    total_step = 100 // args.EF + 1
    sys.stdout = Logger(osp.join(save_path, 'log' + str(args.EF) + time.strftime(".%m_%d_%H:%M:%S") + '.txt'))
    dataf_file = open(osp.join(save_path, 'dataf.txt'), 'a')  # 保存性能数据.  #特征空间中的性能问题.
    data_file = open(osp.join(save_path, 'data.txt'), 'a')  # 保存性能数据.  #特征空间中的性能问题.
    kf_file = open(osp.join(save_path,'kf.txt'),'a')
    # 数据格式为 label_pre_r, select_pre_r,label_pre_t, select_pre_t  ,加上了了tagper的数据.
    tagper_path = osp.join(save_path,'tagper')  #tagper存储路径.
    if not Path(tagper_path).exists():
        os.mkdir(tagper_path)


    '''# 记录配置信息 和路径'''
    print('-'*20+'config_info'+'-'*20)
    config_file = open(osp.join(save_path, 'config.txt'), 'w')
    config_info = str(args).split('(')[1].strip(')').split(',')
    config_info.sort()
    for one in config_info:
        key,value=map(str,one.split('='))
        config_file.write(key.strip()+'='+value.strip('\'')+'\n')
        print(key.strip()+'='+value.strip('\''))
    config_file.write('save_path='+save_path)
    print('save_path='+save_path)
    print('-' * 20 + 'config_info' + '-' * 20)
    config_file.close()

    train_time_file = open(osp.join(save_path, 'time.txt'), 'a')  # 只记录训练所需要的时间.
    # 数据格式为 step_time total_time.
    total_time = 0

    # get all the labeled and unlabeled data for training

    dataset_all = datasets.create(args.dataset, osp.join(data_dir, args.dataset))
    num_all_examples = len(dataset_all.train)
    l_data, u_data = get_init_shot_in_cam1(dataset_all,
                                           load_path="./examples/{}_init_{}.pickle".format(dataset_all.name, args.init),
                                           init=args.init)

    resume_step, ckpt_file = -1, ''
    if args.resume:
        resume_step, ckpt_file = resume(save_path)

        # initial the EUG algorithm
    eug = EUG(batch_size=args.batch_size, num_classes=dataset_all.num_train_ids,
              dataset=dataset_all, l_data=l_data, u_data=u_data, save_path=save_path, max_frames=args.max_frames,
              embeding_fea_size=args.fea, momentum=args.momentum, lamda=args.lamda)
    tagper = EUG(batch_size=args.batch_size, num_classes=dataset_all.num_train_ids,
                 dataset=dataset_all, l_data=l_data, u_data=u_data, save_path=tagper_path,
                 max_frames=args.max_frames,
                 embeding_fea_size=args.fea, momentum=args.momentum, lamda=args.lamda)

    new_train_data = l_data
    unselected_data = u_data
    iter_mode = 2 #迭代模式,确定是否训练tagper
    for step in range(total_step):
        # for resume
        if step < resume_step:
            continue

        ratio = (step + 1) * args.EF / 100
        ratio_t = (step+1+args.t) * args.EF /100
        nums_to_select = int(len(u_data) * ratio)
        nums_to_select_tagper = int(len(u_data) * ratio_t)
        if nums_to_select >= len(u_data):
            break

        #args.vsm_lambda的衰减  0.5 - 0
        vsm_lambda = args.vsm_lambda*step/(1-(total_step/2)) +args.vsm_lambda
        vsm_lambda +=1


        print("Runing: EF={}%, step {}:\t Nums_to_be_select {} \t Ritio \t Logs-dir {}".format(
            args.EF, step, nums_to_select, ratio, save_path))

        # train the model or load ckpt
        start_time = time.time()
        print("training reid model")
        eug.train(new_train_data, unselected_data, step, loss=args.loss, epochs=args.epochs, step_size=args.step_size,
                  init_lr=0.1) if step != resume_step else eug.resume(ckpt_file, step)
        # 只对eug进行性能评估
        # mAP, rank1, rank5, rank10, rank20 = 0, 0, 0, 0, 0
        mAP, rank1, rank5, rank10, rank20 = eug.evaluate(dataset_all.query, dataset_all.gallery)
        # 把数据写到data文件里.
        data_file.write('{} {:.2%} {:.2%} {:.2%} {:.2%} {:.2%}\n'.format(step, mAP, rank1, rank5, rank10, rank20))

        pred_y, pred_score,label_pre,dists= eug.estimate_label_vsm()
        selected_idx = eug.select_top_data_vsm2(pred_score, dists,args.topk,vsm_lambda,min(nums_to_select_tagper,len(u_data)-50) if iter_mode==2 else min(nums_to_select,len(u_data)))   #直接翻两倍取数据. -50个样本,保证unselected_data数量不为0
        new_train_data, unselected_data, select_pre= eug.generate_new_train_data(selected_idx, pred_y)
        raw_label_pre, raw_select_pre = label_pre,select_pre
        t_label_pre,t_select_pre = 0,0
        raw_select_pre_t = 0
        # label_pre_t,select_pre_t=0,0
        if iter_mode==2:
            raw_select_pre_t = raw_select_pre
            print("training tagper model")
            selected_idx = eug.select_top_data_vsm2(pred_score,dists,args.topk,vsm_lambda, min(nums_to_select, len(u_data)))
            _, _, raw_select_pre = eug.generate_new_train_data(selected_idx, pred_y)
            # kf_file.write('{} {:.2%} {:.2%}'.format(step, label_pre, select_pre))

            tagper.resume(osp.join(save_path,'step_{}.ckpt'.format(step)),step)
            tagper.train(new_train_data, unselected_data, step, loss=args.loss, epochs=args.epochs, step_size=args.step_size, init_lr=0.1)

            pred_y, pred_score, label_pre,dists= tagper.estimate_label_vsm()
            selected_idx = tagper.select_top_data_vsm2(pred_score,dists,args.topk,vsm_lambda,min(nums_to_select,len(u_data)))  # 采样目标数量
            new_train_data, unselected_data, select_pre= tagper.generate_new_train_data(selected_idx, pred_y)
            t_label_pre,t_select_pre = label_pre,select_pre
            label_pre,select_pre = t_label_pre,t_select_pre
            if nums_to_select_tagper >=len(u_data):
                iter_mode=1 #切换模式
                print('tagper is stop')
        else: #mode = 1
            # raw_select_pre = raw_select_pre_t
            # raw_select_pre_t = 0
            label_pre,select_pre = raw_label_pre,raw_select_pre

        end_time = time.time()
        step_time = end_time - start_time
        total_time = step_time + total_time
        train_time_file.write('{} {:.6} {:.6}\n'.format(step, step_time, total_time))
        kf_file.write('{} {} {} {:.2%} {:.2%} {:.2%} {:.2%} {:.2%}\n'.format(step,nums_to_select,nums_to_select_tagper,raw_label_pre,raw_select_pre,raw_select_pre_t,t_label_pre,t_select_pre))
        dataf_file.write(
            '{} {:.2%} {:.2%}\n'.format(step, label_pre, select_pre))
    dataf_file.close()
    train_time_file.close()


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Progressive Learning for One-Example re-ID')
    parser.add_argument('-d', '--dataset', type=str, default='mars',
                        choices=datasets.names())
    parser.add_argument('-b', '--batch-size', type=int, default=16)
    parser.add_argument('-f', '--fea', type=int, default=1024)
    parser.add_argument('--EF', type=int, default=10)
    parser.add_argument('--t', type=float, default=2) #不再tagper采样的倍率, 而是表示跨多少个step采样.
    parser.add_argument('--exp_order', type=str, default='0')
    parser.add_argument('--exp_name', type=str, default='atm')
    parser.add_argument('--exp_aim', type=str, default='for paper')
    parser.add_argument('--run_file',type=str,default='train.py')
    parser.add_argument('--log_name',type=str,default='pl_logs')
    parser.add_argument('--topk',type=int,default=2)
    parser.add_argument('--vsm_lambda',type=float,default=0.5)
    parser.add_argument('--resume', type=str, default='Yes')
    parser.add_argument('--max_frames', type=int, default=900)
    parser.add_argument('--loss', type=str, default='ExLoss', choices=['CrossEntropyLoss', 'ExLoss'])
    parser.add_argument('--init', type=float, default=-1)
    parser.add_argument('-m', '--momentum', type=float, default=0.5)
    parser.add_argument('-e', '--epochs', type=int, default=70)
    parser.add_argument('-s', '--step_size', type=int, default=55)
    parser.add_argument('--lamda', type=float, default=0.5)
    main(parser.parse_args())

[((39, 0, 39, 33), 'warnings.filterwarnings', 'warnings.filterwarnings', ({(39, 24, 39, 32): '"""ignore"""'}, {}), "('ignore')", False, 'import warnings\n'), ((58, 14, 58, 45), 're.compile', 're.compile', ({(58, 25, 58, 44): '"""step_(\\\\d+)\\\\.ckpt"""'}, {}), "('step_(\\\\d+)\\\\.ckpt')", False, 'import re\n'), ((63, 12, 63, 32), 'os.listdir', 'os.listdir', ({(63, 23, 63, 31): 'savepath'}, {}), '(savepath)', False, 'import os\n'), ((84, 13, 84, 26), 'pathlib.Path', 'Path', ({(84, 18, 84, 25): '"""/mnt/"""'}, {}), "('/mnt/')", False, 'from pathlib import Path\n'), ((95, 16, 95, 83), 'os.path.join', 'os.path.join', ({(95, 29, 95, 37): 'logs_dir', (95, 39, 95, 51): 'args.dataset', (95, 53, 95, 66): 'args.exp_name', (95, 68, 95, 82): 'args.exp_order'}, {}), '(logs_dir, args.dataset, args.exp_name, args.exp_order)', False, 'import os\n'), ((102, 18, 102, 46), 'os.path.join', 'osp.join', ({(102, 27, 102, 36): 'save_path', (102, 37, 102, 45): '"""tagper"""'}, {}), "(save_path, 'tagper')", True, 'import os.path as osp\n'), ((222, 13, 222, 94), 'argparse.ArgumentParser', 'argparse.ArgumentParser', (), '', False, 'import argparse\n'), ((86, 19, 86, 57), 'pathlib.Path', 'Path', ({(86, 24, 86, 56): '"""/mnt/share/datasets/RE-ID/data"""'}, {}), "('/mnt/share/datasets/RE-ID/data')", False, 'from pathlib import Path\n'), ((89, 19, 89, 66), 'pathlib.Path', 'Path', ({(89, 24, 89, 65): '"""/home/joselyn/workspace/ATM_SERIES/data"""'}, {}), "('/home/joselyn/workspace/ATM_SERIES/data')", False, 'from pathlib import Path\n'), ((98, 22, 98, 54), 'os.path.join', 'osp.join', ({(98, 31, 98, 40): 'save_path', (98, 42, 98, 53): '"""dataf.txt"""'}, {}), "(save_path, 'dataf.txt')", True, 'import os.path as osp\n'), ((99, 21, 99, 52), 'os.path.join', 'osp.join', ({(99, 30, 99, 39): 'save_path', (99, 41, 99, 51): '"""data.txt"""'}, {}), "(save_path, 'data.txt')", True, 'import os.path as osp\n'), ((100, 19, 100, 47), 'os.path.join', 'osp.join', ({(100, 28, 100, 37): 'save_path', (100, 38, 100, 46): '"""kf.txt"""'}, {}), "(save_path, 'kf.txt')", True, 'import os.path as osp\n'), ((104, 8, 104, 29), 'os.mkdir', 'os.mkdir', ({(104, 17, 104, 28): 'tagper_path'}, {}), '(tagper_path)', False, 'import os\n'), ((109, 23, 109, 56), 'os.path.join', 'osp.join', ({(109, 32, 109, 41): 'save_path', (109, 43, 109, 55): '"""config.txt"""'}, {}), "(save_path, 'config.txt')", True, 'import os.path as osp\n'), ((121, 27, 121, 58), 'os.path.join', 'osp.join', ({(121, 36, 121, 45): 'save_path', (121, 47, 121, 57): '"""time.txt"""'}, {}), "(save_path, 'time.txt')", True, 'import os.path as osp\n'), ((127, 48, 127, 80), 'os.path.join', 'osp.join', ({(127, 57, 127, 65): 'data_dir', (127, 67, 127, 79): 'args.dataset'}, {}), '(data_dir, args.dataset)', True, 'import os.path as osp\n'), ((170, 21, 170, 32), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((210, 19, 210, 30), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((224, 32, 224, 48), 'my_reid.datasets.names', 'datasets.names', ({}, {}), '()', False, 'from my_reid import datasets\n'), ((71, 28, 71, 56), 'os.path.join', 'osp.join', ({(71, 37, 71, 45): 'savepath', (71, 47, 71, 55): 'filename'}, {}), '(savepath, filename)', True, 'import os.path as osp\n'), ((103, 11, 103, 28), 'pathlib.Path', 'Path', ({(103, 16, 103, 27): 'tagper_path'}, {}), '(tagper_path)', False, 'from pathlib import Path\n'), ((97, 67, 97, 99), 'time.strftime', 'time.strftime', ({(97, 81, 97, 98): '""".%m_%d_%H:%M:%S"""'}, {}), "('.%m_%d_%H:%M:%S')", False, 'import time\n')]