#!/usr/bin/env python

import numpy as np
import cPickle
import random
import colorsys

# A bit of information about Numpy

# - array[i]: indexing
# - array[i,j]: 2D indexing
# - array[i][j]: 2D indexing (other form)
# - array[i0:i1]: range indexing
# - array[i0:i1, j0:j1]: 2D range indexing

# - np.mean(array): computes the mean of the array
# - np.std(array): computes the standard deviation of the array
# - np.var(array): computes the variance of the array

# - array1 - array2: element by element subtraction
# - array1 * array2: element by element multiplication (not the matrix mult!)
# - array1 * v: multiply each element by a scalar

# ==============================================================================
# =                    Some useful functions for you                           =
# ==============================================================================

def flatten_band(b):
    return b.reshape(b.shape[0] * b.shape[1])
    
def flatten_sort_band(b):
    return np.sort(b.reshape(b.shape[0] * b.shape[1]))

def rgb_to_hsv(r, g, b):
    shape = r.shape
    # reshape input arrays
    R = r.reshape(shape[0] * shape[1])
    G = g.reshape(shape[0] * shape[1])
    B = b.reshape(shape[0] * shape[1])
    # build output arrays
    h = np.zeros(shape[0] * shape[1])
    s = np.zeros(shape[0] * shape[1])
    v = np.zeros(shape[0] * shape[1])
    # fill the output arrays
    for i in xrange(shape[0] * shape[1]):
        h[i], s[i], v[i] = colorsys.rgb_to_hsv(R[i]/255.0, G[i]/255.0, B[i]/255.0)
    # restore the shapes
    h = h.reshape(shape)
    s = s.reshape(shape)
    v = v.reshape(shape)
    return h, s, v

# ==============================================================================
# =                  Functions to extract the features                         =
# =                     GUYS, DO EDIT THIS STUFF!!!                            =
# ==============================================================================

# initial functions
def get_avg_rgb(r, g, b, l):
    return np.mean(r + g + b - 3 * l)
    
def get_std_green(r, g, b, l):
    return np.std(g)
    
def get_var_green(r, g, b, l):
    return np.var(g)
    
def get_var2_green(r, g, b, l):
    m = np.mean(g)
    return np.mean((g - m)**2)
    
def get_avg_red_nooutliers(r, g, b, l):
    r_nooutliers = flatten_sort_band(r)[400:-400]
    return np.mean(r_nooutliers)
    
def get_rquadavg(r, g, b, l):
    return np.mean(r[100:,100:])

# ==============================================================================
# =                           Feature list                                     =
# =                     GUYS, DO EDIT THIS STUFF!!!                            =
# ==============================================================================

# initial state
feature_list = [
    ('get_avg_rgb', get_avg_rgb),
    ('get_std_green', get_std_green),
    ('get_var_green', get_var_green),
    ('get_var2_green', get_var2_green),
    ('get_avg_red_nooutliers', get_avg_red_nooutliers),
    ('get_rquadavg', get_rquadavg),
]

# ==============================================================================
# =                         Build the ARFF file                                =
# =                   *** DO NOT TOUCH THIS STUFF ***                          =
# ==============================================================================

def compute_features(raw_data):
    '''Compute the features'''
    # compute the features from the data
    fdata = []
    img_labels = []
    img_fnames = []
    for img_name, r, g, b, l, img_label in raw_data:
        # convert all values to floats
        r = r.astype(float)
        g = g.astype(float)
        b = b.astype(float)
        l = l.astype(float)
        # compute all features
        fdata_tuple = [fun(r,g,b,l) for _, fun in feature_list]
        # store the feature values
        fdata.append(fdata_tuple)
        # store the label
        img_labels.append(img_label)
        # store the file name
        img_fnames.append(img_name)
        
    # turn the feature data into a numpy array
    fdata = np.array(fdata)
    
    # compute the mean and the standard deviation on a per-column basis
    fdata_mean = fdata.mean(axis = 0)
    fdata_std = fdata.std(axis = 0)
    
    # normalize the data
    fdata = (fdata.astype(float) - fdata_mean) / fdata_std
    
    # return the data
    return [[v for v in dt] + [img_labels[i], img_fnames[i]]
                              for i, dt in enumerate(fdata)]


def make_arff(fdata, arff_fname):
    '''Build the training set and write everything to an ARFF file'''    
    # change the data to ARFF format
    with open(arff_fname, 'w') as outf:
        # write the relation name
        outf.write('@relation is_it_the_sea\n')
        outf.write('\n')
        
        # write the attribute headers
        for feat_name, _ in feature_list:
            outf.write('@attribute %s real\n' % feat_name)
        outf.write('@attribute class {sea, not_sea}\n')
        outf.write('\n')
        
        # write the data
        outf.write('@data\n')
        for i, dt in enumerate(fdata):
            # write the feature data
            outf.write(','.join(['%f' % v for v in dt[:-2]]))
            # write the class label
            outf.write(',' + dt[-2])
            # write the image name
            outf.write('\t% ' + dt[-1] + '.png\n')


if __name__ == '__main__':
    # load the dataset
    with open('all_data.dat') as inf:
        hdr, all_data = cPickle.load(inf)
    
    # compute the features
    all_feat = compute_features(all_data)
        
    # shuffle data
    random.seed(100) # ...but in a predictable fashion!
    random.shuffle(all_feat)
        
    # separate data based on the label
    sea_feat = []
    not_sea_feat = []
    for dt in all_feat:
        if dt[-1] == 'sea':
            sea_feat.append(dt)
        else:
            not_sea_feat.append(dt)
            
    # separate training set and validation set data
    ratio = 0.20
    sea_th = int(len(sea_feat)*ratio)
    not_sea_th = int(len(not_sea_feat)*ratio)
    vs_sea = sea_feat[:sea_th]
    ts_sea = sea_feat[sea_th:]
    vs_not_sea = not_sea_feat[:not_sea_th]
    ts_not_sea = not_sea_feat[not_sea_th:]
    
    # merge the sets
    ts_data = ts_sea + ts_not_sea
    vs_data = vs_sea + vs_not_sea
    
    # shuffle again the data
    random.seed(100) # ...but in a predictable fashion!
    random.shuffle(ts_data)
    random.shuffle(vs_data)
    
    # write the corresponding ARFF files
    make_arff(ts_data, 'ts.arff')
    make_arff(vs_data, 'vs.arff')