import numpy as np
import os
import h5py
import glob
import scipy.misc
import json
import sqlite3 
import tensorflow as tf

"""
Written by Larry Du for Insight Data Science 2017 Remote Program

"""


class ImagesToHdf5:
    def __init__(self,image_dir,output_file,img_dims=[96,96,4],drop_channel=None,num_classes=1,fext='.png'):
        self.image_dir = os.path.abspath(image_dir)
        self.output_file=output_file

        self.drop_channel = drop_channel #0=r,1=g, 2=b, 3=alpha

        
        if self.drop_channel != None:
            #Change image dims removing one channel
            self.img_dims = img_dims[:-1]+[img_dims[-1]-1]
        else:
            self.img_dims=img_dims

        
        self.num_classes=1
        
        self.fext = fext
        self.image_files = glob.glob(self.image_dir+os.sep+'*'+self.fext)
        self.num_examples = len(self.image_files)

        
    def create_hdf5(self):
        with h5py.File(self.output_file,'w') as of:
            of.attrs['num_classes'] = self.num_classes

            num_rows= len(self.image_files)
            num_cols = np.prod(self.img_dims) + 1
            dset = of.create_dataset('toons',
                                     (num_rows,num_cols),
                                     chunks=(64,num_cols),
                                     compression=None,
                                     maxshape = (1000000,num_cols))
            
            
            for i,img in enumerate(self.image_files):
                img_arr = scipy.misc.imread(img) #[96,96,4] RGBA
                if self.drop_channel != None:
                    img_arr = np.delete(img_arr,self.drop_channel,axis=2)

                
                flat_arr = np.reshape(img_arr,num_cols-1)
                flat_arr = np.concatenate((flat_arr,[0]),axis=0) #Last position reserved for label
                dset[i,:] = flat_arr
            print "Finished writing file",self.output_file


       
                
class JsonToonParams:
    def __init__(self,json_file):
        self.filename = json_file
        self.path = os.path.dirname(os.path.abspath(self.filename))
        print "Parsing json file", self.filename
        self.data_key ="toons"
        with open(self.filename,"r") as jf:
            data = json.load(jf)
            self.image_dir = os.path.abspath(data["image_dir"])

            self.save_dir = self.path+os.sep+data['save_dir']
            self.num_epochs = data["num_epochs"]
            self.batch_size = data["batch_size"]
            self.learning_rate=np.float32(data["learning_rate"])
            self.beta1 = np.float32(data["beta1"])



            """DEPRECATED: This is only for training from an HDF5 file"""
            if data['training_file'] != None or data['training_file'] !='':
                try:
                    self.training_file = data['training_file']
                    with h5py.File(self.training_file,'r') as rf:
                        dset = rf[self.data_key]
                        self.num_examples = dset.shape[0]
                        self.record_width = dset.shape[1]
                        #self.num_classes = dset.attrs["num_classes"]
                except:
                    print "Could not load HDF file for parsing num_examples and record width!"

class ToonReaderInput:
    def __init__(self,toon_params,fext='.png'):
        self.params = toon_params
        self.images_dir = self.params.images_dir
        self.fext=fext
        self.image_files = []
        for im_name in glob.glob(images_dir+os.sep+'*'+fext):
            self.image_files.append(im_name)
            

        

        
        
        #self.input_range = input_range
        #self.min_max= self.input_range[1]-self.input_range[0]


                

class PokemonInput:
    """
    Create a feed of pokemon images and labels
    Args:
    	image_dir: directory of images of pokemon with names indexed by id # 
    		e.g. '1.png' is bulbasaur'
    	poke_sqlite: sqlite database file with a 'morphology' column in the 
    		'pokemon' table
    	num_epochs: The number of epochs to specify to tf.train.string_input_producer
    	img_resize [64]: Size in pixels of height and width to resize image to

    Note: It is important that the user does not use tf.nn.sparse_softmax_cross_entropy_with_logits
    with labels generated by this class, as many labels are not mutually exclusive.

    Second note: I (Larry) visually labeled each image as belonging to 10 different 
    non-mutually exclusive classes that do not correspond to pokemon types from any of 
    the franchise games. I did this because for image recognition and generation purposes,
    the official "type classes" do not correspond well with the actual images used.
    Some images in the training set were also flipped or rotated to help the algorithm.

    """
        
    def __init__(self,image_dir,
                 poke_sqlite,
                 num_epochs,
                 image_resize=64,
                 identity_key='id',
                 feature_key="morphology",
                 make_background=None
                 ):

        """
        The order of classes in the following list will define our labels.
        Note that many images belong to multiple categories. 
        This type of problem is also called 'multilabel' classification.
        """
        self.make_background = make_background

        self.feature_key = feature_key
        morphs = ['quadraped','biped','winged','bird','arthropod','serpentine',
                       'aquatic','fish','plant','misc']
        ex_morphs = ['saurian','butterfly','carnivora','humanoid','bird','maus','plant','ungulate','seal','fish','serpent']

        feat_dict = {'morphology':morphs,'exclusive_morph':ex_morphs}
        
        
        #Cast morphs to unicode
        self.morphs = [unicode(m) for m in morphs]
        self.exclusive_morphs = [unicode(e) for e in ex_morphs]
        self.num_classes = len(feat_dict[self.feature_key])
        self.num_epochs=num_epochs
        #self.num_channels=num_channels
        self.identity_key = identity_key
        
        self.image_dir = os.path.abspath(image_dir)
        self.image_resize=image_resize
        self.plite_file = os.path.abspath(poke_sqlite)
        print "Getting images from",self.image_dir
        print "Getting sqlite database file from",self.plite_file
        
        self.pcon = sqlite3.connect(self.plite_file)
        self.pcur=self.pcon.cursor()

            #Select by id number rather than name
            
        id_morph = self.pcur.execute("SELECT "+self.identity_key+","+self.feature_key+
                                         " FROM pokemon WHERE "+self.feature_key+" IS NOT null")
        self.label_dict ={} #key:'bulbasaur',label: [0.5,0,0,0,0,0,0,0,0.5,0]
                            #(bulbasaur is a quadraped,plant, with image 1.png)
        
                            
        for entry in id_morph:
            #Check image_dir for presence of image
            if self.identity_key=='id':
                image_files = [self.image_dir+os.sep+str(entry[0])+'.png']#just one file
            elif self.identity_key=='identifier':#Select picture by name/identifier (e.g. bulbasaur.png)
                image_files = glob.glob(self.image_dir+
                                        os.sep+str(entry[0])+'*.png') #select bulbasaur*.png
            else:
                print "Identity key",self.identity_key,"invalid"
            for image_file in image_files:
                if os.path.isfile(image_file):
                    label = self.class_to_multilabel(entry[1],image_file,feat_dict[feature_key])
                    self.label_dict[image_file] = label
                #print "Image:\t",image_file,"\tlabel:\t",label
        self.num_examples = len(self.label_dict)
        print "Number of examples",self.num_examples
        #Just queue all the examples if the dataset isn't big
        self.perm_indices = np.random.permutation(range(0,self.num_examples))
        self.epoch_tracker = EpochTracker(self.num_examples)
        

        self.image_list = sorted(self.label_dict.keys())
        self.label_list = [self.label_dict[image_file] for image_file in self.image_list]
        self.label_array = np.concatenate(self.label_list,axis=0)
        
        
        #Convert lists to tensors to go into tf.train.slice_input_producer
        self.image_tensors = tf.convert_to_tensor(self.image_list, dtype=tf.string)
        #print "Image tensor shape",self.image_tensors.get_shape().as_list()
        self.label_tensors = tf.convert_to_tensor(self.label_array)
        #print "Label tensor shape",self.image_tensors.get_shape().as_list()
                
        

    def next_batch(self,batch_size,tanh_scale=False):
        do_reset = self.epoch_tracker.increment(batch_size)
        if do_reset:
            #Reset when next pull goes over the limit for current batch size
            self.perm_indices = np.random.permutation(range(self.num_examples))

        batch_start = self.epoch_tracker.cur_index
        batch_end = batch_start+batch_size

        batch_indices = self.perm_indices[batch_start:batch_end]
        #Note, must draw data via shuffled indices

        data = []
        labels =[]
        for bi in batch_indices:
            
            pulled_image = scipy.misc.imread(self.image_list[bi])

            if self.make_background == 'noise':
                pulled_image = self.alpha_to_noise(pulled_image)
            elif self.make_background == 'white':
                pulled_image = self.alpha_to_color(pulled_image,color=(255,255,255))
            
            
            if pulled_image.shape[2]>3:
                pulled_image= pulled_image[:,:,0:3]

            if pulled_image.shape[1] != self.image_resize:
                pulled_image = scipy.misc.imresize(pulled_image,
                                                   (self.image_resize,self.image_resize,
                                                    3),
                                                   interp='nearest')

            #Rescale pixel values to range [-1,1]
            if tanh_scale == True:
                norm_img = [(pulled_image/127.5) - 1]
                data.append(norm_img)
            else:
                #Rescale to [0,1]
                norm_img = [(pulled_image/255.)]
                data.append(norm_img)

            labels.append(self.label_list[bi])

        all_data = np.concatenate(data,axis=0)
        all_labels = np.concatenate(labels,axis=0)

        return all_data,all_labels

        
        


    def alpha_to_color(self,image, color=(255, 255, 255)):
        """Set all fully transparent pixels of an RGBA image to the specified color.
        This is a very simple solution that might leave over some ugly edges, due
        to semi-transparent areas. You should use alpha_composite_with color instead.

        Source: http://stackoverflow.com/a/9166671/284318

        Keyword Arguments:
        image -- PIL RGBA Image object
        color -- Tuple r, g, b (default 255, 255, 255)

        """
        x = np.array(image)
        r, g, b, a = np.rollaxis(x, axis=-1)
        r[a == 0] = color[0]
        g[a == 0] = color[1]
        b[a == 0] = color[2]
        return np.dstack([r, g, b])


    
    def alpha_to_noise(self,image):
        """Set all fully transparent pixels of an RGBA image to random noise.

        Source: http://stackoverflow.com/a/9166671/284318

        Keyword Arguments:
        image -- PIL RGBA Image object
        color -- Tuple r, g, b (default 255, 255, 255)

        """
        x = np.array(image)
        ih = image.shape[0]
        iw = image.shape[1]
        r, g, b, a = np.rollaxis(x, axis=-1)
        mask = (a==0)*1.
        mask_r = mask*np.random.uniform(0,255.,(ih,iw))
        mask_g = mask*np.random.uniform(0,255.,(ih,iw))
        mask_b = mask*np.random.uniform(0,255.,(ih,iw))
        r = r+ mask_r
        g = g+ mask_g
        b = b+ mask_b
        return np.dstack([r, g, b])

    
        
    def class_to_multilabel(self,descriptors,image_file,desc_list):
        """
        Go from 'plant,biped' ---> [0.5,0,0,0,0,0,0,0,0.5,0]
        """
        
        label = np.zeros((1,self.num_classes),dtype=np.float32)
        dlist = descriptors.strip().split(',')
        #print dlist



        for desc in dlist:
            try:
                ind = desc_list.index(desc)
                label[0,ind]=1.
            except:
                "Error, morphology decriptor",desc,"not valid for",image_file

        num_cats = np.sum(label)
        if num_cats != 1.:
            print "Problem with description for",image_file
            print descriptors
        return label/num_cats
                
    def get_print_labels(self):
        print self.label_dict
    

    
    
    def _read_from_queue(self):
        """
        THIS DOES NOT SEEM TO WORK FOR SOME REASON
        Pulls entry from the queue
        Reads/decodes an image file from disk and also returns that image's label
        http://stackoverflow.com/questions/34340489/tensorflow-read-images-with-labels
        """
        
        self.queue = tf.train.slice_input_producer([self.image_tensors,self.label_tensors],
                                                    num_epochs=None,
                                                   shuffle=True)
                                                   

        
        #Pull entry from the queue
        image_name = self.queue[0]
        label = self.queue[1]
        
        
        decoded_image = tf.image.decode_png(image_name,channels=4)

        #print "Dims",decoded_image.get_shape().as_list()
        threec_image = tf.slice(decoded_image,[0,0,0],[-1,-1,3])
        #print "Dims",threec_image.get_shape().as_list()
        
        resized_img = tf.image.resize_images(threec_image,
                               [self.image_resize,self.image_resize],
                               method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)

        #Note: the decode_png op decodes to uint8
        return tf.cast(resized_img,tf.float32),label
    
    
    def image_batch_pipeline(self,batch_size):
        examples_per_epoch = int(self.num_examples//batch_size *batch_size)

        
        read_image,read_label = self._read_from_queue()
        
        """
        THIS DOES NOT SEEM TO WORK FOR SOME REASON
        min_after_dequeue = int(np.floor(examples_per_epoch*.75))
        num_threads = 4
        capacity = min_after_dequeue+num_threads*batch_size
        batch_images,batch_labels = tf.train.shuffle_batch(
                                   [read_image,read_label],
                                    batch_size,
                                    capacity = capacity,
                                    min_after_dequeue = min_after_dequeue,
                                    num_threads=4,
                                    enqueue_many=False,
                                    name="why_wont_i_work")
                                   
        """

        batch_images,batch_labels = tf.train.batch([read_image,read_label],
                                                   batch_size=batch_size)

        return batch_images, batch_labels




 
"""
class ToonInput:

    "Class for pull batches of images from hdf5 files"
#DEPRECATED
# Switched to PokemonInput which reads straight from the folder instead of an HDF5 file

    
    def __init__(self,toon_params):

        #Note batches drawn from this class will be scaled [-1,1.]
        self.params = toon_params
        self.filename = self.params.training_file
        self.hdf5_file=self.filename
        self.data_key="toons"
        #self.input_range = input_range
        #self.min_max= self.input_range[1]-self.input_range[0]

        
        
        with h5py.File(self.filename,'r') as rf:
            dataset = rf[self.data_key]
            self.num_examples = dataset.shape[0]
            self.record_width = dataset.shape[1]


        self.perm_indices = np.random.permutation(range(self.num_examples))
        self.data_key="toons"
        self.epoch_tracker = EpochTracker(self.num_examples)
        self.open()
        
    def open(self):
        self.fhandle = h5py.File(self.filename,'r')
        self.data = self.fhandle[self.data_key]
        
    def close(self):
        self.reader.close()
        
    def next_batch(self,bOAatch_size):
        return self.pull_batch(batch_size)
        
    def pull(self,batch_size):
        return self.pull_batch(batch_size)
    
    def pull_batch(self,batch_size):
        #TODO:batch_start and batch_end
        do_reset = self.epoch_tracker.increment(batch_size)
        if do_reset:
            #Reset when next pull goes over the limit for current batch size
            self.perm_indices = np.random.permutation(range(self.num_examples))
            
        batch_start = self.epoch_tracker.cur_index
        batch_end = batch_start+batch_size

        batch_indices = self.perm_indices[batch_start:batch_end]
        #Note, must draw data via shuffled indices

        data = []
        labels =[]
        for bi in batch_indices:
            #Rescale pixel values to range [-1,1]
            norm_img = [(self.data[bi,:-1]/127.5) - 1]
            data.append(norm_img)
            labels.append([self.data[bi,-1]])
        
        all_data = np.concatenate(data,axis=0)
        all_labels = np.concatenate(labels,axis=0)

        return all_data,all_labels
            
"""

class EpochTracker:
    def __init__(self,num_examples):
        #Reminder: this exists as a seperate class to ToonCollection
        #because the epoch tracking index need to be tracked separately during training
        # and evaluation
        self.num_examples = num_examples
        self.num_epochs = 0 #The number of epochs that have been passed
        self.cur_index = 0 #The index position on current epoch

    def increment(self,increment_size):
        #Returns true if end of current epoch
        new_index = self.cur_index + increment_size
        #Reset epoch counter if end of current epoch has been reached.
        if ( new_index+increment_size >= self.num_examples):
            
            self.num_epochs += 1
            self.cur_index = 0
            #Reshuffle indices
            return True
        else:
            self.cur_index = new_index
            return False
