import argparse
import json
import os

import numpy as np
import torch
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import transforms
from tqdm import tqdm
import albumentations as A
from albumentations.pytorch import ToTensor


from common. dataset import MedicalImageDataset as Dataset
from common.logger import Logger
from common.loss import bce_dice_loss, dice_coef_metric,_fast_hist, jaccard_index
from model.Att_Unet import Att_Unet
from common.utils import log_images


def main(config):
    makedirs(config)
    snapshotargs(config)
    device = torch.device("cpu" if not torch.cuda.is_available() else config.device)

    loader_train, loader_valid = data_loaders(config)
    loaders = {"train": loader_train, "valid": loader_valid}

    unet =Att_Unet()
    unet.to(device)


    best_validation_dsc = 0.0

    optimizer = optim.Adam(unet.parameters(), lr=config.lr,weight_decay=1e-5)
    lr_scheduler= torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', factor=0.5, patience=15, verbose=False)


    logger = Logger(config.logs)
    loss_train = []
    loss_valid = []

    step = 0

    for epoch in tqdm(range(config.epochs), total=config.epochs):
        for phase in ["train", "valid"]:
                if phase == "train":
                    unet.train()
                else:
                    unet.eval()

                validation_pred = []
                validation_true = []
                running_loss = 0.0

                for i, data in enumerate(loaders[phase]):
                    if phase == "train":
                        step += 1

                    x, y_true = data
                    x, y_true = x.to(device), y_true.to(device)

                    optimizer.zero_grad()

                    with torch.set_grad_enabled(phase == "train"):
                            y_pred = unet(x)

                            loss = bce_dice_loss(y_pred, y_true)

                            if phase == "valid":
                                loss_valid.append(loss.item())
                                y_pred_np = y_pred.detach().cpu().numpy()
                                validation_pred.extend(
                                    [y_pred_np[s] for s in range(y_pred_np.shape[0])]
                                )
                                y_true_np = y_true.detach().cpu().numpy()
                                validation_true.extend(
                                    [y_true_np[s] for s in range(y_true_np.shape[0])]
                                )
                                if (epoch % config.vis_freq == 0) or (epoch == config.epochs - 1):
                                    if i * config.batch_size < config.vis_images:
                                        tag = "image/{}".format(i)
                                        num_images = config.vis_images - i * config.batch_size
                                        logger.image_list_summary(
                                            tag,
                                            log_images(x, y_true, y_pred)[:num_images],
                                            step,
                                        )

                            if phase == "train":
                                loss_train.append(loss.item())
                                loss.backward()
                                optimizer.step()
                            running_loss += loss.detach() * x.size(0)

                    if i % 50 == 0:
                        for param_group in optimizer.param_groups:
                            print("Current learning rate is: {}".format(param_group['lr']))



                    if phase == "train" and (step + 1) % 10 == 0:
                        log_loss_summary(logger, loss_train, step)
                        loss_train = []

                print('Epoch [%d/%d], Loss: %.4f, ' %(epoch+1, config.epochs, running_loss/len(loaders[phase].dataset)))

                if phase == "valid":
                    log_loss_summary(logger, loss_valid, step, prefix="val_")
                    mean_dsc,mean_iou = compute_metric(unet,loaders[phase])
                    logger.scalar_summary("val_dsc", mean_dsc, step)
                    logger.scalar_summary("val_iou", mean_iou, step)
                    lr_scheduler.step(mean_dsc)
                    print("\nMean DICE on validation:", mean_dsc)
                    print("Mean IOU on validation:", mean_iou)
                    print("..........................................")

                    if mean_dsc > best_validation_dsc:
                        best_validation_dsc = mean_dsc
                        torch.save(unet.state_dict(), os.path.join(config.weights, "unet.pt"))
                    loss_valid = []



    print("Best validation mean DSC: {:4f}".format(best_validation_dsc))


def data_loaders(config):
    dataset_train, dataset_valid = datasets(config)



    loader_train = DataLoader(
        dataset_train,
        batch_size=config.batch_size,
        num_workers=config.workers
    )
    loader_valid = DataLoader(
        dataset_valid,
        batch_size=config.batch_size,
        num_workers=config.workers

    )

    return loader_train, loader_valid

data_transforms = A.Compose ([
    A.Resize(width = 256, height = 256, p=1.0),
    A.HorizontalFlip(p=0.5),
    A.VerticalFlip(p=0.5),
    A.Rotate((-5,5),p=0.5),
    A.RandomSunFlare(flare_roi=(0, 0, 1, 0.5), angle_lower=0, angle_upper=1,
                                   num_flare_circles_lower=1, num_flare_circles_upper=2,
                                   src_radius=160, src_color=(255, 255, 255),  always_apply=False, p=0.2),
     A.RGBShift (r_shift_limit=10, g_shift_limit=10,
                 b_shift_limit=10, always_apply=False, p=0.2),
    A. ElasticTransform (alpha=2, sigma=15, alpha_affine=25, interpolation=1,
                                      border_mode=4, value=None, mask_value=None,
                                      always_apply=False, approximate=False, p=0.2) ,
    A.Normalize( p=1.0),
    ToTensor(),
])


def datasets(config):
    train = Dataset('train', config.root,
                    transform=data_transforms)


    valid = Dataset('val', config.root,
                    transform=data_transforms)

    return train, valid


def compute_metric(model, loader, threshold=0.3):
    """
    Computes accuracy on the dataset wrapped in a loader

    Returns: accuracy as a float value between 0 and 1
    """
    device = torch.device("cpu" if not torch.cuda.is_available() else config.device)
    #model.eval()
    valloss_one = 0
    valloss_two = 0

    with torch.no_grad():

        for i_step, (data, target) in enumerate(loader):

            data = data.to(device)
            target = target.to(device)


            #prediction = model(x_gpu)

            outputs = model(data)
           # print("val_output:", outputs.shape)

            out_cut = np.copy(outputs.data.cpu().numpy())
            out_cut[np.nonzero(out_cut < threshold)] = 0.0
            out_cut[np.nonzero(out_cut >= threshold)] = 1.0
            hist=_fast_hist(target.data.cpu().numpy(),out_cut,num_classes=2)

            picloss = dice_coef_metric(hist)
            iouloss,_=jaccard_index(hist)
            valloss_one += picloss
            valloss_two +=iouloss


    return valloss_one / i_step,valloss_two/i_step


def log_loss_summary(logger, loss, step, prefix=""):
    logger.scalar_summary(prefix + "loss", np.mean(loss), step)


def makedirs(config):
    os.makedirs(config.weights, exist_ok=True)
    os.makedirs(config.logs, exist_ok=True)


def snapshotargs(config):
    config_file = os.path.join(config.logs, "config.json")
    with open(config_file, "w") as fp:
        json.dump(vars(config), fp)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Finetuning pretrained Unet"
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=16,
        help="input batch size for training (default: 16)",
    )
    parser.add_argument(
        "--epochs",
        type=int,
        default=100,
        help="number of epochs to train (default: 100)",
    )
    parser.add_argument(
        "--lr",
        type=float,
        default=0.001,
        help="initial learning rate (default: 0.001)",
    )
    parser.add_argument(
        "--device",
        type=str,
        default="cuda:0",
        help="device for training (default: cuda:0)",
    )
    parser.add_argument(
        "--workers",
        type=int,
        default=4,
        help="number of workers for data loading (default: 4)",
    )
    parser.add_argument(
        "--vis-images",
        type=int,
        default=100,
        help="number of visualization images to save in log file (default: 200)",
    )
    parser.add_argument(
        "--vis-freq",
        type=int,
        default=10,
        help="frequency of saving images to log file (default: 10)",
    )
    parser.add_argument(
        "--weights", type=str, default="./weights", help="folder to save weights"
    )
    parser.add_argument(
        "--logs", type=str, default="./logs", help="folder to save logs"
    )
    parser.add_argument(
        "--root", type=str, default="./medico2020", help="root folder with images"
    )
    parser.add_argument(
        "--image-size",
        type=int,
        default=256,
        help="target input image size (default: 256)",
    )
    parser.add_argument(
        "--aug-scale",
        type=int,
        default=0.05,
        help="scale factor range for augmentation (default: 0.05)",
    )
    parser.add_argument(
        "--aug-angle",
        type=int,
        default=6,
        help="rotation angle range in degrees for augmentation (default: 15)",
    )
    config = parser.parse_args()
    main(config)