#!/usr/bin/env python

import functools# return altered function, etc.
import math# log, exp
import numpy as np# arrays, etc.
import sys# print without spaces
from termcolor import colored# colours while printing

# main constants
TEST_MODE = False
#TEST_MODE = True
PROGRAM_NAME = "program 1"
TYPE_COL = 29
DATA_FILE =    "../3fgl_full.dat"

#************** Building Tree ****************
# this function splits the array 'D' by the median of the 'column''s column
def splitBranch( D , column ) :
	# 'D' is the array of data, 'column' is a number of column to sort with

	# checks for column to be normal
	if column < 0 :
		print "Error: column =", column
		column = 0
	if column >= len( D [ 0 ] ) :
		print "Error: column =", column
		column = len( D [ 0 ] ) - 1

	# take the column adn get median
	D_col = D [ : , column ]
	m = np.median( D_col )

	return [
		D [ np.where( D_col < m ) ] ,
		D [ np.where ( D_col >= m ) ] ,
		[ column , m ]
	]

# this function builds the tree itself
def buildTree( D , columns ) :
	# 'D' is the array of data, 'columns' is a list of columns to sort by

	# if no columns to split
	if not columns :
		return D

	# create local copy of columns and take the first, also delete it from local array
	columns_local = columns [ : ]
	col = columns_local . pop( 0 )

	T = splitBranch( D , col )

	# if not splitted, this is the leaf
	if ( len ( T [ 0 ] ) < 1 ) or ( len ( T [ 1 ] ) < 1 ) :
		return D

	return [
		buildTree( T [ 0 ] , columns_local ) ,
		buildTree( T [ 1 ] , columns_local ) ,
		T [ 2 ]
	]

# this function sorts the row using the existing tree
def sortWithTree( T , row ) :
	# 'T' is the tree, 'row' is a single row of data
	if ( len( T ) == 3 ) and ( len( T [ 2 ] ) == 2 ) :
		column = T [ 2 ] [ 0 ]
		median = T [ 2 ] [ 1 ]

		if row[ column ] < median :
			return sortWithTree( T [ 0 ] , row )

		return sortWithTree( T [ 1 ] , row )

	#if we are here then there is no tree structure anymore, we are at the final level
	types = T [ : , TYPE_COL ]
	pulsar = len( T [ np.where( types == 1 ) ] )
	blazar = len( T [ np.where( types == -1 ) ] )
	return 1 if pulsar > blazar else -1

# this function returns the classifier -- the function which takes only the row for the parameter and returns its classification
def getClassifier ( D_TRAIN, col_range ) :
	# build a tree

	DTree = buildTree( D_TRAIN , col_range )

	#the first argument is the name of the function, the second is the first parameter of it. We return the function with a predefined first parameter
	return functools.partial( sortWithTree , DTree )


#************** Classic Boosting *************
# this function returns the right classifier depending on the results of the first two
def boostClassifierChoice ( CL1 , CL2 , CL3 , row ) :
	return CL1( row ) if CL1( row ) == CL2 ( row ) else CL3( row )

# the function of testing the classifier
def testClassifier ( CL , D_TESTS ) :
	# run the classificator for each row in the training set 
	D_TEST_RES = np.array( map( CL , D_TESTS ) )

	return [
		len( D_TESTS [ np.where( D_TESTS [ : , TYPE_COL ] != D_TEST_RES ) ] ),# number of errors
		len( D_TESTS ),# number of tests
		len( D_TESTS [ np.where( ( D_TESTS [ : , TYPE_COL ] != D_TEST_RES ) & ( D_TESTS [ : , TYPE_COL ] == -1 ) ) ] ),#number of blazar errors
		len( D_TESTS [ np.where( D_TESTS [ : , TYPE_COL ] == -1 ) ] ),# number of blazars
		len( D_TESTS [ np.where( ( D_TESTS [ : , TYPE_COL ] != D_TEST_RES ) & ( D_TESTS [ : , TYPE_COL ] ==  1 ) ) ] ),# number of pulsar errors
		len( D_TESTS [ np.where( D_TESTS [ : , TYPE_COL ] ==  1 ) ] )# number of pulsars
	]


# the function of boosting itself
def boostClassifierOnce ( CL1 , D_TRAIN , col_range ) :
	# test the classifier
	D_TEST_RES_1 = np.array( map( CL1 , D_TRAIN ) )

	# get all errors
	D_ERROR = D_TRAIN [ np.where ( D_TRAIN [ : , TYPE_COL ] != D_TEST_RES_1 ) ]
	error_count = len ( D_ERROR )

	if error_count < 1 :
		return [ CL1 , False ]

	# get the same number of correct
	D_CORRECT_TMP = D_TRAIN [ np.where ( D_TRAIN [ : , TYPE_COL ] == D_TEST_RES_1 ) ]
	np.random.shuffle ( D_CORRECT_TMP )
	D_CORRECT = D_CORRECT_TMP [ : error_count ]

	# new training set with all the errors and the same number of correct
	D_TRAIN_NEW = np.concatenate( ( D_ERROR , D_CORRECT ) , axis=0 )

	# build the second classifier
	CL2 = getClassifier( D_TRAIN_NEW , col_range )