#!/usr/bin/env python
from sys import argv, stdin, stderr
from svmprob import read_svm, proc_argv, getarg
from string import join
from random import seed,randrange,shuffle
from fcntl import flock, LOCK_EX, LOCK_UN
from socket import gethostname

#coupling methods-------------------------------------------------------
from coupler import ht, markov, pkpd, minpair, vote
from coupler import normalize, coupler_list

#prob-out binary classifiers--------------------------------------------
base_list={}
#dynamic randomForest base
use_RF = 0
def load_RF():
	global use_RF, RFTrain, RFProb, findminmax
	use_RF = 1
	from rndForest import RFTrain, RFProb, findminmax
	base_list['rf']=(RFTrain, RFProb)
#libsvm bases
from svmPlatt import SvmPlattTrain, SvmPlattProb, SvmDeciTrain, SvmDeciTest
base_list['svmplatt']=(SvmPlattTrain,SvmPlattProb)
base_list['svmdeci']=(SvmDeciTrain, SvmDeciTest)
#----------------------------------------------------------------------
from numpy import log
epsilon = 0.0000001
#numerical log2 based
def log2(v):
	return log(float(v))/log(2)

#input comma seperated coupler_str, and return a list of function pointers
def proc_couplers(coupler_str):
	if coupler_str == 'all':
		return coupler_list
	ret = []
	for m in coupler_str.split(','):
		for c in coupler_list:
			if  m == c.__name__:
				ret.append(c)
	if ret == []:
		ret = [coupler_list[3]]
		print "Using default coupler set:", map(lambda x:x.__name__,ret)
	return ret

def proc_bases(base_str):
	ret = {}
	for b in  base_str.split(','):
		if base_list.has_key(b):
			ret[b] = base_list[b]

	if ret == {}:
		ret = {'svmplatt':base_list['svmplatt']}
		print "Using default base", ret.keys()
	return ret

def square(x):
	return x*x


#accept training and testing data with specified 'pairwise coupling' 
#function pointer list, then return a dict correct[couplers]
def testing(couplers, base, train_x, train_y, test_x, test_y, param, outfile=None):

	#XXX: should do checking on coupler<->base
	base_train, base_test  = base

	#divide data according to its class
	cls_x, cls_cnt = {},{}
	#init, find out classes
	for i in range(len(train_x)):
		cls_x[train_y[i]]=[]
		cls_cnt[train_y[i]]=0.0
	#dispatch, split multiclass into binary problems
	for i in range(len(train_x)):
		cls_x[train_y[i]].append(train_x[i])
		cls_cnt[train_y[i]]+=1
	#train, obtain Platt and libsvm_predict model for each binary problem
	all_cls = cls_x.keys()
	all_cls.sort()
	print "Classes dispatched:", all_cls

	#making model
	model = {}  #model[[i,j]] j>i
	for i in all_cls:
		for j in all_cls:
			if j > i:
				model[i,j]=base_train(cls_x[i], cls_x[j], param)
				#SvmPlattTrain(cls_x[i], cls_x[j], (svm_param, platt_cv_fold))
				#RFTrain(cls_x[i], cls_x[j], (mtry, mM))
				#SvmDeciTrain(cls_x[i], cls_x[j], svm_param)
	#counter reset
	correct = {}
	logloss = {}
	SE = {}		#accumlated square error
	for c in couplers:
		correct[c] = 0.
		logloss[c] = 0.
		SE[c] = 0.
	total = len(test_x)
	#begin testing
	prob_bin={}
	votes = {}
	tcount = 0
	for t in range(len(test_x)):
		tcount += 1
		for i in all_cls:
			votes[i] = 0
		#create P^i_{ij} and votes for current test instance
		for i in all_cls:
			for j in all_cls:
				if j>i:
					prob_raw = base_test(test_x[t], model[i,j])
					#truncate
					if prob_raw < epsilon:
						prob_raw = epsilon
					elif prob_raw > (1-epsilon):
						prob_raw = 1-epsilon
					#assign
					prob_bin[i,j]= prob_raw
					prob_bin[j,i]= 1. - prob_raw
						
		#use coupler to count
		for c in couplers:
			mprob = c(prob_bin,cls_cnt)
			for i in all_cls:
				mprob[i] = max(epsilon, mprob[i])
			#logloss \sum \log{p_i_{y_i}}
			logloss[c] += log(mprob[test_y[t]])
			winners = []
			for i in all_cls:
				#take largest prob(s) as predicted_answers
				if mprob[i]==max(mprob.values()):
					winners.append(i)
				#count Square Error
				if test_y[t]==i:
					SE[c]+=square(1.0-mprob[i])
				else:
					SE[c]+=square(mprob[i])
			if len(winners) > 1:
				#continue		#correct require the only winner
				stderr.write(c.__name__ + str(winners))
				shuffle(winners)
			if test_y[t] == winners[0]:
				correct[c] += 1
			if outfile:
				outfile.write("%s %s\n" % (c.__name__ , join(map(lambda x:str(mprob[x]) ,all_cls), ' ')));
		stderr.write('.')

	#*Mean* square error	
	print
	for c in couplers:
		SE[c] = SE[c]/len(all_cls)
    		
	for c in couplers:
		print "SubProb Accuracy %s(%d/%d) %f%% MSE %f Logloss %f" % (c.__name__,correct[c],total ,(100*float(correct[c])/total), SE[c]/total, logloss[c]/total)
	return correct,SE,logloss


def main():
	if len(argv) < 2:
		print """usage: %s 
[-couplers {ht,pkpd,markov,*minpair,vote,dag,ass}]	#couplers separate by comma
[-base {*svmplatt,rf,svmdeci}]				#binary probability classifier
[-cv cv_fold][-logfile logfile][-seed rndseed]		#general parameters
[-svm "SVM Parameters"][-c cost][-g gamma]		#SVM related parameters
[-mtry mtry]						#randomForest parameters
[-o <outfile>][-T <testfile>] <train_file>		#test and traing files""" % argv[0]
		raise SystemExit
	#load_RF if needed
	if getarg('-base', "").find('rf') != -1:
		load_RF()
	#cmdline arguments
	couplers = proc_couplers(getarg('-couplers', ""))
	bases = proc_bases(getarg('-base', ""))
	cv_fold = int(getarg('-cv', '5'))
	logfile = getarg('-logfile', None)
	svm_param, platt_cv_fold, dummy, dummy= proc_argv(('-svm '+getarg('-svm','')).split())
	svm_param.C, svm_param.gamma =  float(getarg('-c', svm_param.C)), float(getarg('-g', svm_param.gamma))
	mtry = int(getarg('-mtry', 1)) 	#default: automatic 
	outfile = getarg('-o', None)		#default: no output of probability
	test_file = getarg('-T',None)		#default: no train file
	train_file = argv[-1]
	seed(int(getarg('-seed', 0)))		#reset random seed
	#read data
	train_x, train_y = read_svm(train_file)
	train_len = len(train_x)
	if test_file:
		test_x, test_y = read_svm(test_file)
		test_len = len(test_x)
		total_data = train_x + test_x
		jobname = 'Testing'
	else:
		test_len = train_len
		total_data = train_x
		jobname = 'Cross Validation'
	if use_RF:
		mM = findminmax(total_data)

	#prepare parameter dict
	params = {}
	params['svmplatt'] = (svm_param, platt_cv_fold)
	params['svmdeci'] = (svm_param)
	if use_RF:
		params['rf'] = (mtry,mM)

	#Take first base as well as its parameter (Extensible)
	base_key, base_func = bases.items()[0]	#workaround, single base
	param = params[base_key]

	#init train
	if outfile:
		outfile = open(outfile,'w')
	total_correct, total_SE, total_logloss = {},{},{}
	for c in couplers:
		total_correct[c],total_SE[c],total_logloss[c]=0., 0., 0.
	#permutation 
	for i in range(train_len):
		j = randrange(i,train_len)
		train_x[i], train_x[j] = train_x[j], train_x[i]
		train_y[i], train_y[j] = train_y[j], train_y[i]

	if test_file:	#test mode
		total_correct, total_SE, total_logloss = testing(couplers, base_func,\
				train_x, train_y, test_x, test_y, param, outfile)
	else:			#CV mode
		for i in range(cv_fold):
			begin = i * train_len/ cv_fold
			end = (i + 1) * train_len / cv_fold
			correct,SE,logloss = testing(couplers, base_func, \
						train_x[:begin] + train_x[end:], \
						train_y[:begin] + train_y[end:], \
						train_x[begin:end],train_y[begin:end], param, outfile)
			for c in couplers:
				total_correct[c] += correct[c]
				total_SE[c] += SE[c]
				total_logloss[c] += logloss[c]
	if outfile:
		outfile.close()

	#only output accuracy of first method to stdout for grid-search
	print "Accuracies:",map(lambda x:(x.__name__,(total_correct[x]/test_len)), couplers)
	print "LogLoss:", map(lambda x:(x.__name__,(total_logloss[x]/test_len)), couplers)
	print "%s MSE %f Accuracy %f%%" % (jobname, total_SE[couplers[0]]/test_len, 100*total_correct[couplers[0]]/test_len)

	#experimental output
	if logfile:
		#separate files for using in Cluster
		logfile= logfile.replace('@host', gethostname())
		print "Output Result:", logfile
		fp = open(logfile , 'a') 
		flock(fp.fileno(), LOCK_EX)
		for c in couplers:
			if base_key == 'svmplatt' or base_key == 'svmdeci' :
				fp.write("%.2f %.2f %f %f %f %s\n" % \
				(log2(svm_param.C), log2(svm_param.gamma), \
				total_correct[c]/test_len, total_SE[c]/test_len, total_logloss[c]/test_len, c.__name__))
			else:
				fp.write("%d %f %f %f %s\n" % (mtry , \
				total_correct[c]/test_len, total_SE[c]/test_len, total_logloss[c]/test_len, c.__name__))
		fp.flush()
		flock(fp.fileno(), LOCK_UN)
		fp.close()

if __name__ == '__main__':
	main()
