# -*- coding: utf-8 -*-
import os
import time
from torch import save, load
from warnings import warn
from torch.nn import Module, DataParallel
from torch.optim import Optimizer
from .config import Config


def get_model(config: Config, **kwargs):
    from . import models
    model = getattr(models, 'DoubleNet')(config, **kwargs)
    # move core to GPU
    if config.use_gpu:
        model = model.cuda(0)
    # parallel processing
    model = DataParallel(model, config.gpu_list)
    # load weights
    if os.path.exists(config.weight_load_path):
        try:
            state_dict = load(config.weight_load_path, map_location=config.map_location)
            model.load_state_dict(state_dict)
            print('Loaded weights from ' + config.weight_load_path)
        except RuntimeError as e:
            print('Failed to load weights file {} because:\n{}'.format(config.weight_load_path, e))
    return model


def make_checkpoint(config, epoch, start_time, loss_val, train_score, val_score, model, optimizer=None):
    # type:(Config,int,float,float,float,float,Module,Optimizer)->None
    """
    generate temporary training process data for resuming by resume_checkpoint()
    """
    save(model.state_dict(), config.temp_weight_path)
    if optimizer is not None and hasattr(optimizer, "state_dict"):
        save(optimizer.state_dict(), config.temp_optim_path)
    with open(config.train_record_file, 'a+') as f:
        elapsed_time = time.time() - start_time,
        record = config.__record_dict__.format(config.init_time, epoch, start_time, elapsed_time, loss_val, train_score,
                                               val_score)
        f.write(record + '\n')


def resume_checkpoint(config: Config, model: Module, optimizer: Optimizer = None) -> int:
    """
    resume training process data from config.logs which generated by make_checkpoint()
    :return number of last epoch
    """
    last_epoch = -1
    temp_weight_path = config.temp_weight_path
    temp_optim_path = config.temp_optim_path
    if os.path.exists(config.train_record_file):
        try:
            with open(config.train_record_file, 'r') as f:
                last = f.readlines()[-1]
                import json
                info = json.loads(last)
                last_epoch = int(info["epoch"])
                last_init = str(info["init"])
                if not os.path.exists(temp_weight_path):
                    temp_weight_path = temp_weight_path.replace(config.init_time, last_init)
                if not os.path.exists(temp_optim_path):
                    temp_optim_path = temp_optim_path.replace(config.init_time, last_init)
            print("Continue train from last epoch %d" % last_epoch)
        except:
            warn("Rename invalid train record file from {} to {}".format(config.train_record_file,
                                                                         config.train_record_file + '.badfile'))
            warn("Can't get last_epoch value, {} will be returned".format(last_epoch))
            os.rename(config.train_record_file, config.train_record_file + '.badfile')
    if os.path.exists(temp_weight_path):
        try:
            model.load_state_dict(load(temp_weight_path))
            print("Resumed weight checkpoint from {}".format(temp_weight_path))
        except:
            warn("Move invalid temp {} weights file from {} to {}".format(type(model), temp_weight_path,
                                                                          temp_weight_path + '.badfile'))
            os.rename(temp_weight_path, temp_weight_path + '.badfile')
    if optimizer is not None and os.path.exists(temp_optim_path):
        try:
            optimizer.load_state_dict(load(temp_optim_path))
            print("Resumed optimizer checkpoint from {}".format(temp_optim_path))
        except:
            warn("Move invalid temp {} weights file from {} to {}".format(type(optimizer), temp_optim_path,
                                                                          temp_optim_path + '.badfile'))
            os.rename(temp_optim_path, temp_optim_path + '.badfile')

    return last_epoch
