import os
import os.path as osp
import random
import matplotlib.pyplot as plt
import pandas as pd
import torch
import torch.nn as nn
from tensorboardX import SummaryWriter
from sklearn.metrics import confusion_matrix
import seaborn as sn
from tqdm import tqdm
import argparse
from torchvision import transforms
from datasets import WasteNetDataset
from torch.utils.data import DataLoader

# 自前の関数群
from apex import amp, optimizers
from utils import *

parser = argparse.ArgumentParser()

parser.add_argument("model_name", type=str) #ex) "efficient_net_b5_1"
parser.add_argument("split_num", type=int)

args = parser.parse_args()

MODEL_NAME = args.model_name
LOGGING_DIR = "undersample+bagging"
SPLIT_NUM = args.split_num
START_SEQUENCE = 0

BATCH_SIZE = 4
IMG_WIDTH = 384
IMG_HEIGHT = 288
SHOW_INFO_FREQ = 100
ALL_SEQUENCE_NUM = 9

# 学習パラメータ系
EPOCH_NUM = 10

# apex関連
# opt_level = 'O1'

if __name__ == "__main__":
    for sequence_num in range(START_SEQUENCE, ALL_SEQUENCE_NUM):
        torch.backends.cudnn.benchmark = True

        # 再現性のためのrandom seedを設定する
        manual_seed = 999
        print("Random Seed: ", manual_seed)
        random.seed(manual_seed)
        torch.manual_seed(manual_seed)

        # どのデバイスで実行するかをきめる
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print(f"Use device: {device}")

        # データ選ぶ部
        # data_table = pd.read_csv("../input/train.csv", header=None).values.tolist()
        # train_dset, test_dset = train_test_split_by_sequence(data_table, test_sequence_num=sequence_num, balanced_test=True)
        train_dset, test_dset = load_dsets_for_bagging(sequence_num, SPLIT_NUM, balanced_test=True)

        if MODEL_NAME.startswith("efficient"):
            model = load_efficientnet("efficientnet-b5", IMG_HEIGHT, IMG_WIDTH)
        elif MODEL_NAME.startswith("resnet"):
            model = load_resnet50()
        elif MODEL_NAME.startswith("densenet"):
            model = load_densenet161()
        elif MODEL_NAME.startswith("deit"):
            model = load_deit(IMG_HEIGHT, IMG_WIDTH)
        elif MODEL_NAME.startswith("coatnet"):
            model = load_CoAtNet()
        elif MODEL_NAME.startswith("effnetv2"):
            model = load_effnetv2()
        elif MODEL_NAME.startswith("mobile"):
            model = load_mobilenet_v2()
        else:
            raise NotImplementedError()

        model.to(device)
        optimizer = torch.optim.Adam(model.parameters())
        # model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level)

        # データローダー部

        if MODEL_NAME.startswith("coatnet"):
            transform = transforms.Compose([
                transforms.ToTensor(),
                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
                transforms.Resize((224, 224))
            ])
        else:
            transform = transforms.Compose([
                transforms.ToTensor(),
                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
            ])

        train_datasets = WasteNetDataset(train_dset, IMG_WIDTH, IMG_HEIGHT, transform)
        test_datasets = WasteNetDataset(test_dset, IMG_WIDTH, IMG_HEIGHT, transform)
        train_loader = DataLoader(train_datasets, batch_size=BATCH_SIZE, shuffle=True, num_workers=2, pin_memory=True)
        test_loader = DataLoader(test_datasets, batch_size=BATCH_SIZE, num_workers=0, pin_memory=True)

        # 損失関数
        criterion = nn.CrossEntropyLoss()

        # tensorboardX
        os.makedirs(osp.join(f"../logs", LOGGING_DIR), exist_ok=True)
        os.makedirs(osp.join(f"../logs", LOGGING_DIR, MODEL_NAME), exist_ok=True)

        os.makedirs(osp.join("../logs", LOGGING_DIR, MODEL_NAME, f"sequence{sequence_num}"), exist_ok=True)
        writer = SummaryWriter(osp.join("../logs", LOGGING_DIR, MODEL_NAME, f"sequence{sequence_num}"))
        classes = ("bad data", "good data")

        for epoch in range(1, EPOCH_NUM + 1):

            # train phase
            model.train()
            running_loss = 0.0
            correct_num = 0
            pred_target_num = 0

            for i, (imgs, labels) in enumerate(tqdm(train_loader)):
                model.zero_grad(set_to_none=True)

                imgs = imgs.to(device)
                labels = labels.to(device)

                pred = model(imgs)
                loss = criterion(pred, labels)
                running_loss += loss.item()

                pred_target_num += pred.shape[0]
                correct_num += torch.sum(torch.argmax(pred, dim=1) == labels).item()

                loss.backward()
                # with amp.scale_loss(loss, optimizer) as scaled_loss:
                #     scaled_loss.backward()

                optimizer.step()

                if i % SHOW_INFO_FREQ == (SHOW_INFO_FREQ - 1):
                    print("train")
                    print(f"[sequence{sequence_num} {epoch}, {i}] {running_loss}")
                    print(f"[avg] {correct_num / pred_target_num}")
                    print()

                    writer.add_scalar("train_loss", running_loss, train_loader.__len__() * (epoch - 1) + i)
                    writer.add_scalar("train_acc", correct_num / pred_target_num, train_loader.__len__() * (epoch - 1) + i)

                    running_loss = 0.0
                    correct_num = 0.0
                    pred_target_num = 0.0

            # test phase
            with torch.no_grad():
                model.eval()
                running_loss = 0.0
                correct_num = 0
                pred_target_num = 0

                y_pred = []
                y_true = []
                for i, (imgs, labels) in enumerate(tqdm(test_loader)):
                    imgs = imgs.to(device)
                    labels = labels.to(device)

                    pred = model(imgs)
                    loss = criterion(pred, labels)
                    running_loss += loss.item()

                    pred_target_num += pred.shape[0]
                    pred_index = torch.argmax(pred, dim=1)
                    correct_num += torch.sum(pred_index == labels).item()

                    y_pred.extend(pred_index.data.cpu().numpy())
                    y_true.extend(labels.data.cpu().numpy())

            print("test")
            print(f"[sequence{sequence_num} {epoch}] {running_loss}")
            print(f"[avg] {correct_num / pred_target_num}")
            print()

            writer.add_scalar("test_loss", running_loss, train_loader.__len__() * (epoch - 1))
            writer.add_scalar("test_acc", correct_num / pred_target_num, train_loader.__len__() * (epoch - 1))
            torch.save(model.state_dict(), osp.join("../logs", LOGGING_DIR, MODEL_NAME, f"sequence{sequence_num}", f"test{epoch}_acc{(correct_num / pred_target_num):.2f}_weights.pth"))

            # 混同行列の保存
            cf_matrix = confusion_matrix(y_true, y_pred)
            df_cm = pd.DataFrame(cf_matrix, index=[f"{i}(label)" for i in classes], columns=[f"{i}(predicted)" for i in classes])
            plt.figure(figsize=(12, 7))
            show_fig = sn.heatmap(df_cm, annot=True).get_figure()
            writer.add_figure("Confusion matrix", show_fig, epoch)
    del model
    torch.cuda.empty_cache()