import pandas as pd 
import numpy as np
import os
import cv2
import random
import matplotlib.pyplot as plt

import torch
from torch.utils.data import Dataset

import albumentations

F = os.path.join(data_dir, files_train[0])

train0 = pd.read_parquet(F)

class BengaliDataset(Dataset):
    def __init__(self, df, img_height, img_width):
        self.df = df
        self.img_height = img_height
        self.img_width = img_width

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

    def __getitem__(self, idx):
        img = self.df.iloc[idx][0:].values.astype(np.uint8)
        img = img.reshape(self.img_height, self.img_width)
        img = 255 - img
        img = (img*(255.0/img.max())).astype(np.uint8)

        return img

HEIGHT = 137
WIDTH = 236

train0.set_index('image_id', inplace=True)
image = BengaliDataset(train0, img_height=HEIGHT, img_width=WIDTH)

from albumentations import HorizontalFlip
from albumentations.core.transforms_interface import ImageOnlyTransform

class AugMix(ImageOnlyTransform):
    """Augmentations mix to Improve Robustness and Uncertainty.
    Args:
        image (np.ndarray): Raw input image of shape (h, w, c)
        severity (int): Severity of underlying augmentation operators.
        width (int): Width of augmentation chain
        depth (int): Depth of augmentation chain. -1 enables stochastic depth uniformly
          from [1, 3]
        alpha (float): Probability coefficient for Beta and Dirichlet distributions.
        augmentations (list of augmentations): Augmentations that need to mix and perform.
    Targets:
        image
    Image types:
        uint8, float32

    Reference:
    |  https://arxiv.org/abs/1912.02781
    |  https://github.com/google-research/augmix
    """

    def __init__(self, width=2, depth=2, alpha=0.5, augmentations=[HorizontalFlip()], always_apply=False, p=0.5):
        super(AugMix, self).__init__(always_apply, p)
        self.width = width
        self.depth = depth
        self.alpha = alpha
        self.augmentations = augmentations
        self.ws = np.float32(np.random.dirichlet([self.alpha] * self.width))
        self.m = np.float32(np.random.beta(self.alpha, self.alpha))

    def apply_op(self, image, op):
        image = op(image=image)["image"]
        return image

    def apply(self, img, **params):
        mix = np.zeros_like(img)
        for i in range(self.width):
            image_aug = img.copy()

            for _ in range(self.depth):
                op = np.random.choice(self.augmentations)
                image_aug = self.apply_op(image_aug, op)

            mix = np.add(mix, self.ws[i] * image_aug, out=mix, casting="unsafe")

        mixed = (1 - self.m) * img + self.m * mix
        if img.dtype in ["uint8", "uint16", "uint32", "uint64"]:
            mixed = np.clip((mixed), 0, 255).astype(np.uint8)
        return mixed

    def get_transform_init_args_names(self):
        return ("width", "depth", "alpha")

augs = [A.HorizontalFlip(always_apply=True),
        A.Blur(always_apply=True),
        A.ShiftScaleRotate(always_apply=True),
        A.GaussNoise(always_apply=True),
        A.Cutout(always_apply=True),
        A.IAAPiecewiseAffine(always_apply=True)]

transforms_train = albumentations.Compose([
    AugMix(width=3, depth=2, alpha=.4, p=1., augmentations=augs),
])

class Aug_BengaliDataset(Dataset):
    def __init__(self, df, img_height, img_width, transform=None):
        self.df = df
        self.img_height = img_height
        self.img_width = img_width
        self.transform = transform

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

    def __getitem__(self, idx):
        img = self.df.iloc[idx][0:].values.astype(np.uint8)
        img = img.reshape(self.img_height, self.img_width)
        img = 255 - img
        img = (img*(255.0/img.max())).astype(np.uint8)
        
        if self.transform is not None:
            res = self.transform(image=img)
            img = res['image']
        else:
            img = img

        return img

aug_image = Aug_BengaliDataset(train0, img_height=HEIGHT, img_width=WIDTH, transform=transforms_train)