#!/usr/bin/python3

"""Train module.

Usage:
    python train.py <task> <model_name>

"""



import argparse
from os.path import isfile, join
from typing import List, Dict, Union
import numpy as np
import warnings
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '7'  # 7

# from scipy import sparse
# 忽略代码中输出的警告
warnings.filterwarnings("ignore")
import torch
import torch.nn as nn
from config import GlobalConfig
from constant import (TEXT_TASK,)
from constant import TASK_STR
from dataset import Dataset
from lib.train import text_train
from util import load_pkl, get_embed_init, load_pkl2
from torch.cuda.amp import autocast as autocast, GradScaler

TASKS: List[str] = list(TASK_STR.values())
# base_dir = '/data/chenxiaolin/dataset_dump/'
base_dir = '/data/chenxiaolin/dataset_user/'

knowledge_dir = '/data/chenxiaolin/knowledge_extract/triplet/'
path_tip_vec = knowledge_dir + 'tip_vec_corresponding.npy'
path_triplet_vec = knowledge_dir + 'triplet_vec_corresponding.npy'
tip_vec = np.load(path_tip_vec).tolist() #tip_vec,triplet_vec
triplet_vec = np.load(path_triplet_vec).tolist()


path_train_context_shuffle = base_dir + 'train_context_shuffle.npy'
path_valid_context_shuffle = base_dir + 'valid_context_shuffle.npy'
path_test_context_shuffle = base_dir + 'test_context_shuffle.npy'

train_context_index = np.load(path_train_context_shuffle).tolist()  #train_context_index,valid_context_index, test_context_index
valid_context_index = np.load(path_valid_context_shuffle).tolist()
test_context_index = np.load(path_test_context_shuffle).tolist()

train_context_dir = base_dir + 'Context/train/'
valid_context_dir = base_dir + 'Context/validate/' # valid_context_dir, test_context_dir
test_context_dir = base_dir + 'Context/test/'

train_img_fea_path = '/data/chenxiaolin/Img_ViT_Feature/train/'
valid_img_fea_path = '/data/chenxiaolin/Img_ViT_Feature/validate/'
test_img_fea_path = '/data/chenxiaolin/Img_ViT_Feature/test/'

# ==origin==
# train_hyper_dir = base_dir + 'Hyper_kg/train/'
# valid_hyper_dir = base_dir + 'Hyper_kg/validate/'
# test_hyper_dir = base_dir + 'Hyper_kg/test/'

# train_hyper_dir = base_dir + 'Hyper_kg_add_V/train/'
# valid_hyper_dir = base_dir + 'Hyper_kg_add_V/validate/'
# test_hyper_dir = base_dir + 'Hyper_kg_add_V/test/'

train_hyper_dir = '/data/chenxiaolin/dataset_dump/Hyper_kg_add_V/train/'
valid_hyper_dir = '/data/chenxiaolin/dataset_dump/Hyper_kg_add_V/validate/'
test_hyper_dir = '/data/chenxiaolin/dataset_dump/Hyper_kg_add_V/test/'

# train_context_example_path = '/data/chenxiaolin/dataset_user/context_top_5_context_dict.npy'
# valid_context_example_path = '/data/chenxiaolin/dataset_user/context_top_5_context_dict_validate.npy'
# test_context_example_path = '/data/chenxiaolin/dataset_user/context_top_5_context_dict_test.npy'

train_context_example_path = '/data/chenxiaolin/dataset_user/User_related_file/context_top_5_context_dict_user_filter_train.npy'
valid_context_example_path = '/data/chenxiaolin/dataset_user/User_related_file/context_top_5_context_dict_user_filter_validate.npy'
test_context_example_path = '/data/chenxiaolin/dataset_user/User_related_file/context_top_5_context_dict_user_filter_test.npy'


def train(task: int, model_file_name: str):
    """Train model.

    Args:
        task (int): Task.
        model_file_name (str): Model file name (saved or to be saved).

    """

    path_dialog_vocab = base_dir + 'dialog_vocab.npy'
    path_glove = base_dir + 'glove.npy'

    dialog_vocab = np.load(path_dialog_vocab, allow_pickle=True).tolist()
    glove = np.load(path_glove, allow_pickle=True).tolist()
    print('vocab_glove loaded!')

    # Dataset wrap.
    # dialog_vacb; global_repre; 超图卷积
    train_dataset = Dataset(task, train_context_dir, train_img_fea_path, train_hyper_dir, dialog_vocab,  train_context_index, train_context_example_path)
    valid_dataset = Dataset(task, valid_context_dir, valid_img_fea_path,  valid_hyper_dir, dialog_vocab, valid_context_index, valid_context_example_path)
    test_dataset = Dataset(task, test_context_dir, test_img_fea_path, test_hyper_dir, dialog_vocab,  test_context_index, test_context_example_path)

    
    vocab_size = 5155 # 4892
    embed_init = get_embed_init(
        glove, vocab_size)
    # 4892, 300
    embed_init = embed_init.to(GlobalConfig.device)


    model_dump_dir = '/data/chenxiaolin/model_file/'
    # Load model file.
    # model_file = join(DatasetConfig.dump_dir, model_file_name)
    model_file = join(model_dump_dir, model_file_name)
    scaler = GradScaler()

    # Task-specific parts---text/image
    if task == TEXT_TASK:
        print('----text_train---------')
        text_train(
                    train_dataset,
                    valid_dataset,
                    test_dataset,                    
                    model_file,
                    dialog_vocab,
                    scaler,
                    embed_init
                )
 

# (:形参数据类型,->返回数据类型)
# -->代表返回数据类型
def parse_cmd() -> Dict[str, List[str]]:
    """Parse commandline parameters.

    Returns:
        Dict[str, List[str]]: Parse result.

    """

    # Definition of argument parser.
    parser = argparse.ArgumentParser(description='Train.')
    parser.add_argument(
        'task',
        metavar='<task>',
        choices=TASKS,
        default='text',
        help='task ({})'.format('/'.join(TASKS))
    )
    parser.add_argument('model_file', metavar='<model_file>',default='model')

    # Namespace -> Dict
    parse_res: Dict[str, Union[List[str], str]] = vars(parser.parse_args())
    return parse_res


def main():
    # Parse commandline parameters and standardize.
    # {'task': 'text', 'model_file': '/home/share/chenxiaolin/MultimodalDialogSystem/Baseline/MAGIC_MM19/model_file/model.json'}
    # 以字典的形式存放了task以及保存模型的数据路径
    # parse_result: Dict[str, Union[List[str], str]] = parse_cmd()
    # print(parse_result)
    # exit()
    # task: int = TEXT_TASK
    model_file = 'model_269'
    # train(task, model_file)
    # parse_result: Dict[str, Union[List[str], str]] = parse_cmd()
    # exit()
    task: int = TEXT_TASK
    # model_file: str = parse_result['model_file']
    train(task, model_file)

if __name__ == '__main__':
    torch.manual_seed(11132)  #7029 90233 902522 902544 90273 10402 11065 11093 11171 11181 11291 12034 12035 12046 120611 12082 121212 30705 40181 40182 40191 40193 50092 50114 50112 50114 50141 50151 50152 50153 50155 502501 50271 60014 60015 60016 60017 60064 60171 80505 80506 90271 90303 10021 10022 10024 10162 10165 10182 10183 10184 10185 10203 10234 11131
    torch.cuda.manual_seed_all(11132)
    main()
