"""
datasets.py
=======================
Houses the DynamicImageDataset class, also functions to help with image color channel normalization, transformers, etc..
"""

import torch
from torchvision import transforms
import os
import dask
#from dask.distributed import Client; Client()
import dask.array as da, pandas as pd, numpy as np
from pathflowai.utils import *
import pysnooper
import nonechucks as nc
from torch.utils.data import Dataset, DataLoader
import random
import albumentations as alb
import copy
from albumentations import pytorch as albtorch
from sklearn.preprocessing import LabelBinarizer
from sklearn.utils.class_weight import compute_class_weight
from pathflowai.losses import class2one_hot
import cv2
from scipy.ndimage.morphology import generate_binary_structure
from dask_image.ndmorph import binary_dilation
cv2.setNumThreads(0)
cv2.ocl.setUseOpenCL(False)


def RandomRotate90():
	"""Transformer for random 90 degree rotation image.

	Returns
	-------
	function
		Transformer function for operation.

	"""
	return (lambda img: img.rotate(random.sample([0, 90, 180, 270], k=1)[0]))

def get_data_transforms(patch_size = None, mean=[], std=[], resize=False, transform_platform='torch', elastic=True, user_transforms=dict()):
	"""Get data transformers for training test and validation sets.

	Parameters
	----------
	patch_size:int
		Original patch size being transformed.
	mean:list of float
		Mean RGB
	std:list of float
		Std RGB
	resize:int
		Which patch size to resize to.
	transform_platform:str
		Use pytorch or albumentation transforms.
	elastic:bool
		Whether to add elastic deformations from albumentations.

	Returns
	-------
	dict
		Transformers.

	"""
	transform_dict=dict(torch=dict(
									colorjitter=lambda kargs: transforms.ColorJitter(**kargs),
									hflip=lambda kargs: transforms.RandomHorizontalFlip(),
									vflip=lambda kargs: transforms.RandomVerticalFlip(),
									r90= lambda kargs: RandomRotate90()
									),
						albumentations=dict(
							huesaturation=lambda kargs: alb.augmentations.transforms.HueSaturationValue(**kargs),
							flip=lambda kargs: alb.augmentations.transforms.Flip(**kargs),
							transpose=lambda kargs: alb.augmentations.transforms.Transpose(**kargs),
							affine=lambda kargs: alb.augmentations.transforms.ShiftScaleRotate(**kargs),
							r90=lambda kargs: alb.augmentations.transforms.RandomRotate90(**kargs),
							elastic=lambda kargs: alb.augmentations.transforms.ElasticTransform(**kargs)
						))
	if 'normalization' in user_transforms:
		mean=user_transforms['normalization'].pop('mean')
		std=user_transforms['normalization'].pop('std')
		del user_transforms['normalization']
	default_transforms=dict() # add normalization custom
	default_transforms['torch']=dict(
							colorjitter=dict(brightness=0.8, contrast=0.8, saturation=0.8, hue=0.5),
							hflip=dict(),
							vflip=dict(),
							r90=dict())
	default_transforms['albumentations']=dict(
							huesaturation=dict(hue_shift_limit=20, sat_shift_limit=30, val_shift_limit=20, p=0.5),
							r90=dict(p=0.5),
							elastic=dict(p=0.5))
	main_transforms = default_transforms[transform_platform] if not user_transforms else user_transforms
	print(main_transforms)
	train_transforms=[transform_dict[transform_platform][k](v) for k,v in main_transforms.items()]
	torch_init=[transforms.ToPILImage(),transforms.Resize((patch_size,patch_size)),transforms.CenterCrop(patch_size)]
	albu_init=[alb.augmentations.transforms.Resize(patch_size, patch_size),
				alb.augmentations.transforms.CenterCrop(patch_size, patch_size)]
	tensor_norm=[transforms.ToTensor(),transforms.Normalize(mean if mean else [0.7, 0.6, 0.7], std if std is not None else [0.15, 0.15, 0.15])] #mean and standard deviations for lung adenocarcinoma resection slides
	data_transforms = { 'torch': {
		'train': transforms.Compose(torch_init+train_transforms+tensor_norm),
		'val': transforms.Compose([
			transforms.ToPILImage(),
			transforms.Resize((patch_size,patch_size)),
			transforms.CenterCrop(patch_size),
			transforms.ToTensor(),
			transforms.Normalize(mean if mean else [0.7, 0.6, 0.7], std if std is not None else [0.15, 0.15, 0.15])
		]),
		'test': transforms.Compose([
			transforms.ToPILImage(),
			transforms.Resize((patch_size,patch_size)),
			transforms.CenterCrop(patch_size),
			transforms.ToTensor(),
			transforms.Normalize(mean if mean else [0.7, 0.6, 0.7], std if std is not None else [0.15, 0.15, 0.15])
		]),
		'pass': transforms.Compose([
			transforms.ToPILImage(),
			transforms.CenterCrop(patch_size),
			transforms.ToTensor(),
		])
	},
	'albumentations':{
	'train':alb.core.composition.Compose(albu_init+train_transforms),
	'val':alb.core.composition.Compose([
		alb.augmentations.transforms.Resize(patch_size, patch_size),
		alb.augmentations.transforms.CenterCrop(patch_size, patch_size)
	]),
	'test':alb.core.composition.Compose([
		alb.augmentations.transforms.Resize(patch_size, patch_size),
		alb.augmentations.transforms.CenterCrop(patch_size, patch_size)
	]),
	'normalize':transforms.Compose([transforms.Normalize(mean if mean else [0.7, 0.6, 0.7], std if std is not None else [0.15, 0.15, 0.15])])
	}}

	return data_transforms[transform_platform]

def create_transforms(mean, std):
	"""Create transformers.

	Parameters
	----------
	mean:list
		See get_data_transforms.
	std:list
		See get_data_transforms.

	Returns
	-------
	dict
		Transformers.

	"""
	return get_data_transforms(patch_size = 224, mean=mean, std=std, resize=True)



def get_normalizer(normalization_file, dataset_opts):
	"""Find mean and standard deviation of images in batches.

	Parameters
	----------
	normalization_file:str
		File to store normalization information.
	dataset_opts:type
		Dictionary storing information to create DynamicDataset class.

	Returns
	-------
	dict
		Stores RGB mean, stdev.

	"""
	if os.path.exists(normalization_file):
		norm_dict = torch.load(normalization_file)
	else:
		norm_dict = {'normalization_file':normalization_file}

	if 'normalization_file' in norm_dict:

		transformers = get_data_transforms(patch_size = 224, mean=[], std=[], resize=True, transform_platform='torch')

		dataset_opts['transformers']=transformers
		#print(dict(pos_annotation_class=pos_annotation_class, segmentation=segmentation, patch_size=patch_size, fix_names=fix_names, other_annotations=other_annotations))

		dataset = DynamicImageDataset(**dataset_opts)#nc.SafeDataset(DynamicImageDataset(**dataset_opts))

		if dataset_opts['classify_annotations']:
			dataset.binarize_annotations()

		dataloader = DataLoader(dataset, batch_size=128, shuffle=True, num_workers=4)

		all_mean = torch.tensor([0.,0.,0.],dtype=torch.float)#[]

		all_std = torch.tensor([0.,0.,0.],dtype=torch.float)

		if torch.cuda.is_available():
			all_mean=all_mean.cuda()
			all_std=all_std.cuda()

		with torch.no_grad():
			for i,(X,_) in enumerate(dataloader): # x,3,224,224
				if torch.cuda.is_available():
					X=X.cuda()
				all_mean += torch.mean(X, (0,2,3))
				all_std += torch.std(X, (0,2,3))

		N=i+1

		all_mean /= float(N) #(np.array(all_mean).mean(axis=0)).tolist()
		all_std /= float(N) #(np.array(all_std).mean(axis=0)).tolist()

		all_mean = all_mean.detach().cpu().numpy().tolist()
		all_std = all_std.detach().cpu().numpy().tolist()

		torch.save(dict(mean=all_mean,std=all_std),norm_dict['normalization_file'])

		norm_dict = torch.load(norm_dict['normalization_file'])
	return norm_dict

def segmentation_transform(img,mask, transformer, normalizer, alb_reduction):
	"""Run albumentations and return an image and its segmentation mask.

	Parameters
	----------
	img:array
		Image as array
	mask:array
		Categorical pixel by pixel.
	transformer :
		Transformation object.

	Returns
	-------
	tuple arrays
		Image and mask array.

	"""
	res=transformer(True, image=img, mask=mask)
	#res_mask_shape = res['mask'].size()
	return normalizer(torch.tensor(np.transpose(res['image']/alb_reduction,axes=(2,0,1)),dtype=torch.float)).float(), torch.tensor(res['mask']).long()#.view(res_mask_shape[0],res_mask_shape[1],res_mask_shape[2])

class DilationJitter:
	def __init__(self, dilation_jitter=dict(), segmentation=True, train_set=False):
		if dilation_jitter and segmentation and train_set:
			self.run_jitter=True
			self.dilation_jitter=dilation_jitter
			self.struct=generate_binary_structure(2,1) #structure=self.struct,
		else:
			self.run_jitter=False


	def __call__(self, mask):
		if self.run_jitter:
			for k in self.dilation_jitter:
				amount_jitter=int(round(max(np.random.normal(self.dilation_jitter[k]['mean'],
														self.dilation_jitter[k]['std']),1)))
				#print((mask==k).compute())
				mask[binary_dilation(mask==k,structure=self.struct,iterations=amount_jitter)]=k

		return mask


class DynamicImageDataset(Dataset):
	"""Generate image dataset that accesses images and annotations via dask.

	Parameters
	----------
	dataset_df:dataframe
		Dataframe with WSI, which set it is in (train/test/val) and corresponding WSI labels if applicable.
	set:str
		Whether train, test, val or pass (normalization) set.
	patch_info_file:str
		SQL db with positional and annotation information on each slide.
	transformers:dict
		Contains transformers to apply on images.
	input_dir:str
		Directory where images comes from.
	target_names:list/str
		Names of initial targets, which may be modified.
	pos_annotation_class:str
		If selected and predicting on WSI, this class is labeled as a positive from the WSI, while the other classes are not.
	other_annotations:list
		Other annotations to consider from patch info db.
	segmentation:bool
		Conducting segmentation task?
	patch_size:int
		Patch size.
	fix_names:bool
		Whether to change the names of dataset_df.
	target_segmentation_class:list
		Now can be used for classification as well, matched with two below options, samples images only from this class. Can specify this and below two options multiple times.
	target_threshold:list
		Sampled only if above this threshold of occurence in the patches.
	oversampling_factor:list
		Over sample them at this amount.
	n_segmentation_classes:int
		Number classes to segment.
	gdl:bool
		Using generalized dice loss?
	mt_bce:bool
		For multi-target prediction tasks.
	classify_annotations:bool
		For classifying annotations.

	"""
	# when building transformers, need a resize patch size to make patches 224 by 224
	#@pysnooper.snoop('init_data.log')
	def __init__(self,dataset_df, set, patch_info_file, transformers, input_dir, target_names, pos_annotation_class, other_annotations=[], segmentation=False, patch_size=224, fix_names=True, target_segmentation_class=-1, target_threshold=0., oversampling_factor=1., n_segmentation_classes=4, gdl=False, mt_bce=False, classify_annotations=False, dilation_jitter=dict(), modify_patches=True):

		#print('check',classify_annotations)
		reduce_alb=True
		self.patch_size=patch_size
		self.input_dir = input_dir
		self.alb_reduction=255. if reduce_alb else 1.
		self.transformer=transformers[set]
		original_set = copy.deepcopy(set)
		if set=='pass':
			set='train'
		self.targets = target_names
		self.mt_bce=mt_bce
		self.set = set
		self.segmentation = segmentation
		self.alb_normalizer=None
		if 'normalize' in transformers:
			self.alb_normalizer = transformers['normalize']
		if len(self.targets)==1:
			self.targets = self.targets[0]
		if original_set == 'pass':
			self.transform_fn = lambda x,y: (self.transformer(x), torch.tensor(1.,dtype=torch.float))
		else:
			if self.segmentation:
				self.transform_fn = lambda x,y: segmentation_transform(x,y, self.transformer, self.alb_normalizer, self.alb_reduction)
			else:
				if 'p' in dir(self.transformer):
					self.transform_fn = lambda x,y: (self.alb_normalizer(torch.tensor(np.transpose(self.transformer(True, image=x)['image']/self.alb_reduction,axes=(2,0,1)),dtype=torch.float)), torch.from_numpy(y).float())
				else:
					self.transform_fn = lambda x,y: (self.transformer(x), torch.from_numpy(y).float())
		self.image_set = dataset_df[dataset_df['set']==set]
		if self.segmentation:
			self.targets='target'
			self.image_set[self.targets] = 1.
		if not self.segmentation and fix_names:
			self.image_set.loc[:,'ID'] = self.image_set['ID'].map(fix_name)
		self.slide_info = pd.DataFrame(self.image_set.set_index('ID').loc[:,self.targets])
		if self.mt_bce and not self.segmentation:
			if pos_annotation_class:
				self.targets = [pos_annotation_class]+list(other_annotations)
			else:
				self.targets = None
		print(self.targets)
		IDs = self.slide_info.index.tolist()
		pi_dict=dict(input_info_db=patch_info_file,
					slide_labels=self.slide_info,
					pos_annotation_class=pos_annotation_class,
					patch_size=patch_size,
					segmentation=self.segmentation,
					other_annotations=other_annotations,
					target_segmentation_class=target_segmentation_class,
					target_threshold=target_threshold,
					classify_annotations=classify_annotations,
					modify_patches=modify_patches)
		self.patch_info = modify_patch_info(**pi_dict)

		if self.segmentation and original_set!='pass':
			#IDs = self.patch_info['ID'].unique()
			self.segmentation_maps = {slide:npy2da(join(input_dir,'{}_mask.npy'.format(slide))) for slide in IDs}
		self.slides = {slide:load_preprocessed_img(join(input_dir,'{}.zarr'.format(slide))) for slide in IDs}
		#print(self.slide_info)
		if original_set =='pass':
			self.segmentation=False
		#print(self.patch_info[self.targets].unique())
		if oversampling_factor > 1:
			self.patch_info = pd.concat([self.patch_info]*int(oversampling_factor),axis=0).reset_index(drop=True)
		elif oversampling_factor < 1:
			self.patch_info = self.patch_info.sample(frac=oversampling_factor).reset_index(drop=True)
		self.length = self.patch_info.shape[0]
		self.n_segmentation_classes = n_segmentation_classes
		self.gdl=gdl if self.segmentation else False
		self.binarized=False
		self.classify_annotations=classify_annotations
		print(self.targets)
		self.dilation_jitter=DilationJitter(dilation_jitter,self.segmentation,(original_set=='train'))
		if not self.targets:
			self.targets = [pos_annotation_class]+list(other_annotations)

	def concat(self, other_dataset):
		"""Concatenate this dataset with others. Updates its own internal attributes.

		Parameters
		----------
		other_dataset:DynamicImageDataset
			Other image dataset.

		"""
		self.patch_info = pd.concat([self.patch_info, other_dataset.patch_info],axis=0).reset_index(drop=True)
		self.length = self.patch_info.shape[0]
		if self.segmentation:
			self.segmentation_maps.update(other_dataset.segmentation_maps)
			#print(self.segmentation_maps.keys())

	def retain_ID(self, ID):
		"""Reduce the sample set to just images from one ID.

		Parameters
		----------
		ID:str
			Basename/ID to predict on.

		Returns
		-------
		self

		"""
		self.patch_info=self.patch_info.loc[self.patch_info['ID']==ID]
		self.length = self.patch_info.shape[0]
		self.segmentation_maps={ID:self.segmentation_maps[ID]}
		return self

	def split_by_ID(self):
		"""Generator similar to groupby, but splits up by ID, generates (ID,data) using retain_ID.

		Returns
		-------
		generator
			ID, DynamicDataset

		"""
		for ID in self.patch_info['ID'].unique():
			new_dataset = copy.deepcopy(self)
			yield ID, new_dataset.retain_ID(ID)

	def select_IDs(self, IDs):
		for ID in IDs:
			if ID in self.patch_info['ID'].unique():
				new_dataset = copy.deepcopy(self)
				yield ID, new_dataset.retain_ID(ID)


	def get_class_weights(self, i=0):#[0,1]
		"""Weight loss function with weights inversely proportional to the class appearence.

		Parameters
		----------
		i:int
			If multi-target, class used for weighting.

		Returns
		-------
		self
			Dataset.

		"""
		if self.segmentation:
			label_counts=self.patch_info[list(map(str,list(range(self.n_segmentation_classes))))].sum(axis=0).values
			freq = label_counts/sum(label_counts)
			weights=1./(freq)
		elif self.mt_bce:
			weights=1./(self.patch_info.loc[:,self.targets].sum(axis=0).values)
			weights=weights/sum(weights)
		else:
			if self.binarized and len(self.targets)>1:
				y=np.argmax(self.patch_info.loc[:,self.targets].values,axis=1)
			elif (type(self.targets)==type('')):
				y=self.patch_info.loc[:,self.targets]
			else:
				y=self.patch_info.loc[:,self.targets[i]]
			y=y.values.astype(int).flatten()
			weights=compute_class_weight(class_weight='balanced',classes=np.unique(y),y=y)
		return weights

	def binarize_annotations(self, binarizer=None, num_targets=1, binary_threshold=0.):
		"""Label binarize some annotations or threshold them if classifying slide annotations.

		Parameters
		----------
		binarizer:LabelBinarizer
			Binarizes the labels of a column(s)
		num_targets:int
			Number of desired targets to preidict on.
		binary_threshold:float
			Amount of annotation in patch before positive annotation.

		Returns
		-------
		binarizer

		"""

		annotations = self.patch_info['annotation']
		annots=[annot for annot in list(self.patch_info.iloc[:,6:]) if annot !='area']
		if not self.mt_bce and num_targets > 1:
			if binarizer == None:
				self.binarizer = LabelBinarizer().fit(annotations)
			else:
				self.binarizer = copy.deepcopy(binarizer)
			self.targets = self.binarizer.classes_
			annotation_labels = pd.DataFrame(self.binarizer.transform(annotations),index=self.patch_info.index,columns=self.targets).astype(float)
			for col in list(annotation_labels):
				if col in list(self.patch_info):
					self.patch_info.loc[:,col]=annotation_labels[col].values
				else:
					self.patch_info[col]=annotation_labels[col].values
		else:
			self.binarizer=None
			self.targets=annots
			if num_targets == 1:
				self.targets = [self.targets[-1]]
			if binary_threshold>0.:
				self.patch_info.loc[:,self.targets]=(self.patch_info[self.targets]>=binary_threshold).values.astype(np.float32)
			print(self.targets)
			#self.patch_info = pd.concat([self.patch_info,annotation_labels],axis=1)
		self.binarized=True
		return self.binarizer

	def subsample(self, p):
		"""Sample subset of dataset.

		Parameters
		----------
		p:float
			Fraction to subsample.

		"""
		np.random.seed(42)
		self.patch_info = self.patch_info.sample(frac=p)
		self.length = self.patch_info.shape[0]

	def update_dataset(self, input_dir, new_db, prediction_basename=[]):
		"""Experimental. Only use for segmentation for now."""
		self.input_dir=input_dir
		self.patch_info=load_sql_df(new_db, self.patch_size)
		IDs = self.patch_info['ID'].unique()
		self.slides = {slide:load_preprocessed_img(join(self.input_dir,'{}.zarr'.format(slide))) for slide in IDs}
		if self.segmentation:
			if prediction_basename:
				self.segmentation_maps = {slide:npy2da(join(self.input_dir,'{}_mask.npy'.format(slide))) for slide in IDs if slide in prediction_basename}
			else:
				self.segmentation_maps = {slide:npy2da(join(self.input_dir,'{}_mask.npy'.format(slide))) for slide in IDs}
		self.length = self.patch_info.shape[0]

	#@pysnooper.snoop("getitem.log")
	def __getitem__(self, i):
		patch_info = self.patch_info.iloc[i]
		ID = patch_info['ID']
		xs = patch_info['x']
		ys = patch_info['y']
		patch_size = patch_info['patch_size']
		if xs==np.nan:
			entire_image=True
		else:
			entire_image=False
		targets=self.targets
		use_long=False
		if not self.segmentation:
			y = patch_info.loc[list(self.targets) if not isinstance(self.targets,str) else self.targets]
			if isinstance(y,pd.Series):
				y=y.values.astype(float)
				if self.binarized and not self.mt_bce and len(y)>1:
					y=np.array(y.argmax())
					use_long=True
			y=np.array(y)
			if not y.shape:
				y=y.reshape(1)
		if self.segmentation:
			arr=self.segmentation_maps[ID]
			if not entire_image:
				arr=arr[xs:xs+patch_size,ys:ys+patch_size]
			arr=self.dilation_jitter(arr)
		y=(y if not self.segmentation else np.array(arr))
		#print(y)
		arr=self.slides[ID]
		if not entire_image:
			arr=arr[xs:xs+patch_size,ys:ys+patch_size,:3]
		image, y = self.transform_fn(arr.compute().astype(np.uint8), y)#.unsqueeze(0) # transpose .transpose([1,0,2])
		if not self.segmentation and not self.mt_bce and self.classify_annotations and use_long:
			y=y.long()
		#image_size=image.size()
		if self.gdl:
			y=class2one_hot(y, self.n_segmentation_classes)
		#	y=one_hot2dist(y)
		return image, y

	def __len__(self):
		return self.length

class NPYDataset(Dataset):
	def __init__(self, patch_info, patch_size, npy_file, transform, mmap=False):
		self.ID=os.path.basename(npy_file).split('.')[0]
		patch_info=patch_info=load_sql_df(patch_info,patch_size)
		self.patch_info=patch_info.loc[patch_info["ID"]==self.ID].reset_index()
		self.X=np.load(npy_file,mmap_mode=(None if not mmap else 'r+'))
		self.transform=transform

	def __getitem__(self,i):
		x,y,patch_size=self.patch_info.loc[i,["x","y","patch_size"]]
		return self.transform(self.X[x:x+patch_size,y:y+patch_size])

	def __len__(self):
		return self.patch_info.shape[0]

	def embed(self,model,batch_size,out_dir):
		Z=[]
		dataloader=DataLoader(self,batch_size=batch_size,shuffle=False)
		n_batches=len(self)//batch_size
		with torch.no_grad():
			for i,X in enumerate(dataloader):
				if torch.cuda.is_available():
					X=X.cuda()
				z=model(X).detach().cpu().numpy()
				Z.append(z)
				print(f"Processed batch {i}/{n_batches}")
		Z=np.vstack(Z)
		torch.save(dict(embeddings=Z,patch_info=self.patch_info),os.path.join(out_dir,f"{self.ID}.pkl"))
		print("Embeddings saved")
		quit()