# -*- coding: utf-8 -*-

"""

Created on Wed Sep  8 10:27:47 2021

@author: m.beuque

"""

from matplotlib.image import imread

import cv2

import tqdm

import pandas as pd

import os

import numpy as np

from torch.utils import data

import torchvision.transforms as transforms

from sklearn.utils import shuffle

import torch

def generate_dataset_tissue_type(main_path,path_data,SEED):

    data = pd.read_csv(path_data,sep = ',' )

    X = []

    y = []

    paths = []

    for slide in tqdm.tqdm(os.listdir(os.path.join(main_path, 'Slides'))):

        tile_path = os.path.join(main_path, 'Slides',slide,'tiles')

        gland = data[(data['labels']=='stroma') & (data['dataset_name']==slide)]['image_name']

        tissue = data[(data['labels']=='epithelial tissue') & (data['dataset_name']==slide)]['image_name']

        gland = list(gland)

        tissue = list(tissue)

        for image_path in gland:

            if os.path.isfile(os.path.join(tile_path, image_path)):

                X.append(imread(os.path.join(tile_path, image_path)))

                y.append(0)

                paths.append(os.path.join(tile_path, image_path))

        for image_path in tissue:

            if os.path.isfile(os.path.join(tile_path, image_path)):

                X.append(imread(os.path.join(tile_path, image_path)))

                y.append(1)

                paths.append(os.path.join(tile_path, image_path))

    X = np.array(X)

    for j, elmt in enumerate(X):

        if elmt.shape !=(96,96,3):

            X[j] = cv2.resize(elmt,(96,96),interpolation = cv2.INTER_CUBIC)

    X,y = shuffle(X,y,random_state=SEED)

    return X, y, paths

def generate_dataset_grade(main_path,path_data): 

    #path_data is the path to the csv containing the information of the ing or testing or validation dataset

    #main_path contains the folder "Slides" were the H&E tiles where stored

    data = pd.read_csv(path_data,sep = ',' )

    X = []

    y = []

    paths = []

    for slide in tqdm.tqdm(os.listdir(os.path.join(main_path, 'Slides'))):

        tile_path = os.path.join(main_path, 'Slides',slide,'tiles')

        healthy = data[(data['labels']=='non-dysplasia') & (data['dataset_name']==slide)]['image_name'] 

        lowgrade = data[(data['labels']=='low grade') & (data['dataset_name']==slide)]['image_name']

        highgrade = data[(data['labels']=='high grade') & (data['dataset_name']==slide)]['image_name']

        healthy = list(healthy)

        lowgrade = list(lowgrade)

        highgrade = list(highgrade)

        for image_path in healthy:

            if os.path.isfile(os.path.join(tile_path, image_path)):

                X.append(imread(os.path.join(tile_path, image_path)))

                y.append("non-dysplasia")

                paths.append(os.path.join(tile_path, image_path))

            else:

                print("error for non-dysplasia")

        for image_path in lowgrade:

            if os.path.isfile(os.path.join(tile_path, image_path)):

                X.append(imread(os.path.join(tile_path, image_path)))

                y.append("low grade")

                paths.append(os.path.join(tile_path, image_path))

            else :

                print("error for low grade")

        for image_path in highgrade:

            if os.path.isfile(os.path.join(tile_path, image_path)):

                X.append(imread(os.path.join(tile_path, image_path)))

                y.append("high grade")

                paths.append(os.path.join(tile_path, image_path))

            else :

                print("error for high grade")

    #rescale the images to the same size

    for j, elmt in enumerate(X):

        if elmt.shape !=(96,96,3):

            X[j] = cv2.resize(elmt,(96,96),interpolation = cv2.INTER_CUBIC)

    X = np.array(X)

    y = np.array(y)

    return X,y,paths

#regular dataset generation

class CancerDataset(data.Dataset):

    'Characterizes a dataset for PyTorch'

    def __init__(self,X_all, y_all, transform = transforms.Compose([transforms.CenterCrop(64),transforms.ToTensor()])):

        'Initialization'

        self.labels = y_all

        self.list_IDs = X_all

        self.transform = transform

        self.image_files_list = [str(s) for s in range(len(self.list_IDs))]

    def __len__(self):

        'Denotes the total number of samples'

        return len(self.list_IDs)

    def __getitem__(self, index):

        'Generates one sample of data'

        # Select sample

        X = self.list_IDs[index]

        X = self.transform(image=X)

        X = X['image']

        # Load data and get label

        y = self.labels[index]

        return X, y

def df_dl_features(X,paths,data_transforms,classifier):

    features = {}

    for i,temp_X in tqdm.tqdm(enumerate(X)):

        tensor_X=data_transforms(image=temp_X)

        tensor_X=tensor_X["image"]

        tensor_X.unsqueeze_(0)

        output=torch.flatten(classifier(tensor_X)).detach().numpy()

        features[paths[i]]=output.flatten()

    features=pd.DataFrame.from_dict(features)

    features=features.T

    return features