# -*- coding: utf-8 -*-

'''

@time: 2019/7/23 19:42

@ author: javis

'''

import torch, time, os, shutil

import models, utils

import numpy as np

import pandas as pd

from tensorboard_logger import Logger

from torch import nn, optim

from torch.utils.data import DataLoader

from dataset import ECGDataset

from config import config

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

torch.manual_seed(41)

torch.cuda.manual_seed(41)

# 保存当前模型的权重，并且更新最佳的模型权重

def save_ckpt(state, is_best, model_save_dir):

    current_w = os.path.join(model_save_dir, config.current_w)

    best_w = os.path.join(model_save_dir, config.best_w)

    torch.save(state, current_w)

    if is_best: shutil.copyfile(current_w, best_w)

def train_epoch(model, optimizer, criterion, train_dataloader, show_interval=10):

    model.train()

    f1_meter, loss_meter, it_count = 0, 0, 0

    for inputs, target in train_dataloader:

        inputs = inputs.to(device)

        target = target.to(device)

        # zero the parameter gradients

        optimizer.zero_grad()

        # forward

        output = model(inputs)

        loss = criterion(output, target)

        loss.backward()

        optimizer.step()

        loss_meter += loss.item()

        it_count += 1

        f1 = utils.calc_f1(target, torch.sigmoid(output))

        f1_meter += f1

        if it_count != 0 and it_count % show_interval == 0:

            print("%d,loss:%.3e f1:%.3f" % (it_count, loss.item(), f1))

    return loss_meter / it_count, f1_meter / it_count

def val_epoch(model, criterion, val_dataloader, threshold=0.5):

    model.eval()

    f1_meter, loss_meter, it_count = 0, 0, 0

    with torch.no_grad():

        for inputs, target in val_dataloader:

            inputs = inputs.to(device)

            target = target.to(device)

            output = model(inputs)

            loss = criterion(output, target)

            loss_meter += loss.item()

            it_count += 1

            output = torch.sigmoid(output)

            f1 = utils.calc_f1(target, output, threshold)

            f1_meter += f1

    return loss_meter / it_count, f1_meter / it_count

def train(args):

    # model

    model = getattr(models, config.model_name)()

    if args.ckpt and not args.resume:

        state = torch.load(args.ckpt, map_location='cpu')

        model.load_state_dict(state['state_dict'])

        print('train with pretrained weight val_f1', state['f1'])

    model = model.to(device)

    # data

    train_dataset = ECGDataset(data_path=config.train_data, train=True)

    train_dataloader = DataLoader(train_dataset, batch_size=config.batch_size, shuffle=True, num_workers=6)

    val_dataset = ECGDataset(data_path=config.train_data, train=False)

    val_dataloader = DataLoader(val_dataset, batch_size=config.batch_size, num_workers=4)

    print("train_datasize", len(train_dataset), "val_datasize", len(val_dataset))

    # optimizer and loss

    optimizer = optim.Adam(model.parameters(), lr=config.lr)

    w = torch.tensor(train_dataset.wc, dtype=torch.float).to(device)

    criterion = utils.WeightedMultilabel(w)

    # 模型保存文件夹

    model_save_dir = '%s/%s_%s' % (config.ckpt, config.model_name, time.strftime("%Y%m%d%H%M"))

    if args.ex: model_save_dir += args.ex

    best_f1 = -1

    lr = config.lr

    start_epoch = 1

    stage = 1

    # 从上一个断点，继续训练

    if args.resume:

        if os.path.exists(args.ckpt):  # 这里是存放权重的目录

            model_save_dir = args.ckpt

            current_w = torch.load(os.path.join(args.ckpt, config.current_w))

            best_w = torch.load(os.path.join(model_save_dir, config.best_w))

            best_f1 = best_w['loss']

            start_epoch = current_w['epoch'] + 1

            lr = current_w['lr']

            stage = current_w['stage']

            model.load_state_dict(current_w['state_dict'])

            # 如果中断点恰好为转换stage的点

            if start_epoch - 1 in config.stage_epoch:

                stage += 1

                lr /= config.lr_decay

                utils.adjust_learning_rate(optimizer, lr)

                model.load_state_dict(best_w['state_dict'])

            print("=> loaded checkpoint (epoch {})".format(start_epoch - 1))

    logger = Logger(logdir=model_save_dir, flush_secs=2)

    # =========>开始训练<=========

    for epoch in range(start_epoch, config.max_epoch + 1):

        since = time.time()

        train_loss, train_f1 = train_epoch(model, optimizer, criterion, train_dataloader, show_interval=100)

        val_loss, val_f1 = val_epoch(model, criterion, val_dataloader)

        print('#epoch:%02d stage:%d train_loss:%.3e train_f1:%.3f  val_loss:%0.3e val_f1:%.3f time:%s\n'

              % (epoch, stage, train_loss, train_f1, val_loss, val_f1, utils.print_time_cost(since)))

        logger.log_value('train_loss', train_loss, step=epoch)

        logger.log_value('train_f1', train_f1, step=epoch)

        logger.log_value('val_loss', val_loss, step=epoch)

        logger.log_value('val_f1', val_f1, step=epoch)

        state = {"state_dict": model.state_dict(), "epoch": epoch, "loss": val_loss, 'f1': val_f1, 'lr': lr,

                 'stage': stage}

        save_ckpt(state, best_f1 < val_f1, model_save_dir)

        best_f1 = max(best_f1, val_f1)

        if epoch in config.stage_epoch:

            stage += 1

            lr /= config.lr_decay

            best_w = os.path.join(model_save_dir, config.best_w)

            model.load_state_dict(torch.load(best_w)['state_dict'])

            print("*" * 10, "step into stage%02d lr %.3ef" % (stage, lr))

            utils.adjust_learning_rate(optimizer, lr)

#用于测试加载模型

def val(args):

    list_threhold = [0.5]

    model = getattr(models, config.model_name)()

    if args.ckpt: model.load_state_dict(torch.load(args.ckpt, map_location='cpu')['state_dict'])

    model = model.to(device)

    criterion = nn.BCEWithLogitsLoss()

    val_dataset = ECGDataset(data_path=config.train_data, train=False)

    val_dataloader = DataLoader(val_dataset, batch_size=config.batch_size, num_workers=4)

    for threshold in list_threhold:

        val_loss, val_f1 = val_epoch(model, criterion, val_dataloader, threshold)

        print('threshold %.2f val_loss:%0.3e val_f1:%.3f\n' % (threshold, val_loss, val_f1))

#提交结果使用

def test(args):

    from dataset import transform

    from data_process import name2index

    name2idx = name2index(config.arrythmia)

    idx2name = {idx: name for name, idx in name2idx.items()}

    utils.mkdirs(config.sub_dir)

    # model

    model = getattr(models, config.model_name)()

    model.load_state_dict(torch.load(args.ckpt, map_location='cpu')['state_dict'])

    model = model.to(device)

    model.eval()

    sub_file = '%s/subA_%s.txt' % (config.sub_dir, time.strftime("%Y%m%d%H%M"))

    fout = open(sub_file, 'w', encoding='utf-8')

    with torch.no_grad():

        for line in open(config.test_label, encoding='utf-8'):

            fout.write(line.strip('\n'))

            id = line.split('\t')[0]

            file_path = os.path.join(config.test_dir, id)

            df = pd.read_csv(file_path, sep=' ').values

            x = transform(df).unsqueeze(0).to(device)

            output = torch.sigmoid(model(x)).squeeze().cpu().numpy()

            ixs = [i for i, out in enumerate(output) if out > 0.5]

            for i in ixs:

                fout.write("\t" + idx2name[i])

            fout.write('\n')

    fout.close()

if __name__ == '__main__':

    import argparse

    parser = argparse.ArgumentParser()

    parser.add_argument("command", metavar="<command>", help="train or infer")

    parser.add_argument("--ckpt", type=str, help="the path of model weight file")

    parser.add_argument("--ex", type=str, help="experience name")

    parser.add_argument("--resume", action='store_true', default=False)

    args = parser.parse_args()

    if (args.command == "train"):

        train(args)

    if (args.command == "test"):

        test(args)

    if (args.command == "val"):

        val(args)