import argparse

import os

import numpy as np

import torch

from skimage.io import imsave

from torch.utils.data import DataLoader

import albumentations as A

from albumentations.pytorch import ToTensor

from torchvision import transforms

from tqdm import tqdm

import cv2

from common. dataset import MedicalImageDataset as Dataset

from common.loss import bce_dice_loss,  dice_coef_metric

from model.Att_Unet import Att_Unet

from common.utils import log_images, gray2rgb,  outline

def main(config):

    makedirs(config)

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

   # logger = Logger(config.logs)

    loader,dataset = data_loader(config)

    with torch.set_grad_enabled(False):

        unet = Att_Unet()

        state_dict = torch.load(config.weights, map_location=device)

        unet.load_state_dict(state_dict)

        unet.eval()

        unet.to(device)

        input_list = []

        pred_list = []

        #true_list = []

        #loss_test = []

        step=0

        for i, data in tqdm(enumerate(loader)):

            step += 1

            x= data

            x= x.to(device)

            y_pred = unet(x)

            #loss = bce_dice_loss(y_pred, y_true)

            #loss_test.append(loss.item())

            y_pred_np = y_pred.detach().cpu().numpy()

            pred_list.extend([y_pred_np[s] for s in range(y_pred_np.shape[0])])

            #y_true_np = y_true.detach().cpu().numpy()

            #true_list.extend([y_true_np[s] for s in range(y_true_np.shape[0])])

            x_np = x.detach().cpu().numpy()

            input_list.extend([x_np[s] for s in range(x_np.shape[0])])

        if config.mask_outline==True:

            for i in range(len(input_list)):

                image = gray2rgb(np.squeeze(input_list[i][0,:,:]))

                image = outline(image, pred_list[i][0,:,:], color=[255, 0, 0])

                image=image.astype("uint8")

                filename="{}.png".format(i)

                filepath = os.path.join(config.predictions, filename)

                imsave(filepath, image)

        else:

            for i in range(len(input_list)):

                 img_name=dataset.imgs[i].split("/" )[-1:][0]

                 org_shape=cv2.imread(dataset.imgs[i]).shape

                 mask=pred_list[i][0,:,:]

                 mask[np.nonzero(mask<0.3)]=0.0

                 mask[np.nonzero(mask>0.3)]=255.0

                 mask=mask.astype("uint8")

                 mask_copy=mask.copy()

                 mk=cv2.resize(mask,(org_shape[0],org_shape[1]),interpolation = cv2.INTER_AREA)

                 filename="{}".format(img_name)

                 filepath = os.path.join(config.predictions, filename)

                 imsave(filepath, mk)

def makedirs(config):

    os.makedirs(config.predictions, exist_ok=True)

data_transforms = A.Compose ([

    A.Resize(width = 256, height = 256, p=1.0),

    A.Normalize( p=1.0),

    ToTensor()

])

def data_loader(config):

    dataset = Dataset('test', config.root,

                    transform=data_transforms

                    )

    loader =DataLoader(

        dataset,

        batch_size=config.batch_size,

        num_workers=4

    )

    return loader,dataset

"""

def compute_iou(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 = 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

            picloss = dice_coef_metric(out_cut, target.data.cpu().numpy())

            valloss += picloss

        #print("Threshold:  " + str(threshold) + "  Validation DICE score:", valloss / i_step)

    return valloss / i_step

"""

"""

def log_loss_summary(logger, loss, step, prefix=""):

    logger.scalar_summary(prefix + "loss", np.mean(loss), step)

"""

if __name__ == "__main__":

    parser = argparse.ArgumentParser(

        description="Inference for segmentation of polyps in GI"

    )

    parser.add_argument(

        "--device",

        type=str,

        default="cuda:0",

        help="device for training (default: cuda:0)",

    )

    parser.add_argument(

        "--batch-size",

        type=int,

        default=32,

        help="input batch size for training (default: 32)",

    )

    parser.add_argument(

        "--weights", type=str, default="./weights/unet.pt", required=True, help="path to weights file"

    )

    parser.add_argument(

        "--root", type=str, default="./medico2020", help="root folder with images"

    )

    parser.add_argument(

        "--mask_outline", type=bool, default=False, help="If True ,draws border line of mask else returns binary mask"

    )

    parser.add_argument(

        "--image-size",

        type=int,

        default=256,

        help="target input image size (default: 256)",

    )

    parser.add_argument(

        "--predictions",

        type=str,

        default="./predictions",

        help="folder for saving images with prediction outlines",

    )

    parser.add_argument(

        "--vis-images",

        type=int,

        default=20,

        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(

        "--logs", type=str, default="./test_logs", help="folder to save logs"

    )

    config = parser.parse_args()

    main(config)