import numpy as np

import pandas as pd

import os

import glob

import torch

import torch.nn as nn

import random

import itertools

from sklearn.metrics import roc_auc_score

import torch.nn.functional as F

from scipy import interpolate

data_path_img = '<path-to-binary-image-files>'

data_file     = '<path-and_name-of-csv-datafile>'

image_size    = 310 # this is the size (310x310) of the MIP images used in the paper. Adjust to fit your images.

def get_datasets_singleview(transform=None, norm=None, balance=False, split_index=0):

    split = 'split'+str(split_index)

    df = pd.read_csv(data_file)

    # Balance weight

    weight_neg_pos = [1-(df.target==0).sum()/len(df), 1-(df.target==1).sum()/len(df)]    

    # Read split

    df_train = df[df[split]=='train'].drop(df.filter(regex='split').columns,axis=1)

    train_dset = dataset_singleview(df_train, transform=transform, norm=norm)

    trainval_dset = dataset_singleview_center(df_train, transform=None, norm=norm)

    # Val split

    df_val = df[df[split]=='val'].drop(df.filter(regex='split').columns,axis=1)

    val_dset = dataset_singleview_center(df_val, transform=None, norm=norm)

    # Test split

    df_test = df[df[split]=='test'].drop(df.filter(regex='split').columns,axis=1)

    test_dset = dataset_singleview_center(df_test, transform=None, norm=norm)

    return train_dset,trainval_dset,val_dset,test_dset,weight_neg_pos

def get_bbox(img):

    rows = np.any(img, axis=1)

    cols = np.any(img, axis=0)

    rmin, rmax = np.where(rows)[0][[0, -1]]

    cmin, cmax = np.where(cols)[0][[0, -1]]

    return img[rmin:rmax, cmin:cmax]

def pad2square_random(image, size):

    out = np.zeros((size,size))

    # Sample offset

    maxr = size - image.shape[0]

    maxc = size - image.shape[1]

    offsetc = np.random.randint(0, maxc)

    offsetr = np.random.randint(0, maxr)

    # Place image

    out[offsetr:offsetr+image.shape[0], offsetc:offsetc+image.shape[1]] = image

    return out

def pad2square_center(image, size):

    # Place image

    out = np.zeros((size,size))

    out[int((size-image.shape[0])/2):int((size-image.shape[0])/2)+image.shape[0],int((size-image.shape[1])/2):int((size-image.shape[1])/2)+image.shape[1]] = image

    return out

def clip_and_normalize_SUVimage(img):

    mu  = 2.13

    std = 3.39

    q   = 30.00

    img = np.clip(img,0.,q)

    return (img-mu)/std

def get_image(df, transform, norm):

    name = glob.glob(os.path.join(data_path_img,df.filename))

    if not name:

        print('File not found:',name)

    img  = np.fromfile(os.path.join(data_path_img, name[0]), dtype='float32')

    img  = np.reshape( img,[df.matrix_size_1, df.matrix_size_2])

    # Find bbox

    img = get_bbox(img)

    # Pad randomly

    img = pad2square_random(img, image_size)

    # Norm

    if norm:

        img = clip_and_normalize_SUVimage(img)

    # Make Tensors

    img = torch.FloatTensor(img).unsqueeze(0)

    if transform is not None:

        img = [transform(x) for x in img]

        img = torch.stack(img)

    return img

def get_image_center(df, transform, norm):

    name = glob.glob(os.path.join(data_path_img,df.filename))

    if not name:

        print('File not found:',name)

    img  = np.fromfile(os.path.join(data_path_img, name[0]), dtype='float32')

    img  = np.reshape( img,[df.matrix_size_1, df.matrix_size_2])

    # Find bbox

    img = get_bbox(img)

    # Pad randomly

    img = pad2square_center(img, image_size)

    # Norm

    if norm:

        img = clip_and_normalize_SUVimage(img)

    # Make Tensors

    img = torch.FloatTensor(img).unsqueeze(0)

    if transform is not None:

        img = [transform(x) for x in img]

        img = torch.stack(img)

    return img

class dataset_singleview(torch.utils.data.Dataset):

    def __init__(self, df, transform=None, norm=False):

        self.df = df.copy()

        self.transform = transform

        self.norm = norm

    def errors(self, probs):

        df = self.df.copy()

        df['p'] = probs

        df['pred'] = (df.p >= 0.5).astype(int)

        fpr = ((df.pred!=df.target) & (df.target==0)).sum() / (df.target==0).sum()

        fnr = ((df.pred!=df.target) & (df.target==1)).sum() / (df.target==1).sum()

        ber = (fpr + fnr) / 2.

        ## Calculate auc

        auc = roc_auc_score(df.target, df.p)

        return auc, ber, fpr, fnr

    def __getitem__(self, index): 

        df = self.df.iloc[index]

        # Read image

        img = get_image(df, self.transform, self.norm)

        return img, df.target

    def __len__(self):

        return len(self.df)

class dataset_singleview_center(torch.utils.data.Dataset):

    def __init__(self, df, transform=None, norm=False):

        self.df = df.copy()

        self.transform = transform

        self.norm = norm

    def errors(self, probs):

        df = self.df.copy()

        df['p'] = probs

        df['pred'] = (df.p >= 0.5).astype(int)

        fpr = ((df.pred!=df.target) & (df.target==0)).sum() / (df.target==0).sum()

        fnr = ((df.pred!=df.target) & (df.target==1)).sum() / (df.target==1).sum()

        ber = (fpr + fnr) / 2.

        ## Calculate auc

        auc = roc_auc_score(df.target, df.p)

        return auc, ber, fpr, fnr

    def __getitem__(self, index): 

        df = self.df.iloc[index]

        # Read image

        img = get_image_center(df, self.transform, self.norm)

        return img, df.target

    def __len__(self):

        return len(self.df)

class RandomFlip(object):

    """Randomly flip the 2D image.

    """

    def __call__(self, image):

        # Random flip: none, 0=vertical, 1=horizontal

        flip = random.choice((None,0,1))

        if flip is not None:

            if flip==0:

                image   = image[range(image.shape[flip]-1,-1,-1),:]

            elif flip==1:

                image   = image[:,range(image.shape[flip]-1,-1,-1)]       

        return image

class RandomFlipLeftRight(object):

    """Randomly flip all channels of the 2D image.

    """

    def __call__(self, image):

        # Random flip: none, 0=vertical, 1=horizontal

        flip = random.choice((None,1))

        if flip is not None:

            image = image[:,range(image.shape[1]-1,-1,-1)]

        return image

class RandomRot90(object):

    """Randomly rotate the 2D image by n*90 degrees.

    """

    def __call__(self, image):

        # Random  90 rotation

        rot = random.randint(0,3)

        if rot != 0:

            image = torch.rot90(image, rot, (0,1))

        return image

class RandomScale(object):

    """Randomly scale the 2D image.

    """

    def __call__(self, image):

        scale = np.random.uniform(low=0.85, high=1.15, size=1)

        image = image*scale[0]

        return image

class RandomNoise(object):

    """Randomly gauss noise the 2D image.

    """

    def __call__(self, image):

        noise = random.choice((None,1))

        if noise is not None:

            image[image<0] = 0

            level = np.random.uniform(low=0.001, high=0.02, size=1)

            sigma = np.random.uniform(low=0.01, high=0.1, size=1)

            sigma = sigma[0]*image+level[0]

            gauss = torch.normal(0,sigma)

            image = image + gauss

            image[image<0] = 0

        return image