import glob
import os
import torch
from torch.utils.data import Dataset, DataLoader
import numpy as np
import matplotlib.image as mpimg
import pandas as pd
import cv2
from scipy.ndimage import rotate
import random
from PIL import Image
from torchvision import transforms
import math

class FacialKeypointsDataset(Dataset):
    """Face Landmarks dataset."""

    def __init__(self, csv_file, root_dir, transform=None):
        """
        Args:
            csv_file (string): Path to the csv file with annotations.
            root_dir (string): Directory with all the images.
            transform (callable, optional): Optional transform to be applied
                on a sample.
        """
        self.key_pts_frame = pd.read_csv(csv_file)
        self.root_dir = root_dir
        self.transform = transform
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'

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

    def __getitem__(self, idx):
        image_name = os.path.join(self.root_dir,
                                self.key_pts_frame.iloc[idx, 0])
        
        image = mpimg.imread(image_name)
        
        # if image has an alpha color channel, get rid of it
        if(image.shape[2] == 4):
            image = image[:,:,0:3]
        
        key_pts = self.key_pts_frame.iloc[idx, 1:].as_matrix()
        key_pts = key_pts.astype('float').reshape(-1, 2)
        sample = {'image': image, 'keypoints': key_pts}

        if self.transform:
            sample = self.transform(sample)

        return sample
    

    
# tranforms

class Normalize(object):
    """Convert a color image to grayscale and normalize the color range to [0,1]."""        

    def __call__(self, sample):
        image, key_pts = sample['image'], sample['keypoints']
        
        image_copy = np.copy(image)
        key_pts_copy = np.copy(key_pts)

        # convert image to grayscale
        image_copy = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
        
        # scale color range from [0, 255] to [0, 1]
        image_copy=  image_copy/255.0
            
        
        # scale keypoints to be centered around 0 with a range of [-1, 1]
        # mean = 100, sqrt = 50, so, pts should be (pts - 100)/50
        key_pts_copy = (key_pts_copy - 100)/50.0


        return {'image': image_copy, 'keypoints': key_pts_copy}


class Rescale(object):
    """Rescale the image in a sample to a given size.

    Args:
        output_size (tuple or int): Desired output size. If tuple, output is
            matched to output_size. If int, smaller of image edges is matched
            to output_size keeping aspect ratio the same.
    """

    def __init__(self, output_size):
        assert isinstance(output_size, (int, tuple))
        self.output_size = output_size

    def __call__(self, sample):
        image, key_pts = sample['image'], sample['keypoints']

        h, w = image.shape[:2]
        if isinstance(self.output_size, int):
            if h > w:
                new_h, new_w = self.output_size * h / w, self.output_size
            else:
                new_h, new_w = self.output_size, self.output_size * w / h
        else:
            new_h, new_w = self.output_size

        new_h, new_w = int(new_h), int(new_w)

        img = cv2.resize(image, (new_w, new_h))
        
        # scale the pts, too
        key_pts = key_pts * [new_w / w, new_h / h]

        return {'image': img, 'keypoints': key_pts}


class RandomCrop(object):
    """Crop randomly the image in a sample.

    Args:
        output_size (tuple or int): Desired output size. If int, square crop
            is made.
    """

    def __init__(self, output_size):
        assert isinstance(output_size, (int, tuple))
        if isinstance(output_size, int):
            self.output_size = (output_size, output_size)
        else:
            assert len(output_size) == 2
            self.output_size = output_size

    def __call__(self, sample):
        image, key_pts = sample['image'], sample['keypoints']

        h, w = image.shape[:2]
        new_h, new_w = self.output_size

        top = np.random.randint(0, h - new_h)
        left = np.random.randint(0, w - new_w)

        image = image[top: top + new_h,
                      left: left + new_w]

        key_pts = key_pts - [left, top]

        return {'image': image, 'keypoints': key_pts}


class ToTensor(object):
    """Convert ndarrays in sample to Tensors."""

    def __call__(self, sample):
        image, key_pts = sample['image'], sample['keypoints']
         
        # if image has no grayscale color channel, add one
        if(len(image.shape) == 2):
            # add that third color dim
            image = image.reshape(image.shape[0], image.shape[1], 1)
            
        # swap color axis because
        # numpy image: H x W x C
        # torch image: C X H X W
        image = image.transpose((2, 0, 1))
        
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
        img_torch = torch.from_numpy(image)
        img_torch = img_torch.to(device)
        key_pts_torch = torch.from_numpy(key_pts)
        key_pts_torch = key_pts_torch.to(device)
        
        return {'image': img_torch,
                'keypoints': key_pts_torch}
    
class RandomRotation(object):
    """Applies random rotation - either no rotation, angle, or -angle with equal probability."""
    # Be careful, as it is hard to train the model where angle is large. 
    # For simplicity, keep angle at =<5 deg. For generalization, keep angle at ~45 deg.
    def __init__(self, angle=5):
        self.angle = angle
        
    def __call__(self, sample):
        img, key_pts = sample["image"], sample["keypoints"]

        img_copy = np.copy(img)
        key_pts_copy = np.copy(key_pts)

        k = random.choice([0,1,2])
        if k==0:
            return {"image":img_copy, "keypoints":key_pts_copy}
        if k==1:
            img_rot, new_key_pts = self.rot(img_copy, key_pts_copy, angle=self.angle)
            return {"image":img_rot, "keypoints":new_key_pts}
        if k==2:
            img_rot, new_key_pts = self.rot(img_copy, key_pts_copy, angle=-self.angle)
            return {"image":img_rot, "keypoints":new_key_pts}

    def rot(self, image, key_pts, angle):
        img_rot = rotate(image, angle, reshape=False) 
        center = (np.array(image.shape[:2][::-1]))/2
        a = (angle/180) * math.pi
        new_key_pts = np.zeros_like(key_pts)
        for i in range(len(key_pts)):
            key_pt = key_pts[i]
            org = key_pt-center
            new_coords = np.array([org[0]*np.cos(a) + org[1]*np.sin(a),
            -org[0]*np.sin(a) + org[1]*np.cos(a) ])
            new_key_pts[i] = np.round(new_coords + center)
        return img_rot, new_key_pts
    
class RandomFlip(object):
    """Applies random flip - either no flip, horizontal, vertical, or both with equal probability."""
    # Uncomment k==2 and k==3 cases and add 2 and 3 to random.choice in order to allow vertical and horizontal+vertical flips.
    # Be careful, as it is harder to train the model with vertically flipped outputs
    def __call__(self, sample):
        img, key_pts = sample["image"], sample["keypoints"]

        img_copy = np.copy(img)
        key_pts_copy = np.copy(key_pts)

        k = random.choice([0,1])
        if k==0:
            return {"image":img_copy, "keypoints":key_pts_copy}
        if k==1:
            im_flip, new_key_pts = self.flip(img_copy, key_pts_copy, option=1)
            return {"image":im_flip, "keypoints":new_key_pts}
#         if k==2:
#             im_flip, new_key_pts = self.flip(img_copy, key_pts_copy, option=0)
#             return {"image":im_flip, "keypoints":new_key_pts}
#         if k==3:
#             im_flip, new_key_pts = self.flip(img_copy, key_pts_copy, option=-1)
#             return {"image":im_flip, "keypoints":new_key_pts}
        
    def flip(self, img, key_pts, option):
        img_flip = cv2.flip(img, option)
    
        if option==1:
            key_pts[:,0] = img.shape[1] - key_pts[:,0]
#         if option==0:
#             key_pts[:,1] = img.shape[0] - key_pts[:,1]
#         if option==-1:
#             key_pts = np.flip(img.shape[:2], 0) - key_pts

        return img_flip, key_pts

    
class ColorJitter(object):
    """Apply color jitter - random change in brightness, contrast, and saturation in range 0.3."""
    
    def __call__(self, sample):
        img, key_pts = sample['image'], sample['keypoints']
        
        color_jitter = transforms.ColorJitter(
            brightness=0.3,
            contrast=0.3,
            saturation=0.3,)
        
        img_copy = np.copy(img)
        img_jitter = np.array(color_jitter(Image.fromarray(img_copy)))
        
        return {'image': img_jitter, 'keypoints': key_pts}
