#!/usr/bin/python

import sys, os
import cv2
import numpy as np
import scipy.ndimage as nd
import csv
import math
import joblib

from scipy.stats.mstats import mquantiles, kurtosis, skew

def getCurrentFolderLocation():
	return "/home/rkjha/Documents/Copy/CS365/project/code/kaggle-galaxy-master/data_folder_new"

def getNameField():
    return [ 'GalaxyID','Class1.1','Class1.2','Class1.3','Class2.1','Class2.2','Class3.1','Class3.2','Class4.1','Class4.2','Class5.1','Class5.2','Class5.3','Class5.4','Class6.1','Class6.2','Class7.1','Class7.2','Class7.3','Class8.1','Class8.2','Class8.3','Class8.4','Class8.5','Class8.6','Class8.7','Class9.1','Class9.2','Class9.3','Class10.1','Class10.2','Class10.3','Class11.1','Class11.2','Class11.3','Class11.4','Class11.5','Class11.6']

def getNameClasses():
    return np.array(map(lambda x: int(x[len('Class'):-len('.1')] ), getNameField()[1:] ), dtype='int')

def saturateImage( img, q0=0.01, q1=0.99 ):
    img = img.astype('float')
    if q0 is None:
        q0 = 0
    if q1 is None:
        q1 = 1
    q = mquantiles(img[np.nonzero(img)].flatten(),[q0,q1])
    img[img<q[0]] = q[0]
    img[img>q[1]] = q[1]
    return img
    
def boundingBox( img, crop=False ):
    pts = np.transpose(np.nonzero(img>0))
    y_min, x_min = pts.min(axis=0)
    y_max, x_max = pts.max(axis=0)
    if not crop:
        return ( x_min, y_min, x_max, y_max )
    else:
        return img[y_min:y_max+1, x_min:x_max+1]

def makeCompose(m1,m2):
    n1 = np.eye(3,dtype='float32')
    n2 = np.eye(3,dtype='float32')
    n1[:2] = m1
    n2[:2] = m2
    n3 = np.dot(n1,n2)
    return n3[:2]

def rescaleImage( img, min=0, max=255 ):
    img = img.astype('float')
    img -= img.min()
    img /= img.max()
    img = img * (max - min) + min
    return img

def rotateImage( img, (x,y), angle, interpolation=cv2.INTER_LINEAR ):
    cx = float(img.shape[1])/2
    cy = float(img.shape[0])/2    
    translation_mat = np.array([[1,0,cx-x],[0,1,cy-y]],dtype='float32')
    rotation_mat = cv2.getRotationMatrix2D((cx,cy),angle,1.0)
    m = makeCompose(rotation_mat, translation_mat)
    res = cv2.warpAffine(img, m, img.shape[:2],flags=interpolation)
    return res

def recenterImage( img, (x,y), interpolation=cv2.INTER_LINEAR ):
    cx = float(img.shape[1])/2
    cy = float(img.shape[0])/2    
    translation_mat = np.array([[1,0,cx-x],[0,1,cy-y]],dtype='float32')
    res = cv2.warpAffine(img, translation_mat, img.shape[:2],flags=interpolation)
    return res

def gini(x,f=""):
    x = x.flatten()
    n = len(x)
    s = x.sum()
    r = np.argsort(np.argsort(-x)) 
    if s == 0 or n == 0:
        return 1.0
    else:
        return 1.0 - (2.0 * (r*x).sum() + s)/(n*s)

def fit_ellipse( points, factor=1.96, f="" ):
    points = points.astype('float')

    center = points.mean(axis=0)
    points -= center
    U, S, V = np.linalg.svd(points, full_matrices=False)

    S /= np.sqrt(len(points)-1)
    S *= factor
    angle = math.atan2(V[0,1],V[0,0])/math.pi*180
    
    return (center,2*S,angle)

def unique_indices_images(M):
    unique_index = np.unique( M.dot(np.random.rand(M.shape[1])),
                                  return_index=True)[1]
    return unique_index

def get_entropy(img):
    hist = cv2.calcHist([img],[0],None,[256],[0,256])
    hist = hist.ravel()/hist.sum()
    logs = np.log2(hist+0.00001)
    return -1 * (hist*logs).sum()

def gaussian(shape):
    x, y = np.mgrid[-shape[0]//2+1:shape[0]//2+1, -shape[1]//2+1:shape[1]//2+1]
    g = np.exp(-(x**2/(2.0*float(shape[0]//2//2)**2)+y**2/(2.0*float(shape[1]//2//2)**2)))
    return g / g.max()

def get_light_radius_images(img,r=[0.1,0.8]):
    img = img.astype('float')
    idx = np.nonzero(img)
    s = img[idx].sum()
    mask = np.ones(img.shape)
    mask[img.shape[0]/2,img.shape[1]/2] = 0
    edt = nd.distance_transform_edt( mask)
    edt[edt>=img.shape[1]/2] = 0
    edt[img==0] = 0
    q = mquantiles(edt[np.nonzero(edt)].flatten(),r)
    res = []
    for q0 in q:
        res.append( img[edt<q0].sum()/s )
    return res

def getColorImagesHistogram(img_color):
    hist = np.array( [ np.bincount( img_color[:,:,0].flatten(),
                          minlength=256 ),
                       np.bincount( img_color[:,:,1].flatten(),
                          minlength=256 ),
                       np.bincount( img_color[:,:,2].flatten(),
                          minlength=256 ) ], dtype='float').flatten()
    
    norm = np.linalg.norm(hist)
    if norm > 0:
        hist /= norm

    return hist 

def getRandomColorImages(labels):
    idx = np.nonzero(labels)
    nb_labels = labels.max()
    colors = np.random.random_integers(0,255,size=(nb_labels+1,3))
    seg = np.zeros( (labels.shape[0],labels.shape[1],3), dtype='uint8' )
    seg[idx] = colors[labels[idx].astype('int')]
    return seg
    
def largeConComponent( img, labels, nb_labels ):
    sizes =  np.bincount( labels.flatten(),
                          minlength=nb_labels+1 )
    centers = nd.center_of_mass( img, labels, range(1,nb_labels+1) )
    distances = map( lambda (y,x):
                         (img.shape[0]/2-y)**2+(img.shape[1]/2-x)**2,
                     centers )
    distances = [1.0] + distances
    distances = np.array(distances)
    sizes[0] = 0
    sizes[sizes<20] = 0
    sizes = sizes/(distances+0.000001)
    best_label = np.argmax( sizes )
    thresholded = (labels==best_label)*255
    
    return thresholded

def get_ImageFeatures( f, dsift_learn=False, debug=False, tiny_img=False, tiny_grad=False, image_statistics=False, color_histogram=False, orientation_histogram=False, transform=0):
    img_color = cv2.imread( f ).astype('float')
    if transform == 1:
        img_color = img_color[::-1,:].copy()
    if transform == 2:
        img_color = img_color[:,::-1].copy()

    original_img = img_color.copy()
    img_color = cv2.GaussianBlur(img_color,(0,0),2.0)
    
    img = cv2.cvtColor( img_color.astype('uint8'),
                        cv2.COLOR_BGR2GRAY ).astype('float')
    
    t = np.max(np.median(img_color[np.nonzero(img)],axis=0))
    
    img_color[img_color<t] = t
    img_color = rescaleImage(img_color).astype('uint8')
    
    img = cv2.cvtColor( img_color.astype('uint8'),
                        cv2.COLOR_BGR2GRAY )

    saturatedImage = saturateImage(img,q1=0.75)
    labels,nb_maxima = nd.label(saturatedImage==saturatedImage.max(), output='int64')

    thresholded = largeConComponent( img, labels, nb_maxima )
    center = nd.center_of_mass( img, thresholded )

    original_size = img_color.shape[0]*img_color.shape[1]
    original_shape = img_color.shape[:2]

    idx = np.nonzero(original_img.mean(axis=2))
    nonzero = original_img[idx]
    mean = nonzero.mean(axis=0)
    std = nonzero.std(axis=0)
    mean_center = original_img[thresholded>0].mean(axis=0)
    ImageFeatures = np.array( [ mean_center[0],mean_center[1],mean_center[2], mean[0],mean[1],mean[2], std[0],std[1],std[2], kurtosis(nonzero[:,0]), kurtosis(nonzero[:,1]), kurtosis(nonzero[:,2]), skew(nonzero[:,0]), skew(nonzero[:,1]), skew(nonzero[:,2]), gini(nonzero[:,0],f), gini(nonzero[:,1],f), gini(nonzero[:,1],f)], dtype='float')
    
    img_color = recenterImage( img_color, (center[1],center[0]), interpolation=cv2.INTER_LINEAR )

    img = cv2.cvtColor(img_color.astype('uint8'), cv2.COLOR_BGR2GRAY)

    center_offset = np.linalg.norm(np.array([center[0],center[1]],dtype='float')
                                   - np.array(original_shape,dtype='float')/2)
            
    
    thresholded = cv2.adaptiveThreshold(img, maxValue=255, adaptiveMethod=cv2.ADAPTIVE_THRESH_GAUSSIAN_C, thresholdType=cv2.THRESH_BINARY, blockSize=301, C=0)

    labels, nb_labels = nd.label( thresholded, output='int64' )
        
    thresholded = largeConComponent( img, labels, nb_labels )
        
    XY = np.transpose( np.nonzero(thresholded) )[:,::-1]
    ((cx,cy),(w,h),angle) = fit_ellipse(XY,f=f)
    eccentricity = h/w
    
    # rotate image
    img_color = rotateImage( img_color, (cx,cy), angle, interpolation=cv2.INTER_LINEAR )

    thresholded = rotateImage( thresholded, (cx,cy), angle, interpolation=cv2.INTER_NEAREST )

    cx = float(img.shape[1])/2
    cy = float(img.shape[0])/2        
    
    # crop
    img_color = img_color[max(0,int(cy-h/2)):min(img.shape[0],int(cy+h/2+1)),
                          max(0,int(cx-w/2)):min(img.shape[1],int(cx+w/2+1))]
    thresholded = thresholded[max(0,int(cy-h/2)):min(img.shape[0],int(cy+h/2+1)),
                              max(0,int(cx-w/2)):min(img.shape[1],int(cx+w/2+1))]

    color_hist = getColorImagesHistogram(img_color)
    if color_histogram:
        return color_hist

    if orientation_histogram:
        return get_orientation_histogram(img_color)    

    img = cv2.cvtColor( img_color, cv2.COLOR_BGR2GRAY ).astype('float')
    img = rescaleImage(img)

    saturatedImage = saturateImage(img,q1=0.95)
    labels,nb_maxima = nd.label(saturatedImage==saturatedImage.max(), output='int64')
    
    img_thumbnail = cv2.resize( img.astype('uint8'), (64,64),
                                interpolation=cv2.INTER_AREA )

    if tiny_img:
            return img_thumbnail.flatten()

    grad_color = nd.gaussian_gradient_magnitude(img_color,1.0)
    grad_img = rescaleImage(img_color[:,:,0]+img_color[:,:,2])

    grad_thumbnail = cv2.resize( grad_img, (64,64),
                                 interpolation=cv2.INTER_AREA )

    if tiny_grad == True:
        return grad_thumbnail.flatten()

    mirror_grad_img = grad_img[:,::-1]
    chirality = np.sum(np.sqrt( (grad_img - mirror_grad_img)**2 ))/ (grad_img.sum()) 

    size_to_ellipse = float(thresholded.sum()) / (math.pi * w * h / 4)
    box_to_image = float(img.shape[0]*img.shape[1]) / original_size

    idx = np.nonzero(thresholded)
    mean = img_color[idx].mean(axis=0)
    grey_mean = img[idx].mean()
    img_center = img[img.shape[0]/2-img.shape[0]/4:img.shape[0]/2+img.shape[0]/4,
                     img.shape[1]/2-img.shape[1]/4:img.shape[1]/2+img.shape[1]/4]
    img_center_color = img_color[img.shape[0]/2-img.shape[0]/4:img.shape[0]/2+img.shape[0]/4,
                                 img.shape[1]/2-img.shape[1]/4:img.shape[1]/2+img.shape[1]/4]
    center_mean = img_center[np.nonzero(img_center)].mean()
    center_mean_color = img_center_color[np.nonzero(img_center)].mean(axis=0)
    color_ImageFeatures = [
        img_color[img_color.shape[0]/2,
                                 img_color.shape[1]/2,0],
                       img_color[img_color.shape[0]/2,
                                 img_color.shape[1]/2,1],
                       img_color[img_color.shape[0]/2,
                                 img_color.shape[1]/2,2],
                      mean[0],mean[1],mean[2],
                       center_mean_color[0],center_mean_color[1],center_mean_color[2],
                       float(img[img.shape[0]/2,
                                 img.shape[1]/2])/grey_mean,
                       float(center_mean)/grey_mean]

    entropy = get_entropy(img.astype('uint8'))

    light_radius_images = get_light_radius_images(img)
    ImageFeatures = np.append( ImageFeatures, [ eccentricity, w,h, thresholded.sum(), entropy, chirality, size_to_ellipse, box_to_image, center_offset, light_radius_images[0], light_radius_images[1], nb_maxima, kurtosis(img_color[idx][:,0]), kurtosis(img_color[idx][:,1]), kurtosis(img_color[idx][:,2]), skew(img_color[idx][:,0]), skew(img_color[idx][:,1]), skew(img_color[idx][:,2]), gini(img_color[idx][:,0],f), gini(img_color[idx][:,1],f), gini(img_color[idx][:,2],f), kurtosis(img[idx]), skew(img[idx]), gini(img[idx],f)])
    ImageFeatures = np.append( ImageFeatures, color_ImageFeatures )
                           
    m = cv2.moments( thresholded.astype('uint8' ), binaryImage=True )
    hu1 = cv2.HuMoments( m )

    m = cv2.moments( img, binaryImage=False )
    hu2 = cv2.HuMoments( m )
    
    m = cv2.moments( grad_img, binaryImage=False )
    hu3 = cv2.HuMoments( m )    

    hu = np.append( hu1.flatten(), hu2.flatten() )
    hu = np.append( hu.flatten(), hu3.flatten() )
    ImageFeatures = np.append( ImageFeatures, hu.flatten() )

    ImageFeatures = np.append( ImageFeatures, hu )

    if image_statistics:
        return ImageFeatures

    average_prediction = np.zeros( 37, dtype='float' )
    
    if not debug:
        image_statistics = ImageFeatures
        for ImageClass in xrange(1,12):
            scaler = joblib.load(getCurrentFolderLocation()+"/scaler_statistics_Class"+ str(ImageClass)+"_")
            clf = joblib.load(getCurrentFolderLocation()+"/svm_statistics_Class"+ str(ImageClass)+"_")
            ImageFeatures = np.append( ImageFeatures, clf.predict_proba(scaler.transform(image_statistics)))

        average_prediction += ImageFeatures[-37:]
        
        img_thumbnail = img_thumbnail.flatten()
        for ImageClass in xrange(1,12):
            pca = joblib.load(getCurrentFolderLocation()+"/pca_img_Class"+ str(ImageClass)+"_")
            thumbnail_pca = pca.transform(img_thumbnail)

            clf = joblib.load(getCurrentFolderLocation()+"/pca_img_SVM_Class" + str(ImageClass)+"_")
            ImageFeatures = np.append( ImageFeatures,
                                  clf.predict_proba(thumbnail_pca).flatten() )

        average_prediction += ImageFeatures[-37:]
        
        for ImageClass in xrange(1,12):
            pca = joblib.load(getCurrentFolderLocation()+"/pca_color_Class"+ str(ImageClass)+"_")
            hist_pca = pca.transform(color_hist)

            clf = joblib.load(getCurrentFolderLocation()+"/pca_color_SVM_Class" + str(ImageClass)+"_")
            ImageFeatures = np.append( ImageFeatures,
                                  clf.predict_proba(hist_pca).flatten() )            

        average_prediction += ImageFeatures[-37:]
        average_prediction /= 3
        ImageFeatures = np.append( ImageFeatures, average_prediction )
        
    return ImageFeatures

def read_responses( f ):
    reader = csv.DictReader( open( f, 'rb' ) )
    responses = []
    ids = []
    for line in reader:
        ids.append( line['GalaxyID'] )
        r = []
        for k in getNameField():
            if k == 'GalaxyID':
                continue
            r.append( float(line[k]) )
        responses.append( r )
    return np.array(responses,dtype='float32'),np.array(ids)

def read_responses_as_dict( f ):
    reader = csv.DictReader( open( f, 'rb' ) )
    responses = []
    for line in reader:
        responses.append( line )
    return responses

if __name__ == '__main__':
    f = sys.argv[1]

    print get_ImageFeatures(f,debug=False,tiny_img=False)

