#!/usr/bin/env python
from random import shuffle,seed
from rndForest import * #this automatically initialize R-interface
from sys import stderr,stdout, argv
from svmprob import read_svm, getarg
seed(0)
if len(argv) < 3:
	print "%s [-logfile logfile][-cv nr_fold] <train_file> <test_file>" % argv[0]
	raise SystemExit

nr_fold = int(getarg('-cv', 5))
logfile = getarg('-logfile', None)
#read data
train_x, train_y = read_svm(argv[-2])
test_x,test_y = read_svm(argv[-1])

mM = findminmax(train_x + test_x)
n_attr = mM[1] - mM[0] + 1
# default:
# 	max(floor(ncol(x)/3), 1), floor(sqrt(ncol(x))))
# search space:
#	floor(c(1, sqrt(p), p/3, p/2, p))   
r.unique.local_mode(BASIC_CONVERSION)
r.sort.local_mode(BASIC_CONVERSION)

last = -1
bestcverr = len(train_x) #all wrong
#prepare training & testing data
f=r.sort(r.unique(train_y))
rx = svm2matrix(train_x,mM)
ry =  r.as_factor(train_y)
tx = svm2matrix(test_x,mM)
rlen = len(train_y)
tlen = len(test_y)
#prepare data in folds
#randomforest require label be factor()
cvrx, cvry, cvtx, cvty = [], [], [], []
for i in range(nr_fold):
	begin = i * (rlen) / nr_fold
	end = (i + 1) * (rlen) / nr_fold
	cvrx.append(svm2matrix(train_x[:begin]+train_x[end:],mM))
	cvry.append(r.factor(train_y[:begin]+train_y[end:],levels=f))
	cvtx.append(svm2matrix(train_x[begin:end],mM))
	cvty.append(train_y[begin:end])

#begin Experiment##################
def rftest(X,Y,TX,TY,MTRY):
      errcount = 0
      model=r.randomForest(x=X,y=Y, mtry = MTRY)
      pred=r.predict(model, TX, type='response')
      for j in range(len(TY)):
              if(TY[j] == int(pred[j])):
                      stderr.write('.')
              else:
                      errcount+=1
                      stderr.write('*')
      return errcount

for m in (1,int(sqrt(n_attr)), n_attr/3, n_attr/2, n_attr):
	if m == last:	#avoid doing same mtry twice
		continue
	cverr = 0
	for i in range(nr_fold):
		stderr.write(str(i))
		cverr += rftest(cvrx[i], cvry[i], cvtx[i], cvty[i], m)

	if cverr < bestcverr:
		bestcverr = cverr
		bestcvm = m
	stdout.write('\nCV:mtry = %d, error = %d\n' % (m, cverr))
	last = m
#begin testing
model=r.randomForest(x=rx,y=ry, mtry = bestcvm)
pred=r.predict(model, tx, type='prob')

#make a dict of class labels
all_cls = {}
for i in train_y + test_y:
	all_cls[i] = 1.0
all_cls = all_cls.keys()
all_cls.sort()

#error accounting
correct = 0.0
SE = 0.0
for i in range(tlen):
	p = pred[i]  # p is [prob_1, prob_2 .... prob_n]
	maxp = max(p)
	winners = []
	for j in range(len(all_cls)):
		if p[j] ==  maxp:
			winners.append(all_cls[j])
		if test_y[i] == all_cls[j]:
			SE += (1.0 - p[j])*(1.0 - p[j])
		else:
			SE += p[j]*p[j]
	if len(winners) > 1:
		print "more than one winner:",winners
		shuffle(winners)
	if test_y[i] == winners[0]:
		correct += 1
stdout.write("Testing %s MSE %f Accuracy %f%% mtry=%d\n" % \
	('pure-RF', SE/tlen,100*correct/tlen, bestcvm))
if logfile:
		logfile= logfile.replace('@host', gethostname())
		print "Output Result:", logfile
		fp = open(logfile , 'a') 
		flock(fp.fileno(), LOCK_EX)
		fp.write("%d %f %f %s\n" % (mtry , correct/tlen, SE/tlen, 'pure-RF'))
		fp.flush()
		fp.close()
		flock(fp.fileno(), LOCK_UN)
