'''
#=============================================================================
#     FileName: pls_train.py
#         Desc: 
#       Author: jlpeng
#        Email: jlpeng1201@gmail.com
#     HomePage: 
#      Created: 2014-08-21 14:10:24
#   LastChange: 2014-08-22 13:12:34
#      History:
#=============================================================================
'''
import sys
import pickle
from getopt import getopt
import numpy as np
try:
    from sklearn.cross_decomposition import PLSRegression
except ImportError:
    from sklearn.pls import PLSRegression
from sklearn.base import clone
from sklearn.cross_validation import KFold

train_y_file = None
nijr_file = None
n_comp = 0
cv = 0
verbose = False
model_file = None
USE_KFOLD = False

def parse_arguments(argv):
    options,args = getopt(argv[1:], 'n:', ['try=','nijr=','cv=','verbose'])
    global train_y_file
    global nijr_file
    global n_comp
    global cv
    global verbose
    global model_file
    for opt,val in options:
        if opt == '--try':
            train_y_file = val
        elif opt == '--nijr':
            nijr_file = val
        elif opt == '-n':
            n_comp = int(val)
            assert n_comp > 0
        elif opt == '--cv':
            cv = int(val)
            assert cv > 1
        elif opt == '--verbose':
            verbose = True
        else:
            print >>sys.sderr, "Error: invalid option",opt
            sys.exit(1)
    if train_y_file is None:
        print >>sys.stderr, "Error: '--try' is needed"
        sys.exit(1)
    if nijr_file is None:
        print >>sys.stderr, "Error: '--nijr' should be given"
        sys.exit(1)
    if cv > 0:
        if n_comp==0:
            print >>sys.stderr, "Error: '-n' is needed when '--cv' is given"
            sys.exit(1)
    else:
        assert len(args) == 1
        model_file = args[0]

    
def exit_with_help(name):
    print "\nUsage:"
    print "  %s [options] output.model"%name
    print "\n[options]"
    print "  --try   file: each line should be `name y-value`"
    print "  --nijr  file: generated by get_nijr.py"
    print "  -n    n_comp: specify the number of components to be used"
    print "                if not given, then it'll be optimized using cross-validation"
    print "                to get the number of components yielding smallest RMSE"
    print "  --cv       n: <optional>"
    print "                do {n}-fold cross-validation"
    print "                in this case, '-n' must be given"
    print "  --verbose   : if given, additional information will be displayed"
    print "                e.g. the predicting result of each sample will be displayed"
    print "                     show the number of components being tested"
    print "                     ..."
    print ""
    sys.exit(1)


def main(argv=sys.argv):
    if len(argv) < 2:
        exit_with_help(argv[0])
    
    #parse options
    parse_arguments(argv)

    global train_y_file
    global nijr_file
    global n_comp
    global cv
    global verbose
    global model_file

    #read training set
    train_list = []
    train_ys = []
    inf = open(train_y_file,'r')
    for line in inf:
        if line.startswith('#'):
            continue
        name,val = line.split()
        train_list.append(name)
        train_ys.append(float(val))
    inf.close()
    train_ys = np.asarray(train_ys)
    if verbose:
        print "train_ys.shape:",train_ys.shape
    
    #read nijr
    #each item `i,j,n`  ==> Xs[][j*len(DC_TYPES)+i] = n
    inf = open(nijr_file,'r')
    line = inf.readline()
    num_DC_TYPES = int(line.split(',')[0].split('=')[-1])
    num_bins = int(line.split('=')[-1])
    train_Xs = np.zeros((len(train_list), num_DC_TYPES*num_bins))
    count = 0
    for line in inf:
        if line.startswith('#'):
            continue
        line = line.split()
        if line[0] in train_list:
            k = train_list.index(line[0])
            for item in line[1:]:
                i,j,n = map(int, item.split(','))
                index = j*num_DC_TYPES+i
                train_Xs[k, index] = n
            count += 1
        else:
            print line[0],"not found in",train_y_file
            sys.exit(1)
    inf.close()
    if count != len(train_list):
        print >>sys.stderr, "Error: number of samples(%d) in %s not equal to those(%d) in %s"%(count, nijr_file, len(train_list), train_y_file)
        sys.exit(1)
    if verbose:
        print "train_Xs.shape:",train_Xs.shape

    #do cross-validation only
    if cv > 0:
        print "do %d-fold cross-validation"%cv
        actualY,predY = doCV(train_Xs, train_ys, n_comp, cv)
        mae = calcMAE(actualY,predY)
        rmse = calcRMSE(actualY,predY)
        r = calcR(actualY, predY)
        print "\n  statistics of %d-fold CV"%cv
        print "  MAE=%g, RMSE=%g, r=%g\n"%(mae, rmse, r)
        return 

    #train PLS model
    if n_comp > 0:
        if n_comp > train_Xs.shape[1]:
            print "Error: the number of components specified(%d) is larger than the number of predictors(%d)"%(n_comp, train_Xs.shape[1])
            sys.exit(1)
        print "\nPLS will be trained using %d components"%n_comp
        pls1 = PLSRegression(n_components=n_comp)
    else:
        n_components,best_val = search_n_components(train_Xs, train_ys, verbose)
        print "\nPLS will be trained using %d components, with RMSE=%g of 5-fold CV"%(n_components, best_val)
        pls1 = PLSRegression(n_components=n_components)

    pls1.fit(train_Xs,train_ys)
    print "\nto apply the PLS model to training set"
    predict_and_display(pls1, train_Xs, train_ys, train_list, verbose)

    #save model
    outf = open(model_file,'w')
    pickle.dump(pls1, outf)
    outf.close()


calcMAE = lambda actualY, predictY: np.mean(np.abs(actualY-predictY))
calcRMSE = lambda actualY, predictY: np.sqrt(np.mean(np.power(actualY-predictY,2)))
calcR = lambda actualY, predictY: np.corrcoef(actualY, predictY)[0][1]

def predict_and_display(pls1, Xs, ys, names, verbose):
    predY = pls1.predict(Xs)
    predY = np.ndarray.flatten(predY)
    rmse = calcRMSE(ys, predY)
    mae = calcMAE(ys, predY)
    r = calcR(ys, predY)
    if verbose:
        print "name actualY predictY"
        for i in xrange(len(names)):
            print names[i],ys[i],predY[i]
        print ""
    print "  RMSE=%g, MAE=%g, r=%g\n"%(rmse, mae, r)

def search_n_components(Xs, Ys, verbose=False):
    best_val = 1e38
    best_n = 0
    if verbose:
        print "\nsearch for the best number of components"
    for i in xrange(1, Xs.shape[1]+1):
        if verbose:
            print "> to use %d components"%i,
        actualY,predY = doCV(Xs, Ys, i, 5)
        val = calcRMSE(actualY, predY)
        r = calcR(actualY, predY)
        if verbose:
            print "  => result of 5-fold CV: RMSE=%g, r=%g"%(val,r)
        if val < best_val:
            best_val = val
            best_n = i
    if verbose:
        print "best_n_components=%d, best_val=%g"%(best_n, best_val)
    return best_n,best_val

def doCV(Xs, Ys, n_components, nfold):
    global USE_KFOLD
    if USE_KFOLD:
        print "using KFold"
        kf = KFold(Xs.shape[0], n_folds=nfold, shuffle=True)
    else:
        kf = MyFold(Xs.shape[0], nfold)
    actualY = []
    predictY = []
    for train,test in kf:
        X_train,X_test = Xs[train],Xs[test]
        y_train,y_test = Ys[train],Ys[test]
        pls1 = PLSRegression(n_components=n_components)
        pls1.fit(X_train,y_train)
        y_test_pred = pls1.predict(X_test)
        actualY.extend(list(y_test))
        predictY.extend(list(np.ndarray.flatten(y_test_pred)))
    return np.asarray(actualY), np.asarray(predictY)

class MyFold:
    def __init__(self, length, nfolds):
        self.length = length
        self.nfolds = nfolds
        self.i = 0

    def __iter__(self):
        return self

    def next(self):
        self.i += 1
        if self.i > self.nfolds:
            raise StopIteration
        tr_idx = []
        te_idx = []
        for i in xrange(self.length):
            if i%(self.nfolds) == self.i-1:
                te_idx.append(i)
            else:
                tr_idx.append(i)
        return (tr_idx,te_idx)

if __name__ == "__main__":
    main()

