# -*- coding: utf-8 -*-
import os
import cv2
from sklearn.ensemble import RandomForestRegressor
import csv_io
import math
import pickle
import numpy as np

def main():
    train = []
    image_files = sorted(os.listdir('./train_part1'))
    #image_files = sorted(os.listdir('../../data/a1'))
    count = 0
    
    print "Iraining Images"
    for img in image_files:
        tmp = cv2.imread('./train_part1/' + img)
        #tmp = cv2.imread('../../data/a1/' + img)
        tmp = cv2.resize(tmp, (25, 25), interpolation = cv2.INTER_CUBIC)
        #nlist = [float(j) for j in list(tmp.ravel())]
        l = np.prod(tmp.shape)
        tmp = np.reshape(tmp, l)        
        train.append(tmp)
        count += 1
        
    #tr = csv_io.read_data("../../data/t1.csv")
    tr = csv_io.read_data("./training_solutions_rev1.csv")
    target = [[0 for j in range(37)] for i in range(count)]
   
   
    print len(tr[0]), len(tr), count
    for i in range(37):
        for j in range(count):
            #print i,j
            target[j][i] = tr[j][i+1]
    #print "tr, target"
    #print tr, target
    #test = csv_io.read_data("../Data/test.csv")
    test = []
    count = 0
    image_files = sorted(os.listdir('./test_part1'))
    #image_files = sorted(os.listdir('../../data/a2'))
    
    print "Test Images"
    for img in image_files:
        tmp = cv2.imread('./test_part1/' + img)
        #tmp = cv2.imread('../../data/a2/' + img)
        tmp = cv2.resize(tmp, (25, 25), interpolation = cv2.INTER_CUBIC)
        #nlist = [float(j) for j in list(tmp.ravel())]
        l = np.prod(tmp.shape)
        tmp = np.reshape(tmp, l)        
        test.append(tmp)
        count += 1        
        
    predicted_probs = [[0 for j in range(count)] for i in range(37)]
    for i in range(37):
        print "Training", i
        rf = RandomForestRegressor()
        tar1 = [0 for p in range(len(target))]
        for j in range(len(target)):        
            tar1[j]=target[j][i]
        rf.fit(train, tar1)
        pickle.dump(rf, open( "./pickle/rf"+str(i)+".p", "wb" ) )               
        
        print "Predict", i
        pred1 = rf.predict(test)
        for j in range(len(pred1)):
            predicted_probs[i][j]=pred1[j]
        
    #print predicted_probs
    predicted = [[] for i in range(len(predicted_probs[0]))]
    for i in range(len(predicted_probs[0])):
        predicted[i] = [x[i] for x in predicted_probs]
    #print predicted
    for i in range(len(predicted)):
        for j in range(len(predicted[0])):
            predicted[i][j] = "%f" %predicted[i][j]
    #predicted_probs = [x[1] for x in predicted_probs]
    csv_io.write_delimited_file("rf_benchmark.csv", predicted)

if __name__=="__main__":
    main()
