""" Dataset classes for MBZUAI- BiomedIA Fetal Ultra Sound datasets
"""

import os, sys
import json, glob
import random
import PIL.Image
import h5py
import pandas as pd
import numpy as np

import torch
from torch.utils.data import Dataset, WeightedRandomSampler
import torchvision.transforms as torch_transforms
from typing import List, Dict

##---------------------- Generals -----------------------------------------------

def filter_dataframe(self, df, filtering_dict):
    """ Usage:
    {"blacklist":{'class':["4ch"],"machine_type":["Voluson E8","Voluson S10 Expert","V830"]}}
    """

    if "blacklist" in filtering_dict and "whitelist" in filtering_dict:
        raise  Exception("Hey, decide between whitelisting or blacklisting,"+\
                         "Can't do both! remove either one")

    if "blacklist" in filtering_dict:
        print("blacklisting...")
        blacklist_dict = filtering_dict["blacklist"]
        new_df = df
        for k in blacklist_dict.keys():
            for val in blacklist_dict[k]:
                new_df = new_df[new_df[k] != val]

    elif "whitelist" in filtering_dict:
        print("whitelisting...")
        whitelist_dict = filtering_dict["whitelist"]
        new_df_list = []
        for k in whitelist_dict.keys():
            for val in whitelist_dict[k]:
                new_df_list.append(df[df[k] == val])
        new_df = pd.concat(new_df_list).drop_duplicates().reset_index(drop=True)

    else:
        print("No filtering of data done, Peace!")
        new_df = df

    return new_df


def get_class_weights(targets, nclasses):
    """
    Sample level weights fro balanced Loss statergy or data sampling
    targets: assumed to be Long ints representing class from dataset
    """

    n_target = len(targets)
    count_per_class = np.zeros(nclasses, dtype=int)
    for c in targets:
        count_per_class[c] += 1
    count_per_class[count_per_class==0] = n_target

    # for passing to Loss funcs
    weight_per_class = np.zeros(nclasses, dtype=float)
    for i in range(nclasses):
        weight_per_class[i] = float(n_target) / float(count_per_class[i])

    # for passing to sampler
    weight_samplewise = np.zeros(n_target, dtype=float)
    for idx, tgt in enumerate(targets):
        weight_samplewise[idx] = weight_per_class[tgt]

    return weight_per_class, weight_samplewise



## =============================================================================
## Classification


class ClassifyDataFromCSV(Dataset):
    def __init__(self, root_folder, csv_path, transform = None,
                        filtering_dict: Dict[str,Dict[str,List]] = None,
                        ):
        """
        """
        self.root_folder = root_folder
        self.df = pd.read_csv(csv_path)

        ## Filter based on some condition in dataframes
        if filtering_dict: self.df = filter_dataframe(self.df, filtering_dict)

        self.class_to_idx ={c:i for i, c in enumerate(sorted(set(
                                    self.df["class"])))}
        self.images_path =  [ os.path.join(root_folder, p)
                             for p in self.df["image_path"] ]
        self.targets =list(map(lambda x: self.class_to_idx[x],
                                    list(self.df["class"]) ))

        if transform: self.transform = transform
        else: self.transform = torch_transforms.ToTensor()

        print("Class Indexing:", self.class_to_idx)

    def __len__(self):
        return len(self.images_path)

    def __getitem__(self, index):
        imgpath = self.images_path[index]
        target = self.targets[index]
        image = PIL.Image.open(imgpath).convert("RGB")
        image = self.transform(image)
        return image, target





##================ US Video Frames Loader ======================================


class FetalUSFramesDataset(torch.utils.data.Dataset):
    """ Treats Video frames as Independant images for trainng purposes
    """
    def __init__(self, images_folder=None, hdf5_file=None,
                        transform = None,
                        load2ram = False, frame_skip=None):
        """
        """
        self.load2ram = load2ram
        self.frame_skip = frame_skip
        #tobedefined
        self.image_paths= []
        self.image_frames= []
        self.get_image_func = None
        ##-----

        if transform: self.transform = transform
        else: self.transform = torch_transforms.ToTensor()

        if hdf5_file:       self._hdf5file_handler(hdf5_file)
        elif images_folder:  self._imagefolder_handler(images_folder)
        else: raise Exception("No Data info to load")


    # for image folder handling
    def _imagefolder_handler(self, images_folder):
        def __get_image_lazy(index):
            return PIL.Image.open(self.image_paths[index]).convert("RGB")
        def __get_image_eager(index):
            return self.image_frames[index]

        self.image_paths = sorted(glob.glob(images_folder+"/**/*.png"))

        self.get_image_func = __get_image_lazy
        if self.load2ram:
            self.image_frames = [ __get_image_lazy(i)
                                    for i in range(len(self.image_paths))]
            self.get_image_func = __get_image_eager

        print("Frame Skip is not implemented")

    # for hdf5 file handling
    def _hdf5file_handler(self, hdf5_file):
        def __get_image_lazy(index):
            k, i = self.image_paths[index]
            arr = self.hdfobj[k][i]
            return PIL.Image.fromarray(arr).convert("RGB")

        def __get_image_eager(index):
            return self.image_frames[index]

        self.hdfobj = h5py.File(hdf5_file,'r')
        for k in self.hdfobj.keys():
            for i in range(self.hdfobj[k].shape[0]):
                if i % self.frame_skip: continue
                self.image_paths.append([k, i])

        self.get_image_func = __get_image_lazy
        if self.load2ram:
            self.image_frames = [ __get_image_lazy(i)
                                    for i in range(len(self.image_paths))]
            self.get_image_func = __get_image_eager



    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, index):
        image = self.get_image_func(index)
        image = self.transform(image)
        return image

    def get_info(self):
        print(self.get_image_func)
        return {
            "DataSize": self.__len__(),
            "Transforms": str(self.transform),
        }