import argparse
import glob
import json
import os
import multiprocessing
from importlib import import_module

import sys
sys.path.append(os.path.dirname(os.path.abspath(os.path.dirname(__file__))))

import numpy as np
import torch
from torch.optim.lr_scheduler import StepLR
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from sklearn.metrics import classification_report, f1_score, confusion_matrix

from dataset import SubDataset
from loss import create_criterion
from util import read_json, update_argument, draw_confusion_matrix, seed_everything, get_lr, increment_path, grid_image


def train(data_dir, model_dir, args):
    seed_everything(args.seed)

    base_save_dir = increment_path(os.path.join(model_dir, args.name))
    os.makedirs(base_save_dir, exist_ok=True)

    # -- settings
    use_cuda = torch.cuda.is_available()
    device = torch.device("cuda" if use_cuda else "cpu")

    tasks = ["Mask", "Gender", "Age"]

    with open(os.path.join(base_save_dir, 'config.json'), 'w', encoding='utf-8') as f:
            json.dump(vars(args), f, ensure_ascii=False, indent=4)

    for i, task in enumerate(tasks):
        model = None
        torch.cuda.empty_cache()
        print("="*20 + f" TASK : {task} " + "="*20)
       
        save_dir = os.path.join(base_save_dir, task)
        # -- dataset
        dataset_module = getattr(import_module("dataset"), args.dataset[i])     # 리스트 형태로 입력된 dataset을 순서대로 사용합니다.
        dataset = dataset_module(
            data_dir=data_dir,
        )
        num_classes = dataset.num_classes

        # -- augmentation
        transform_module = getattr(import_module("dataset"), args.augmentation)  # default: BaseAugmentation
        transform = transform_module(
            resize=args.resize,
            mean=dataset.mean,
            std=dataset.std,
        )
        validation_transform_module = getattr(import_module("dataset"), args.valid_augmentation)  # default: BaseAugmentation
        validation_transform = validation_transform_module(
            resize=args.resize,
            mean=dataset.mean,
            std=dataset.std,
        )
        
        # -- data_loader
        train_set, val_set = dataset.split_dataset()
        train_set = SubDataset(train_set, transform=transform)
        val_set = SubDataset(val_set, transform=validation_transform)

        train_loader = DataLoader(
            train_set,
            batch_size=args.batch_size,
            num_workers=multiprocessing.cpu_count()//2,
            shuffle=True,
            pin_memory=use_cuda,
            drop_last=True,
        )

        val_loader = DataLoader(
            val_set,
            batch_size=args.valid_batch_size,
            num_workers=multiprocessing.cpu_count()//2,
            shuffle=False,
            pin_memory=use_cuda,
            drop_last=True,
        )
        
        # -- model
        model_module = getattr(import_module("model"), args.model[i])  # default: BaseModel
        model = model_module(
            args.config,
            num_classes
        ).to(device)
        model = torch.nn.DataParallel(model)

        # -- loss & metric
        criterion = create_criterion(args.criterion, classes=num_classes)  # default: cross_entropy
        opt_module = getattr(import_module("torch.optim"), args.optimizer)  # default: SGD
        optimizer = opt_module(
            filter(lambda p: p.requires_grad, model.parameters()),
            lr=args.lr,
            weight_decay=5e-4
        )
        scheduler = StepLR(optimizer, args.lr_decay_step, gamma=0.5)

        # -- logging
        logger = SummaryWriter(log_dir=save_dir)
       

        best_val_acc = 0
        best_val_loss = np.inf
        best_val_f1_score = 0

        for epoch in range(args.epochs[i]):
            # train loop
            model.train()
            loss_value = 0
            matches = 0
            for idx, train_batch in enumerate(train_loader):
                inputs, labels = train_batch
                inputs = inputs.to(device)
                labels = labels.to(device)

                optimizer.zero_grad()

                outs = model(inputs)
                preds = torch.argmax(outs, dim=-1)
                loss = criterion(outs, labels)

                loss.backward()
                optimizer.step()

                loss_value += loss.item()
                matches += (preds == labels).sum().item()
                if (idx + 1) % args.log_interval == 0:
                    train_loss = loss_value / args.log_interval
                    train_acc = matches / args.batch_size / args.log_interval
                    current_lr = get_lr(optimizer)
                    print(
                        f"Epoch[{epoch}/{args.epochs[i]}]({idx + 1}/{len(train_loader)}) || "
                        f"training loss {train_loss:4.4} || training accuracy {train_acc:4.2%} || lr {current_lr}"
                    )
                    logger.add_scalar("Train/loss", train_loss, epoch * len(train_loader) + idx)
                    logger.add_scalar("Train/accuracy", train_acc, epoch * len(train_loader) + idx)

                    loss_value = 0
                    matches = 0

            scheduler.step()

            # val loop
            with torch.no_grad():
                print("Calculating validation results...")
                model.eval()
                val_loss_items = []
                val_acc_items = []
                predicts, answers = [], []
                figure = None
                for inputs, labels in val_loader:
                    inputs = inputs.to(device)
                    labels = labels.to(device)

                    outs = model(inputs)
                    preds = torch.argmax(outs, dim=-1)

                    loss_item = criterion(outs, labels).item()
                    acc_item = (labels == preds).sum().item()
                    val_loss_items.append(loss_item)
                    val_acc_items.append(acc_item)

                    if figure is None:
                        inputs_np = torch.clone(inputs).detach().cpu().permute(0, 2, 3, 1).numpy()
                        inputs_np = dataset_module.denormalize_image(inputs_np, dataset.mean, dataset.std)
                        figure = grid_image(
                            inputs_np, labels, preds, n=16, shuffle=(args.dataset not in ["MaskSplitByProfileDataset", "DJ_SeparatedDataset"])
                        )

                    answers.extend(labels.cpu().numpy())
                    predicts.extend(preds.cpu().numpy())

                val_loss = np.sum(val_loss_items) / len(val_loader)
                val_acc = np.sum(val_acc_items) / len(val_set)
                val_f1_score = f1_score(answers, predicts, average='macro')
                best_val_loss = min(best_val_loss, val_loss)

                classification_result = classification_report(answers, predicts)

                if val_f1_score > best_val_f1_score:
                    print(f"New best model for val f1 score : {val_f1_score:4.2}! saving the best model..")
                    print(f"New best model for val accuracy : {val_acc:4.2%}!")
                    torch.save(model.module.state_dict(), f"{save_dir}/best.pth")
                    best_val_f1_score = val_f1_score
                    best_val_acc = val_acc

                    draw_confusion_matrix(answers, predicts, save_dir, num_classes)
                    with open(f"{save_dir}/classification_result_of_best_model.txt", "w") as f:
                        f.write(classification_result)
                
                torch.save(model.module.state_dict(), f"{save_dir}/last.pth")
                print(
                    f"[Val] acc : {val_acc:4.2%}, loss: {val_loss:4.2}, f1 score: {val_f1_score:4.2} || "
                    f"best acc : {best_val_acc:4.2%}, best loss: {best_val_loss:4.2}, best f1: {best_val_f1_score:4.2}"
                )
                print(classification_result)
                logger.add_scalar("Val/loss", val_loss, epoch)
                logger.add_scalar("Val/accuracy", val_acc, epoch)
                logger.add_figure("results", figure, epoch)
                print()


if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    from dotenv import load_dotenv
    import os
    load_dotenv(verbose=True)

    # Data and model checkpoints directories
    parser.add_argument('--seed', type=int, default=42, help='random seed (default: 42)')
    parser.add_argument('--epochs', type=int, default=10, help='number of epochs to train (default: 10)')
    parser.add_argument('--dataset', type=str, default='MaskSplitByProfileDataset', help='dataset augmentation type (default: MaskSplitByProfileDataset)')
    parser.add_argument('--augmentation', type=str, default='BaseAugmentation', help='data augmentation type (default: BaseAugmentation)')
    parser.add_argument('--valid_augmentation', type=str, default='BaseAugmentation', help='validation data augmentation type (default: BaseAugmentation)')
    parser.add_argument("--resize", nargs="+", type=int, default=[224, 224], help='resize size for image when training (defualt: [224, 224])')
    parser.add_argument('--batch_size', type=int, default=64, help='input batch size for training (default: 64)')
    parser.add_argument('--valid_batch_size', type=int, default=128, help='input batch size for validing (default: 128)')
    parser.add_argument('--model', type=str, default='BaseModel', help='model type (default: BaseModel)')
    parser.add_argument('--optimizer', type=str, default='Adam', help='optimizer type (default: Adam)')
    parser.add_argument('--lr', type=float, default=1e-3, help='learning rate (default: 1e-3)')
    parser.add_argument('--val_ratio', type=float, default=0.2, help='ratio for validaton (default: 0.2)')
    parser.add_argument('--criterion', type=str, default='cross_entropy', help='criterion type (default: cross_entropy) [focal, f1, label_smoothing]')
    parser.add_argument('--lr_decay_step', type=int, default=20, help='learning rate scheduler deacy step (default: 20)')
    parser.add_argument('--log_interval', type=int, default=20, help='how many batches to wait before logging training status')
    parser.add_argument('--name', default='exp', help='model save at {SM_MODEL_DIR}/{name}')
    parser.add_argument('--config', '-c', default='./config_multi_model.json', type=str, help='config file path 건드릴 필요 없음. (default: ./config.json)')

    # Container environment
    parser.add_argument('--data_dir', type=str, default=os.environ.get('SM_CHANNEL_TRAIN', '/opt/ml/input/data/train/images'))
    parser.add_argument('--model_dir', type=str, default=os.environ.get('SM_MODEL_DIR', './model'))

    args = parser.parse_args()
    configs = read_json(args.config)
    args = update_argument(args, configs)
    print(args)

    data_dir = args.data_dir
    model_dir = args.model_dir

    train(data_dir, model_dir, args)