main.py

"""
main.py

此模块定义了模型训练和评估的核心逻辑。它包含了数据生产者和消费者的实现，训练和评估模型的函数，以及模型准备和保存的功能。通过使用多线程和多进程技术，该模块能够高效地处理大规模数据集，并在训练过程中动态评估模型性能。整个流程包括数据加载、模型训练、性能评估以及结果记录等关键步骤，为整个项目的运行提供了基础支持。
"""

from concurrent import futures
from concurrent.futures.thread import ThreadPoolExecutor
import json
from logging import FileHandler, Formatter
import logging
import os
from queue import Queue
import random
import shutil
from threading import Lock
import traceback
import tracemalloc

import numpy
from torch import multiprocessing
import torch
from torch.nn import utils
from torch.optim.adamw import AdamW

import config
import constants
from data import build_dataset, get_data_meta
from model import FMLSTMAttentionModel, FMLlamaModel


def dev(problem_type, out_mode, model_mode, log_mode, models, trainable_models, group_size, queue_size, thread_size,
        data_size, sensitive_rate, data, feature_size, seq_size, buffer_size, batch_size, num_workers, is_multiprocess,
        device, epoch, dev_type):
    """
    评测模型。

    :param problem_type: bool - 为分类任务时为 True，为回归任务时为 False
    :param out_mode: bool - 为整体模型输出时为 True，为单独模型输出时为 False
    :param model_mode: int - 为 FMLSTMAttentionModel 时为 0，为 FMLlamaModel 时为 1
    :param log_mode: bool - 为 True 时记录日志，为 False 时不记录日志
    :param models: list - 模型列表
    :param trainable_models: dict - 可训练模型字典
    :param group_size: int - 组大小
    :param queue_size: int - 队列大小
    :param thread_size: int - 线程大小
    :param data_size: int - 数据大小
    :param sensitive_rate: float - 敏感率
    :param data: list - 数据列表
    :param feature_size: int - 特征大小
    :param seq_size: int - 序列大小
    :param buffer_size: int - 缓冲区大小
    :param batch_size: int - 批次大小
    :param num_workers: int - 工作进程数
    :param is_multiprocess: bool - 是否使用多进程
    :param device: torch.device - 模型运行的设备
    :param epoch: int - 轮次
    :param dev_type: str - 评测类型

    :return: tuple - 包含两个字典，分别为准确率和损失，字典的键为模型代码，值为对应的平均准确率和平均损失
    """
    i = 0
    j = 0
    group_num = int(data_size / group_size)
    group_remainder = data_size % group_size
    if group_remainder != 0:
        group_num += 1
    mark_eval(models)
    accuracy = {}
    loss = {}
    logger = logging.getLogger()
    if log_mode:
        log_name = 'log'
        if dev_type is not None:
            log_name = log_name + '_' + dev_type
        handlers = logger.handlers
        for handler in handlers:
            logger.removeHandler(handler)
            del handler
        del handlers
        handler = FileHandler(log_name + '.txt', mode='a', encoding=constants.DEFAULT_CHARSET)
        handler.setLevel(constants.LOG_LEVEL)
        handler.setFormatter(Formatter(constants.LOG_FORMAT))
        logger.addHandler(handler)
        del handler
    while i < data_size:
        collect_future_list = []
        pending_data_queue = Queue(queue_size)
        with ThreadPoolExecutor(max_workers=thread_size) as executor:
            j += 1
            curr_size = group_size if j < group_num or group_remainder == 0 else group_remainder
            logger.info('[{} Epoch {}, Task Num {}, Group Num {}/{}]'.format(dev_type, epoch, curr_size, j, group_num))
            collect_future_list.append(
                executor.submit(consume_data_2_dev_model, out_mode, model_mode, dev_type, epoch, i, data_size,
                                trainable_models, curr_size, pending_data_queue, accuracy, loss, collect_future_list,
                                executor, logger))
            while i < data_size:
                data_path = data[i]
                code = 'all' if out_mode else data_path[-6:]
                executor.submit(product_data, problem_type, model_mode, code, trainable_models.get(code)[0],
                                sensitive_rate, data_path, feature_size, seq_size, buffer_size, batch_size, num_workers,
                                is_multiprocess, device, pending_data_queue)
                i += 1
                if i % group_size == 0:
                    break
            collect_futures = futures.as_completed(collect_future_list)
            for collect_future in collect_futures:
                collect_result = collect_future.result()
                del collect_result
                del collect_future
            del collect_futures
            del collect_future_list
            del pending_data_queue
    del logger
    return dict_value_mean(accuracy), dict_value_mean(loss)


def dict_value_mean(material):
    """
    计算字典中每个键对应的值的平均值，并返回一个新的字典，其中键不变，值为对应的平均值。

    :param material: dict - 输入字典

    :return: dict - 包含平均值的新字典
    """
    result = {}
    material_list = material.items()
    for key, value in material_list:
        mean_value = numpy.mean(value)
        del value
        result[key] = mean_value
        del mean_value
        del key
    del material_list
    return result


def train(run_mode, problem_type, out_mode, model_mode, log_mode, epoch_num, group_size, queue_size, thread_size,
          feature_size, seq_size, buffer_size, batch_size, num_workers, is_multiprocess, device, models,
          trainable_models, sensitive_rate, train_data, model_path, train_mode, result_data, dev_data=None,
          dev_type=None):
    """
    训练模型。

    :param run_mode: int - 运行模式，0 表示正常模式，1 表示内存分析模式
    :param problem_type: bool - 为分类任务时为 True，为回归任务时为 False
    :param out_mode: bool - 为整体模型输出时为 True，为单独模型输出时为 False
    :param model_mode: int - 为 FMLSTMAttentionModel 时为 0，为 FMLlamaModel 时为 1
    :param log_mode: bool - bool - 为 True 时记录日志，为 False 时不记录日志
    :param epoch_num: int - 轮次数量
    :param group_size: int - 组大小
    :param queue_size: int - 队列大小
    :param thread_size: int - 线程数量
    :param feature_size: int - 特征大小
    :param seq_size: int - 序列大小
    :param buffer_size: int - 缓冲区大小
    :param batch_size: int - 批次大小
    :param num_workers: int - 工作进程数量
    :param is_multiprocess: bool - 是否使用多进程
    :param device: torch.device - 模型运行的设备
    :param models: list - 模型列表
    :param trainable_models: dict - 可训练模型字典
    :param sensitive_rate: float - 敏感率
    :param train_data: list - 训练数据
    :param model_path: str - 模型路径
    :param train_mode: bool - 为接着原来的模型继续训练时为 True，为重新训练模型时为 False
    :param result_data: str - 结果数据路径
    :param dev_data: list - 验证数据
    :param dev_type: str - 验证数据类型

    :return: tuple - 包含两个字典，分别为准确率和损失，字典的键为模型代码，值为对应的平均准确率和平均损失
    """
    epoch = 0
    train_data_size = len(train_data)
    group_num = int(train_data_size / group_size)
    group_remainder = train_data_size % group_size
    if group_remainder != 0:
        group_num += 1
    dev_data_size = len(dev_data)
    grad_norm_max = settings.grad_norm_max
    dev_eval_best = None
    eval_name = 'eval'
    log_name = 'log' if log_mode else None
    if dev_type is not None:
        eval_name = eval_name + '.' + dev_type
        if log_mode:
            log_name = log_name + '_' + dev_type
    eval_data = os.path.join(result_data, eval_name)
    if train_mode and os.path.exists(eval_data):
        with open(eval_data, 'r', encoding=constants.DEFAULT_CHARSET) as file:
            dev_eval_best = json.load(file)
        epoch = dev_eval_best['epoch']
        if epoch is None:
            epoch = 0
    else:
        dev_eval_best = {}
    logger = logging.getLogger()
    if log_mode:
        handlers = logger.handlers
        for handler in handlers:
            logger.removeHandler(handler)
            del handler
        del handlers
        handler = FileHandler(log_name + '.txt', mode='a', encoding=constants.DEFAULT_CHARSET)
        handler.setLevel(constants.LOG_LEVEL)
        handler.setFormatter(Formatter(constants.LOG_FORMAT))
        logger.addHandler(handler)
        del handler
    while epoch < epoch_num:
        i = 0
        j = 0
        mark_train(models)
        epoch_accuracy = {}
        epoch_loss = {}
        while i < train_data_size:
            collect_future_list = []
            pending_data_queue = Queue(queue_size)
            with ThreadPoolExecutor(max_workers=thread_size) as executor:
                j += 1
                curr_size = group_size if j < group_num or group_remainder == 0 else group_remainder
                logger.info('[Train Epoch {}, Task Num {}, Group Num {}/{}]'.format(epoch, curr_size, j, group_num))
                collect_future_list.append(
                    executor.submit(consume_data_2_train_model, out_mode, model_mode, grad_norm_max, epoch, i,
                                    train_data_size, trainable_models, curr_size, pending_data_queue, epoch_accuracy,
                                    epoch_loss, collect_future_list, executor, logger))
                while i < train_data_size:
                    data_path = train_data[i]
                    code = 'all' if out_mode else data_path[-6:]
                    executor.submit(product_data, problem_type, model_mode, code, trainable_models.get(code)[0],
                                    sensitive_rate, data_path, feature_size, seq_size, buffer_size, batch_size,
                                    num_workers, is_multiprocess, device, pending_data_queue)
                    del code
                    del data_path
                    i += 1
                    if i % group_size == 0:
                        break
                collect_futures = futures.as_completed(collect_future_list)
                for collect_future in collect_futures:
                    collect_result = collect_future.result()
                    del collect_result
                    del collect_future
                del collect_futures
                del collect_future_list
                del pending_data_queue

        dev_accuracy = None
        dev_loss = None
        if dev_data is not None:
            dev_accuracy, dev_loss = dev(problem_type, out_mode, model_mode, log_mode, models, trainable_models,
                                         group_size, queue_size, thread_size, dev_data_size, sensitive_rate, dev_data,
                                         feature_size, seq_size, buffer_size, batch_size, num_workers, is_multiprocess,
                                         device, epoch, dev_type)
        code_accuracies = epoch_accuracy.items()
        for code, accuracy in code_accuracies:
            code_model_state_dict = trainable_models.get(code)[1].state_dict()
            torch.save(code_model_state_dict, os.path.join(model_path, code + '.snapshoot'))
            train_accuracy = numpy.mean(accuracy)
            train_loss = numpy.mean(epoch_loss.get(code))
            msg = '{} [Train Epoch {}/{}] train_accuracy: {:.10f} train_loss: {:.10f}'.format(code, epoch, epoch_num,
                                                                                              train_accuracy,
                                                                                              train_loss)
            del accuracy
            code_eval_best = dev_eval_best.get(code)
            if dev_data is not None:
                code_loss = dev_loss.get(code)
                code_accuracy = dev_accuracy.get(code)
                msg += ' dev_accuracy: {:.10f} dev_loss: {:.10f}'.format(code_accuracy, code_loss)
                if code_eval_best is None:
                    dev_eval_best[code] = [train_accuracy, train_loss, code_accuracy, code_loss]
                    torch.save(code_model_state_dict, os.path.join(model_path, code))
                    msg += ' save_best_model: *'
                elif code_loss < code_eval_best[3]:
                    code_eval_best[0] = train_accuracy
                    code_eval_best[1] = train_loss
                    code_eval_best[2] = code_accuracy
                    code_eval_best[3] = code_loss
                    torch.save(code_model_state_dict, os.path.join(model_path, code))
                    msg += ' save_best_model: *'
                del code_accuracy
                del code_loss
            else:
                if code_eval_best is None:
                    dev_eval_best[code] = [train_accuracy, train_loss]
                    torch.save(code_model_state_dict, os.path.join(model_path, code))
                    msg += ' save_best_model: *'
                elif train_loss < code_eval_best[1]:
                    code_eval_best[0] = train_accuracy
                    code_eval_best[1] = train_loss
                    torch.save(code_model_state_dict, os.path.join(model_path, code))
                    msg += ' save_best_model: *'
            del code_eval_best
            del code_model_state_dict
            del train_accuracy
            del train_loss
            del code
            logger.info(msg)
            del msg
        del code_accuracies
        del epoch_accuracy
        del epoch_loss

        if run_mode == 1:
            k = 0
            top_stats = tracemalloc.take_snapshot().statistics("lineno")
            stats_size = len(top_stats)
            top50_stats = top_stats[:50]
            top_stats_size = len(top50_stats)
            while k < top_stats_size:
                logger.info('[Train Epoch {}, top {}/{}]: {}'.format(epoch, k + 1, stats_size, top50_stats[k]))
                k += 1

        epoch += 1

        dev_eval_best['epoch'] = epoch if epoch < epoch_num else 0
        with open(eval_data, 'w', encoding=constants.DEFAULT_CHARSET) as file:
            file.write(json.dumps(dev_eval_best))

    del dev_eval_best
    del logger


def product_data(problem_type, model_mode, code, meta, sensitive_rate, data_path, feature_size, seq_size, buffer_size,
                 batch_size, num_workers, is_multiprocess, device, pending_data_queue):
    """
    生产用于训练或者评测模型的数据集。

    :param problem_type: bool - 为分类任务时为 True，为回归任务时为 False
    :param model_mode: int - 为 FMLSTMAttentionModel 时为 0，为 FMLlamaModel 时为 1
    :param code: str - 标识
    :param meta: dict - 包含数据特征范围的元数据字典
    :param sensitive_rate: float - 敏感率
    :param data_path: str - 数据文件路径
    :param feature_size: int - 输入特征的维度大小
    :param seq_size: int - 序列长度（时间步数）
    :param buffer_size: int - 数据加载时的缓冲区大小，通常用于控制数据预处理和加载的效率
    :param batch_size: int - 每个训练或评测批次的数据量大小
    :param num_workers: int - 工作进程数
    :param is_multiprocess: bool - 如果设置为 True ，则数据加载器在数据集被使用一次后不会关闭工作进程。这使得能够保持工作进程中的数据集实例处于活跃状态。（默认值：False）
    :param device: torch.device - 数据加载到的设备
    :param pending_data_queue: queue.Queue - 用于存放生产的数据集的队列，消费者线程或进程将从该队列中获取数据集进行训练或评测

    :return: bool - 操作是否成功
    """
    pending_data_queue.put(
        build_dataset(problem_type, model_mode, code, meta, sensitive_rate, data_path, feature_size, seq_size,
                      buffer_size, batch_size, num_workers, is_multiprocess, device, True))
    return True


def consume_data_2_dev_model(out_mode, model_mode, dev_type, epoch, j, data_size, trainable_models, curr_size,
                             pending_data_queue, accuracy, loss, collect_future_list, executor, logger):
    """
    获取数据后评测模型。
    
    :param out_mode: bool - 为整体模型输出时为 True，为单独模型输出时为 False
    :param model_mode: int - 为 FMLSTMAttentionModel 时为 0，为 FMLlamaModel 时为 1
    :param dev_type: str - 评测类型
    :param epoch: int - 当前训练轮次
    :param j: int - 当前处理的文件索引
    :param data_size: int - 数据集大小
    :param trainable_models: dict - 可训练模型字典
    :param curr_size: int - 当前处理的数据量
    :param pending_data_queue: queue.Queue - 用于存放生产的数据集的队列，消费者线程或进程将从该队列中获取数据集进行训练或评测
    :param accuracy: dict - 存储每个模型的准确率
    :param loss: dict - 存储每个模型的损失值
    :param collect_future_list: list - 用于收集异步任务结果的列表
    :param executor: concurrent.futures.thread.ThreadPoolExecutor - 线程池执行器，用于管理和执行异步任务
    :param logger: logging.Logger - 日志记录器

    :return: bool - 操作是否成功
    """
    i = 0
    while i < curr_size:
        code, data_loader = pending_data_queue.get()
        i += 1
        j += 1
        logger.info('[{} Epoch {}, File {}/{}]'.format(dev_type, epoch, j, data_size))
        code_accuracy = accuracy.get(code)
        if code_accuracy is None:
            code_accuracy = []
            accuracy[code] = code_accuracy
        code_loss = loss.get(code)
        if code_loss is None:
            code_loss = []
            loss[code] = code_loss
        if out_mode:
            if model_mode == 2:
                for feat_value, label in data_loader:
                    dev_model(trainable_models.get(code)[1], None, feat_value, label, code_accuracy, code_loss, None)
                    del feat_value
                    del label
            else:
                for feat_index, feat_value, label in data_loader:
                    dev_model(trainable_models.get(code)[1], feat_index, feat_value, label, code_accuracy, code_loss,
                              None)
                    del feat_index
                    del feat_value
                    del label
        else:
            lock = Lock()
            if model_mode == 2:
                for feat_value, label in data_loader:
                    collect_future_list.append(
                        executor.submit(dev_model, trainable_models.get(code)[1], None, feat_value, label,
                                        code_accuracy, code_loss, lock))
                    del feat_value
                    del label
            else:
                for feat_index, feat_value, label in data_loader:
                    collect_future_list.append(
                        executor.submit(dev_model, trainable_models.get(code)[1], feat_index, feat_value, label,
                                        code_accuracy, code_loss, lock))
                    del feat_index
                    del feat_value
                    del label
        del data_loader
        del code
    return True


def consume_data_2_train_model(out_mode, model_mode, grad_norm_max, epoch, j, train_data_size, trainable_models,
                               curr_size, pending_data_queue, epoch_accuracy, epoch_loss, collect_future_list, executor,
                               logger):
    """
    获取数据后训练模型。

    :param out_mode: bool - 为整体模型输出时为 True，为单独模型输出时为 False
    :param model_mode: int - 为 FMLSTMAttentionModel 时为 0，为 FMLlamaModel 时为 1
    :param grad_norm_max: float - 梯度范数的最大值
    :param epoch: int - 当前训练轮次
    :param j: int - 当前处理的文件索引
    :param train_data_size: int - 训练数据集大小
    :param trainable_models: dict - 可训练模型字典
    :param curr_size: int - 当前处理的数据量
    :param pending_data_queue: queue.Queue - 用于存放生产的数据集的队列，消费者线程或进程将从该队列中获取数据集进行训练或评测
    :param epoch_accuracy: dict - 存储每个模型在当前轮次的准确率
    :param epoch_loss: dict - 存储每个模型在当前轮次的损失值
    :param collect_future_list: list - 用于收集异步任务结果的列表
    :param executor: concurrent.futures.thread.ThreadPoolExecutor - 线程池执行器，用于管理和执行异步任务
    :param logger: logging.Logger - 日志记录器

    :return: bool - 操作是否成功
    """
    i = 0
    while i < curr_size:
        code, data_loader = pending_data_queue.get()
        i += 1
        j += 1
        logger.info('[Train Epoch {}, File {}/{}]'.format(epoch, j, train_data_size))
        code_accuracy = epoch_accuracy.get(code)
        if code_accuracy is None:
            code_accuracy = []
            epoch_accuracy[code] = code_accuracy
        code_loss = epoch_loss.get(code)
        if code_loss is None:
            code_loss = []
            epoch_loss[code] = code_loss
        _, model, optimizer = trainable_models.get(code)
        if out_mode:
            if model_mode == 2:
                for feat_value, label in data_loader:
                    train_model(grad_norm_max, model, optimizer, None, feat_value, label, code_accuracy, code_loss,
                                None, logger)
                    del feat_value
                    del label
            else:
                for feat_index, feat_value, label in data_loader:
                    train_model(grad_norm_max, model, optimizer, feat_index, feat_value, label, code_accuracy,
                                code_loss, None, logger)
                    del feat_index
                    del feat_value
                    del label
        else:
            lock = Lock()
            if model_mode == 2:
                for feat_value, label in data_loader:
                    collect_future_list.append(
                        executor.submit(train_model, grad_norm_max, model, optimizer, None, feat_value, label,
                                        code_accuracy, code_loss, lock, logger))
                    del feat_value
                    del label
            else:
                for feat_index, feat_value, label in data_loader:
                    collect_future_list.append(
                        executor.submit(train_model, grad_norm_max, model, optimizer, feat_index, feat_value, label,
                                        code_accuracy, code_loss, lock, logger))
                    del feat_index
                    del feat_value
                    del label
        del data_loader
        del code
    return True


def dev_model(model, feat_index, feat_value, label, code_accuracy, code_loss, lock):
    """
    评测模型。

    :param model: 待评测的模型
    :param feat_index: 一个形状为 (batch_size, seq_size, field_size) 的 LongTensor，其中包含对嵌入表中元素的索引
    :param feat_value: 具有相同形状的 Tensor，其中包含特征值 (floats)。通常为 0/1 值或特征值的缩放形式
    :param label: 真实标签，对于分类任务通常为一个形状为 (batch_size,) 的 LongTensor，对于回归任务通常为一个形状为 (batch_size,) 的 FloatTensor
    :param code_accuracy: list - 存放当前标识下准确率的列表
    :param code_loss: list - 存放当前标识下损失值的列表
    :param lock: 待评测模型的锁

    :return: bool - 操作是否成功
    """
    if lock is None:
        with torch.no_grad():
            local_accuracy = None
            local_loss = None
            if feat_index is None:
                local_accuracy, local_loss = model.predict_and_evaluate(feat_value, label)
            else:
                local_accuracy, local_loss = model.predict_and_evaluate(feat_index, feat_value, label)
                del feat_index
            code_accuracy.append(local_accuracy.item())
            del local_accuracy
            code_loss.append(local_loss.item())
            del local_loss
    else:
        with torch.no_grad():
            local_accuracy = None
            local_loss = None
            with lock:
                if feat_index is None:
                    local_accuracy, local_loss = model.predict_and_evaluate(feat_value, label)
                else:
                    local_accuracy, local_loss = model.predict_and_evaluate(feat_index, feat_value, label)
                    del feat_index
                code_accuracy.append(local_accuracy.item())
                del local_accuracy
                code_loss.append(local_loss.item())
                del local_loss
    del feat_value
    del label
    return True


def train_model(grad_norm_max, model, optimizer, feat_index, feat_value, label, code_accuracy, code_loss, lock, logger):
    """
    训练模型。

    :param grad_norm_max: float - 梯度范数的最大值
    :param model: 待训练的模型
    :param optimizer: 优化器
    :param feat_index: 一个形状为 (batch_size, seq_size, field_size) 的 LongTensor，其中包含对嵌入表中元素的索引
    :param feat_value: 具有相同形状的 Tensor，其中包含特征值 (floats)。通常为 0/1 值或特征值的缩放形式
    :param label: 真实标签，对于分类任务通常为一个形状为 (batch_size,) 的 LongTensor，对于回归任务通常为一个形状为 (batch_size,) 的 FloatTensor
    :param code_accuracy: list - 存放当前标识下准确率的列表
    :param code_loss: list - 存放当前标识下损失值的列表
    :param lock: 训练模型的锁
    :param logger: 日志记录器

    :return: bool - 操作是否成功
    """
    if lock is None:
        model.zero_grad()
        local_accuracy = None
        local_loss = None
        if feat_index is None:
            local_accuracy, local_loss = model.predict_and_evaluate(feat_value, label)
        else:
            local_accuracy, local_loss = model.predict_and_evaluate(feat_index, feat_value, label)
            del feat_index
        if torch.isnan(local_loss).any():
            logger.error('NaN detected in local_loss!')
        local_loss.backward()
        utils.clip_grad_norm_(model.parameters(), grad_norm_max)
        code_accuracy.append(local_accuracy.item())
        del local_accuracy
        code_loss.append(local_loss.item())
        del local_loss
        optimizer.step()
    else:
        with lock:
            model.zero_grad()
            local_accuracy = None
            local_loss = None
            if feat_index is None:
                local_accuracy, local_loss = model.predict_and_evaluate(feat_value, label)
            else:
                local_accuracy, local_loss = model.predict_and_evaluate(feat_index, feat_value, label)
                del feat_index
            if torch.isnan(local_loss).any():
                logger.error('NaN detected in local_loss!')
            local_loss.backward()
            utils.clip_grad_norm_(model.parameters(), grad_norm_max)
            code_accuracy.append(local_accuracy.item())
            del local_accuracy
            code_loss.append(local_loss.item())
            del local_loss
            optimizer.step()
    del feat_value
    del label
    return True


def save_models(model_list, model_path):
    """
    保存模型字典里的全部模型。

    :param model_list: dict - 待保存的全部模型
    :param model_path: 模型路径
    """
    for code, trainable_model in model_list:
        torch.save(trainable_model[1].state_dict(), os.path.join(model_path, code))


def mark_train(models):
    """
    标记全部的模型为训练模式。

    :param models: 待标记的全部模型
    """
    for model in models:
        model.train()


def mark_eval(models):
    """
    标记全部的模型为评测模式。

    :param models: 待标记的全部模型
    """
    for model in models:
        model.eval()


def prepare_models(out_mode, train_mode, no_decay_flags, learn_rate, weight_decay, betas, eps, model_mode, problem_type,
                   feature_size, field_size, seq_size, device, settings, meta_path, model_path):
    """
    准备全部的模型。

    :param out_mode: bool - 为整体模型输出时为 True，为单独模型输出时为 False
    :param train_mode: bool - 为接着原来的模型继续训练时为 True，为重新训练模型时为 False
    :param no_decay_flags: list - 不需要权重衰减的参数名称标识列表
    :param learn_rate: float - 学习率
    :param weight_decay: float - 权重衰减系数
    :param betas: tuple[float, float] - AdamW 优化器的 beta 参数，通常为 (0.9, 0.999)
    :param eps: float - AdamW 优化器的 eps 参数
    :param model_mode: int - 为 FMLSTMAttentionModel 时为 0，为 FMLlamaModel 时为 1
    :param problem_type: bool - 为分类任务时为 True，为回归任务时为 False
    :param feature_size: int - 输入特征的维度大小
    :param field_size: int - 每个时间步长的字段数量
    :param seq_size: int - 序列长度（时间步数）
    :param device: torch.device - 模型运行的设备
    :param settings: 配置对象
    :param meta_path: str - 元数据路径
    :param model_path: str - 模型路径

    :return: tuple[list, dict] - 模型列表和可训练模型字典
    """
    models = []
    trainable_models = {}
    if out_mode:
        model = prepare_model(train_mode, no_decay_flags, learn_rate, weight_decay, betas, eps, model_mode,
                              problem_type, feature_size, field_size, seq_size, device, settings,
                              os.path.join(meta_path, 'all'), os.path.join(model_path, 'all'))
        models.append(model[1])
        trainable_models['all'] = model
    else:
        codes = os.listdir(meta_path)
        for code in codes:
            model = prepare_model(train_mode, no_decay_flags, learn_rate, weight_decay, betas, eps, model_mode,
                                  problem_type, feature_size, field_size, seq_size, device, settings,
                                  os.path.join(meta_path, code), os.path.join(model_path, code))
            models.append(model[1])
            trainable_models[code] = model
    return models, trainable_models


def prepare_model(train_mode, no_decay_flags, learn_rate, weight_decay, betas, eps, model_mode, problem_type,
                  feature_size, field_size, seq_size, device, settings, meta_path, model_path):
    """
    准备单个模型。

    :param train_mode: bool - 为接着原来的模型继续训练时为 True，为重新训练模型时为 False
    :param no_decay_flags: list - 不需要权重衰减的参数名称标识列表
    :param learn_rate: float - 学习率
    :param weight_decay: float - 权重衰减系数
    :param betas: tuple[float, float] - AdamW 优化器的 beta 参数，通常为 (0.9, 0.999)
    :param eps: float - AdamW 优化器的 eps 参数
    :param model_mode: int - 为 FMLSTMAttentionModel 时为 0，为 FMLlamaModel 时为 1
    :param problem_type: bool - 为分类任务时为 True，为回归任务时为 False
    :param feature_size: int - 输入特征的维度大小
    :param field_size: int - 每个时间步长的字段数量
    :param seq_size: int - 序列长度（时间步数）
    :param device: torch.device - 模型运行的设备
    :param settings: 配置对象
    :param meta_path: str - 元数据路径
    :param model_path: str - 模型路径

    :return: tuple - 包含三个元素的元组，第一个元素是模型的元数据，第二个元素是准备好的模型对象，第三个元素是优化器对象
    """
    no_decay_params = []
    other_params = []
    model = FMLlamaModel(problem_type, feature_size, field_size, seq_size, settings).to(
        device) if model_mode == 1 else FMLSTMAttentionModel(problem_type, feature_size, field_size, seq_size,
                                                             settings).to(device)
    if train_mode and os.path.exists(model_path):
        snapshoot = model_path + '.snapshoot'
        if os.path.exists(snapshoot):
            model.load_state_dict(torch.load(snapshoot))
        else:
            model.load_state_dict(torch.load(model_path))
    named_parameters = model.named_parameters()
    for name, parameter in named_parameters:
        if any(no_decay_flag in name for no_decay_flag in no_decay_flags):
            no_decay_params.append(parameter)
        else:
            other_params.append(parameter)
    optimizer = AdamW(
        [
            {
                'params': other_params,
                'weight_decay': weight_decay
            },
            {
                'params': no_decay_params,
                'weight_decay': 0.0
            }
        ],
        lr=learn_rate,
        betas=betas,
        eps=eps
    )
    return get_data_meta(meta_path)['all'], model, optimizer


def copy_models(out_mode, meta_path, src_model_path, model_path):
    """
    复制全部的模型。

    :param out_mode: bool - 为整体模型输出时为 True，为单独模型输出时为 False
    :param meta_path: str - 元数据路径
    :param src_model_path: str - 源模型路径
    :param model_path: str - 目标模型路径
    """
    if out_mode:
        copy_model('all', src_model_path, model_path)
    else:
        codes = os.listdir(meta_path)
        for code in codes:
            copy_model(code, src_model_path, model_path)


def copy_model(model_name, src_model_path, model_path):
    """
    复制单个模型。

    :param model_name: str - 模型名称
    :param src_model_path: str - 源模型路径
    :param model_path: str - 目标模型路径
    """
    dst_path = os.path.join(model_path, model_name)
    if not os.path.exists(dst_path):
        shutil.copyfile(os.path.join(src_model_path, model_name), dst_path)


if __name__ == '__main__':
    settings = config.parsers()
    num_workers = settings.num_workers
    is_multiprocess = num_workers >= 1
    if is_multiprocess:
        multiprocessing.freeze_support()
    multiprocessing.set_start_method('spawn', force=True)
    if is_multiprocess:
        multiprocessing.set_sharing_strategy('file_system')

    run_mode = settings.run_mode
    if run_mode == 1:
        tracemalloc.start()

    random.seed(settings.random_seed)
    train_data = settings.final_train_data
    train_data_files = None
    dev_data = settings.final_dev_data
    dev_data_files = None
    test_data = settings.final_test_data
    test_data_files = None

    dev_type = None
    problem_type = settings.problem_type
    model_mode = settings.model_mode
    sensitive_rate = settings.sensitive_rate
    group_size = settings.group_size
    queue_size = settings.queue_size
    thread_size = settings.thread_size
    out_mode = settings.out_mode
    feature_size = settings.feature_size
    field_size = settings.field_size
    meta_data = settings.meta_data
    train_stage = settings.train_stage
    if train_stage == 0:
        dev_type = 'Dev'
        meta_data = meta_data + '_' + dev_type
    elif train_stage == 1:
        dev_type = 'Test'
        meta_data = meta_data + '_' + dev_type
    if out_mode:
        feature_size += 2
        field_size += 1
        train_data_files = os.listdir(train_data)
        random.shuffle(train_data_files)
        train_data_files = [os.path.join(train_data, file_name) for file_name in train_data_files]
        dev_data_files = os.listdir(dev_data)
        random.shuffle(dev_data_files)
        dev_data_files = [os.path.join(dev_data, file_name) for file_name in dev_data_files]
        test_data_files = os.listdir(test_data)
        random.shuffle(test_data_files)
        test_data_files = [os.path.join(test_data, file_name) for file_name in test_data_files]
    else:
        train_data_files = []
        dev_data_files = []
        test_data_files = []
        codes = os.listdir(meta_data)
        for code in codes:
            train_data_files.append(os.path.join(train_data, code))
            dev_data_files.append(os.path.join(dev_data, code))
            test_data_files.append(os.path.join(test_data, code))
    test_data_size = len(test_data_files)
    seq_size = settings.seq_size
    buffer_size = settings.buffer_size
    epoch_num = settings.epoch_num
    batch_size = settings.batch_size
    learn_rate = settings.learn_rate
    weight_decay = settings.weight_decay
    betas = (settings.beta_1, settings.beta_2)
    eps = settings.eps
    device = settings.device
    train_mode = settings.train_mode
    log_mode = settings.log_mode
    result_data = settings.result_data
    no_decay_flags = ['bias', 'norm.weight']

    if train_stage == 0:
        dev_model_best = settings.dev_model_best
        models, trainable_models = prepare_models(out_mode, train_mode, no_decay_flags, learn_rate, weight_decay, betas,
                                                  eps, model_mode, problem_type, feature_size, field_size, seq_size,
                                                  device, settings, meta_data, dev_model_best)
        try:
            train(run_mode, problem_type, out_mode, model_mode, log_mode, epoch_num, group_size, queue_size,
                  thread_size, feature_size, seq_size, buffer_size, batch_size, num_workers, is_multiprocess, device,
                  models, trainable_models, sensitive_rate, train_data_files, dev_model_best, train_mode, result_data,
                  dev_data_files, dev_type)
            test_accuracy, test_loss = dev(problem_type, out_mode, model_mode, log_mode, models, trainable_models,
                                           group_size, queue_size, thread_size, test_data_size, sensitive_rate,
                                           test_data_files, feature_size, seq_size, buffer_size, batch_size,
                                           num_workers, is_multiprocess, device, epoch_num - 1, dev_type)
            logger = logging.getLogger()
            code_accuracies = test_accuracy.items()
            for code, accuracy in code_accuracies:
                logger.info('{} last epoch: test_accuracy:{:.10f}, test_loss:{:.10f}'.format(code, accuracy,
                                                                                             test_loss.get(code)))
        except BaseException as e:
            traceback.print_exc()
    else:
        test_model_best = settings.test_model_best
        if train_stage == 1:
            copy_models(out_mode, meta_data, settings.dev_model_best, test_model_best)
        models, trainable_models = prepare_models(out_mode, True, no_decay_flags, learn_rate, weight_decay, betas, eps,
                                                  model_mode, problem_type, feature_size, field_size, seq_size, device,
                                                  settings, meta_data, test_model_best)
        try:
            if train_stage == 1:
                train(run_mode, problem_type, out_mode, model_mode, log_mode, epoch_num, group_size, queue_size,
                      thread_size, feature_size, seq_size, buffer_size, batch_size, num_workers, is_multiprocess,
                      device, models, trainable_models, sensitive_rate, dev_data_files, test_model_best, train_mode,
                      result_data, test_data_files, dev_type)
            else:
                train(run_mode, problem_type, out_mode, model_mode, log_mode, epoch_num, group_size, queue_size,
                      thread_size, feature_size, seq_size, buffer_size, batch_size, num_workers, is_multiprocess,
                      device, models, trainable_models, sensitive_rate, test_data_files, settings.online_model,
                      train_mode, result_data)
        except BaseException as e:
            traceback.print_exc()