import json
import os
import time
import torch
import numpy as np
from typing import Optional
from torch.utils.data import ConcatDataset, DataLoader, Subset
from abfml.param.param import Param
from abfml.core.model.builder import build_model
from abfml.logger.loggers import Logger, log_data_info, log_logo
from abfml.data.read_data import ReadData
from abfml.train.trainer import train_loop, valid_loop
from abfml.core.optimizer.optim_init import optim_init
from abfml.core.optimizer.learn_rate import calc_lr_param
import torch.multiprocessing


def train_mlff(
        *,
        INPUT: str,
        init_model: Optional[str],
        restart: Optional[str],
        **kwargs):
    with open(INPUT, 'r') as file:
        input_json = json.load(file)
    param_class = Param(input_dict=input_json)

    logger = Logger(param_class.GlobalSet.logger_file).logger
    log_logo(logger=logger)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    num_threads = torch.get_num_threads()
    num_worker = int(num_threads / 2)
    # torch.set_num_threads(1)
    logger.info(f"num_threads : {num_threads}, device : {device}")
    torch.set_printoptions(precision=10)
    if param_class.DataSet.ratio is None:
        logger.info("+-------------------------------------- train data file ---------------------------------------+")
        train_dataclass = ReadData(filename=param_class.DataSet.train_file,
                                   cutoff=param_class.GlobalSet.cutoff,
                                   neighbor=param_class.GlobalSet.neighbor,
                                   file_format=param_class.DataSet.file_format)
        log_data_info(logger, train_dataclass)
        train_data = ConcatDataset(train_dataclass.create_dataset())

        logger.info("+-------------------------------------- valid data file ---------------------------------------+")
        valid_dataclass = ReadData(filename=param_class.DataSet.valid_file,
                                   cutoff=param_class.GlobalSet.cutoff,
                                   neighbor=param_class.GlobalSet.neighbor,
                                   file_format=param_class.DataSet.file_format)
        log_data_info(logger, valid_dataclass)
        valid_data = ConcatDataset(valid_dataclass.create_dataset())
    else:
        logger.info("+----------------------------------------- data file -------------------------------------------+")
        total_dataclass = ReadData(filename=param_class.DataSet.train_file,
                                   cutoff=param_class.GlobalSet.cutoff,
                                   neighbor=param_class.GlobalSet.neighbor,
                                   file_format=param_class.DataSet.file_format)
        log_data_info(logger, total_dataclass)
        total_data = ConcatDataset(total_dataclass.create_dataset())
        total_indices = np.arange(len(total_data))
        if param_class.DataSet.shuffle:
            logger.info(f"| Because of shuffle is set to true, ")
            logger.info(f"| The training and validation sets will be randomly selected from all the data.")
            np.random.shuffle(total_indices)
        else:
            logger.info(f"| Because of shuffle is set to false, preserve the order of the dataset, scaled")
        split_point = int(len(total_data) * param_class.DataSet.ratio)
        train_indices, valid_indices = total_indices[:split_point], total_indices[split_point:]
        train_data = Subset(total_data, train_indices)
        valid_data = Subset(total_data, valid_indices)
    train_data_load = DataLoader(train_data, batch_size=param_class.TrainSet.batch_size,
                                 shuffle=True, num_workers=num_worker)
    valid_data_load = DataLoader(valid_data, batch_size=param_class.TrainSet.batch_size,
                                 shuffle=False, num_workers=num_worker)

    model = build_model(param_class, train_data, init_model, restart, logger)

    if param_class.GlobalSet.dtype == "float32":
        model.to(torch.float32)
    else:
        model.to(torch.float64)
    model.to(device=device)

    if param_class.TrainSet.freeze_param is not None:
        logger.info(f"+----------------------------------- freeze or train param ------------------------------------+")
        namelist = [name for name, _ in model.named_parameters()]
        if not all(name in namelist for name in param_class.TrainSet.freeze_param[1:]):
            raise Exception("Error: Some parameters were not found in the model.")
        if param_class.TrainSet.freeze_param[0] == "freeze":
            logger.info(f"| The following parameters will be frozen.")
            for name, param in model.named_parameters():
                if name in param_class.TrainSet.freeze_param[1:]:
                    param.requires_grad = False
        elif param_class.TrainSet.freeze_param[0] == "train":
            logger.info(f"| The following parameters will be trained.")
            for name, param in model.named_parameters():
                if name not in param_class.TrainSet.freeze_param:
                    param.requires_grad = False
        else:
            raise Exception(f"Keyword {param_class.TrainSet.freeze_param[0]} error, flags only allowed freeze and train!")
        for name in param_class.TrainSet.freeze_param[1:]:
            logger.info(f"| {name:<60s}")
        logger.info(f"+----------------------------------------------------------------------------------------------+")

    iters_step: int = 0
    start_epoch: int = 1
    if restart is not None:
        check_point_class = torch.load(restart)
        optimizer = check_point_class["optimizer"]
        start_epoch = check_point_class["epoch"]
        iters_step = check_point_class["iters_step"]
    else:
        optimizer = optim_init(model=model, param=param_class.TrainSet)
        calc_lr_param(lr_param=param_class.LrSet,
                      epochs=param_class.TrainSet.epoch,
                      train_iters=int(len(train_data)/param_class.TrainSet.batch_size),
                      logger_name=param_class.GlobalSet.logger_file)

    logger.info(r"########################################--Start  Train--########################################")
    logger.info(f"+------------------------------------------Start Loop------------------------------------------+")
    time_loop_start = time.time()
    for epoch in range(start_epoch, param_class.TrainSet.epoch + 1):
        logger.info(f"|-- Train Epoch:{epoch:>3d}, batch size:{train_data_load.batch_size:>3d}, "
                    f"total iters: {len(train_data_load):>6d} --------------------------------------|")
        train_loss_dict = train_loop(data_load=train_data_load,
                                     model=model,
                                     optimizer=optimizer,
                                     iters_step=iters_step,
                                     config=param_class)

        iters_step = iters_step + len(train_data_load)

        logger.info(f"|-- Valid Epoch:{epoch:>3d}, batch size:{valid_data_load.batch_size:>3d}, "
                    f"total iters: {len(valid_data_load):>6d}                                       |")
        valid_loss_dict, _ = valid_loop(data_load=valid_data_load,
                                        model=model,
                                        logger_name=param_class.GlobalSet.logger_file,
                                        print_freq=param_class.TrainSet.print_freq,
                                        save_predict=False)

        # ************************************** Start writing loss file ************************************** #
        if not os.path.exists(param_class.LossSet.loss_file) or epoch == 1:
            header = '{:>8}'.format("Epoch")
            header_time = ""
            for train_loss_key in train_loss_dict.keys():
                if train_loss_key == "time":
                    header_time = '{:>16s}(s)'.format("Total_Time")
                else:
                    header += '{:>18s}{:>18s}'.format("T_RMSE_" + train_loss_dict[train_loss_key].name,
                                                      "V_RMSE_" + train_loss_dict[train_loss_key].name)
            header += header_time
            with open(param_class.LossSet.loss_file, 'w') as fp:
                fp.write(header)

        summaries = f"\n{epoch:>8d}"
        summaries_time = ""
        for train_loss_key in train_loss_dict.keys():
            if train_loss_key == "time":
                total_time = '{:6.2f}'.format(train_loss_dict["time"].sum + valid_loss_dict["time"].sum)
                summaries_time = '{:>16s}'.format(total_time)
            else:
                summaries += '{:>18s}'.format(format(train_loss_dict[train_loss_key].avg, '.5e'))
                summaries += '{:>18s}'.format(format(valid_loss_dict[train_loss_key].avg, '.5e'))
        summaries += summaries_time
        with open(param_class.LossSet.loss_file, 'a') as fp:
            fp.write(summaries)
        # **************************************  Stop writing loss file ************************************** #

        if epoch % param_class.TrainSet.save_check == 0:
            check_point_path = "checkpoint_" + str(epoch) + ".ckpt"
        elif epoch == param_class.TrainSet.epoch:
            check_point_path = "checkpoint_final.ckpt"
        else:
            check_point_path = "checkpoint_latest.ckpt"

        torch.save({
            'epoch': epoch,
            'iters_step': iters_step,
            'model': model,
            'optimizer': optimizer
        }, check_point_path)
    try:
        # move model to CPU before scripting
        model_cpu = model.to("cpu")
        jit_model = torch.jit.script(model_cpu)
        jit_model.save(param_class.TrainSet.model_file)
    except Exception as e:
        torch.save(model, param_class.TrainSet.model_file+'.pth')
        logger.info(f"| Information: Failed to save the model via torchscript, will save the model by torch.save! \n"
                    f"               The model can't be loaded in lammps, but it can be used in ASE")
        logger.warning(f"{e}")
    time_loop_end = time.time()
    logger.info(f"+------------------------------------------ End Loop ------------------------------------------+")
    logger.info(f"| Total training time:{time_loop_end - time_loop_start:8.2f}s, total training iters:{iters_step:>8d}")
    logger.info(r"######################################## --End  Train-- ########################################")

    return
