import sys
import os
import json
import tensorflow as tf

workspace_path = '/cfs/cfs-6kz7e0an/mecreeli/content-understanding-UDA'
python = "/data/anaconda3/bin/python"
out_eval = './output_eval'
out_tmp = './output_tmp'
out = './output'
rub_path = '/cfs/cfs-6kz7e0an/mecreeli/rub/0'
pre_conf = '/cfs/cfs-6kz7e0an/mecreeli/content-understanding-UDA/demos/restores/multimodal_predict/strategy.config'
fea_conf = '/cfs/cfs-6kz7e0an/mecreeli/content-understanding-UDA/demos/restores/multimodal_train/feature.config'
res_pre = '/cfs/cfs-6kz7e0an/mecreeli/content-understanding-UDA/res_predict'

def run_init():
    # for running on venus
    os.system('pip install pandas --user -i https://mirrors.tencent.com/pypi/simple/')
    os.system('pip install jieba --user -i https://mirrors.tencent.com/pypi/simple/')
    os.system('pip install torch==1.2.0 torchvision==0.4.0 --user -i https://mirrors.tencent.com/pypi/simple/')
    os.system('pip install einops --user -i https://mirrors.tencent.com/pypi/simple/')
    os.system('pip install pytorch_pretrained_bert --user -i https://mirrors.tencent.com/pypi/simple/')
    # os.system('pip install efficientnet_pytorch --user -i https://mirrors.tencent.com/pypi/simple/')

def run_multimodal_feat_record(python, workspace_path):
    sup_tf_path = os.path.join(workspace_path, './demos/restores/multimodal_train/train_data.tf_record')
    unsup_tf_path = os.path.join(workspace_path, './demos/restores/multimodal_train/unsup_data.tf_record')
    aug_unsup_tf_path = os.path.join(workspace_path, './demos/restores/multimodal_train/aug_unsup_data.tf_record')
    # 查詢是否存在先前的tf_record并刪除
    # if os.path.exists(sup_tf_path):
    os.system('rm -rf ' + sup_tf_path)
    os.system('rm -rf ' + unsup_tf_path)
    os.system('rm -rf ' + aug_unsup_tf_path)
    path = workspace_path + '/strategy/'
    os.system(python + ' ' + path + 'multimodal_feat_record_strategy.py')

def run_multimodal_train(python, workspace_path):
    path = workspace_path + '/strategy/'
    os.system(python + ' ' + path + 'multimodal_train_strategy.py')

def run_multimodal_eval(python, workspace_path):
    path = workspace_path + '/strategy/'
    os.system(python + ' ' + path + 'multimodal_eval_strategy.py')

def run_multimodal_predict(python, workspace_path):
    path = workspace_path + '/strategy/'
    os.system(python + ' ' + path + 'multimodal_predict_strategy.py')


def main():
    run_init()

    for tmp in range(1000, 4001, 1000):
        train_data = 'corpus/train-rank_1-29/train_rank_{}'.format(tmp)

        with open(fea_conf, 'rb') as f:
            data_fea = json.load(f)
        data_fea['TFExampleExtractor']['file_dict']['train_data'] = train_data
        with open(fea_conf, 'w', encoding='utf-8') as f:
            json.dump(data_fea, f)

        with open(res_pre, 'a') as f:
            f.write('train_rank_{}'.format(tmp)+'\n')

        # 特征提取
        run_multimodal_feat_record(python, workspace_path)
        # 訓(xùn)練
        run_multimodal_train(python, workspace_path)
        # 預(yù)測(cè)
        models_path = os.path.join(workspace_path, 'output')
        os.system('rm -rf ' + os.path.join(workspace_path, 'output_al'))
        os.system('mv ' + models_path + ' ' + os.path.join(workspace_path, 'output_al'))
        for i in range(500, 500*8+1, 500):
            with open(pre_conf, 'rb') as f:
                data = json.load(f)
            data['init_checkpoint'] = './output_al/model.ckpt-{}'.format(i)
            with open(pre_conf, 'w', encoding='utf-8') as f:
                json.dump(data, f)
            run_multimodal_predict(python, workspace_path)
        # last_ckpt 預(yù)測(cè)
        with open(pre_conf, 'rb') as f:
            data = json.load(f)
        data['init_checkpoint'] = tf.train.latest_checkpoint('/cfs/cfs-6kz7e0an/mecreeli/content-understanding-UDA/output_al').replace('/cfs/cfs-6kz7e0an/mecreeli/content-understanding-UDA/output_al', './output_al')
        with open(pre_conf, 'w', encoding='utf-8') as f:
            json.dump(data, f)
        run_multimodal_predict(python, workspace_path)

if __name__ == '__main__':
    main()