from __future__ import print_function, division
import os
import torch
import pandas as pd
from skimage import io, transform
import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils

# Ignore warnings
import warnings

plt.ion()   # interactive mode



class GoogleDatasetFour(Dataset):
    """Contains images and their coordinates"""

    def __init__(self, csv_file, root_dir, transform=None):
        """
        Input:
            csv_file (string): Path to the csv file with annotations.
            root_dir (string): Directory with folders of all the images.
            transform (callable, optional): Optional transform to be applied
                on a sample.
        """
        self.google_frame = pd.read_csv(csv_file)
        self.root_dir = root_dir
        self.transform = transform

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

    def __getitem__(self, idx):
        """
         Returns dictionary of the form
              {
                 '0':image 0.jpg
                 '90':image 90.jpg
                 '180':image 180.jpg
                 '270':image 270.jpg
                 'coordinates':(longitude, latitude)
             }
         for the given index idx.
         """

        if torch.is_tensor(idx):
            idx = idx.tolist()

        folder_name = os.path.join(self.root_dir,
                                self.google_frame.iloc[idx, 0])

        coordinates = self.google_frame.iloc[idx, 1:]
        coordinates = np.array([coordinates])
        angles = ('0', '90', '180','270')
        sample = {'coordinates':coordinates}
        for angle in angles:
            img_name = os.path.join(folder_name, angle+'.jpg')
            image = io.imread(img_name)
            sample.update({angle:image})
            if self.transform:
                sample = self.transform(sample)
        return sample








