import os
from os.path import join
import numpy as np
from PIL import Image
import cv2
from tqdm import tqdm
from random import choice
from mediaug.image_utils import read_png, save_img
from mediaug.download import get_data_cache
from random import randint
# TODO: Documentation

class DataPoint:

    def __init__(self, img_path, mask_path, _class, _id=None):
        self.img_path = img_path
        self.mask_path = mask_path
        self._class = _class
        if _id is None:
            self.id = img_path.split('.')[0]

    @property
    def img(self):
        return read_png(self.img_path)

    @property
    def pil_img(self):
        return Image.fromarray(read_png(self.img_path))

    @property
    def mask(self):
        return read_png(self.mask_path)

    @property
    def pil_mask(self):
        return Image.fromarray(read_png(self.mask_path))

    def __repr__(self):
        return f'<img_path: {self.img_path}>\n<mask_path: {self.mask_path}>'


class Dataset:
    """Dataset object for managing image augmentation

    Attributes:
        data_path (str): Path to the data directory root
    """

    def __init__(self, data_path=None, classes=None):
        self.data_path = data_path
        if not os.path.exists(data_path) and classes is not None:
            self._create_empty_dataset(classes)
        if not os.path.exists(data_path) and classes is None:
            raise ValueError('No data in path or classes.')
        self._parse(data_path)
    
    def _parse(self, data_path):
        self.data = {}
        categories =  [x for x in os.listdir(data_path) if not x.startswith('.')]
        self.data = {key:[] for key in categories}
        for c in categories:
            cur_dir = join(data_path, c)
            for base_name in os.listdir(join(cur_dir, 'image')):
                name = base_name.split('.')[0]
                dp = DataPoint(join(cur_dir, 'image', base_name),
                                join(cur_dir, 'mask', base_name), c, name)
                self.data[c].append(dp)

    def _create_empty_dataset(self, classes):
        os.mkdir(self.data_path)
        self.data = {key:[] for key in classes}
        for _class in classes:
            os.mkdir(join(self.data_path, _class))
            os.mkdir(join(self.data_path, _class, 'image'))
            os.mkdir(join(self.data_path, _class, 'mask'))


    def add_datapoint(self, dp):
        self.data[dp._class].append(dp)

    
    def random_sample(self):
        _class = choice(self.classes)
        return choice(self.data[_class])


    def add_data(self, img, mask, _class, name):
        img_path = save_img(img, join(self.data_path, _class, 'image', f'{name}.png'))
        mask_path = save_img(mask, join(self.data_path, _class, 'mask', f'{name}.png'))
        self.data[_class].append(DataPoint(img_path, mask_path, _class))


    def get_data(self, _id):
        """ Gets a datapoint by id """
        raise NotImplementedError 


    def get_array(self, num_samples=-1, n_last=False, greyscale=False):
        """ This is of the form:
        (x_train, y_train), (x_test, y_test)
        ex: (num_samples, 32, 32, 3)
        (num_samples, 1)
        """
        images = []
        masks = []
        for c in tqdm(self.classes):
            for dp in tqdm(self.data[c][:num_samples]):
                if greyscale == True:
                    images.append(cv2.cvtColor(dp.img, cv2.COLOR_BGR2GRAY))
                    masks.append(cv2.cvtColor(dp.mask, cv2.COLOR_BGR2GRAY))
                else:
                    images.append(dp.img)
                    masks.append(dp.mask)
        images = np.array(images)
        masks = np.array(masks)
        if n_last:
            images = np.moveaxis(images, 0, -1)
            masks = np.moveaxis(masks, 0, -1)
        return images, masks
    

    @property
    def classes(self):
        return list(self.data.keys())

    @property
    def size(self):
        size = 0
        for c in self.classes:
            size += len(self.data[c])
        return size

    def __getitem__(self, arg):
        return self.data[arg]