# -*- coding: UTF-8 -*-

import os
from itertools import chain
from optparse import OptionParser
from time import time
from collections import Counter
import re

import nltk
import sklearn
import scipy.stats
import sys

from sklearn.externals import joblib
from sklearn.metrics import make_scorer
from sklearn.cross_validation import cross_val_score
from sklearn.grid_search import RandomizedSearchCV

import sklearn_crfsuite
from sklearn_crfsuite import scorers
from sklearn_crfsuite import metrics

from pandas import DataFrame as DF
from nltk.corpus import stopwords


# Objective
# Training and evaluation of CRFs with sklearn-crfsuite.
#
# Input parameters
# --inputPath=PATH    	Path of training and test data set
# --trainingFile        File with training data set
# --testFile        	File with test data set
# --outputPath=PATH    	Output path to place output files
# --nameGrid            Number of run
# --version    		Version Report
# --nrules		Number of crf transitions

# Output
# 1) Best model
# 2) Report

# Examples
# python training_validation_v10.py
# --inputPath 			/home/egaytan/automatic-extraction-growth-conditions/CRF/data-sets
# --trainingFile 		training-data-set-70.txt
# --testFile			test-data-set-30.txt
# --outputPath 			/home/egaytan/automatic-extraction-growth-conditions/CRF/
# --version    		    _v1
# --nrules			50
# python3 training_validation_v9.py --inputPath /home/egaytan/automatic-extraction-growth-conditions/CRF/data-sets --trainingFile training-data-set-70.txt --testFile test-data-set-30.txt --outputPath /home/egaytan/automatic-extraction-growth-conditions/CRF/ --nameGrid Run1 --version _v1 --S1 --S2 --S3 --nrules 50

##################################################################
#                             FEATURES                           # 
##################################################################

#================== COMPLETE WORD FEATURES ======================#

def isGreek(word):
    ## Complete word are greek letters
    alphabet = ['Α','Β','Γ','Δ','Ε','Ζ','Η','Θ','Ι','Κ','Λ','Μ','Ν','Ξ','Ο','Π','Ρ','Σ','Τ','Υ','Φ','Χ','Ψ','Ω',
    'α','β','γ','δ','ε','ζ','η','θ','ι','κ','λ','μ','ν','ξ','ο','π','ρ','ς','σ','τ','υ','φ','χ','ψ','ω']
    if word in alphabet:
        return True
    else:
        return False 

#================ INNER OF THE WORD FEATURES ====================#

def hNumber(word):
    ## Al leats has one greek letter
    for l in word:
        if l.isdigit():
            return True
    return False

def symb(word):
    n=0
    ## At least a not alphanumeric character
    for l in word:
        if l.isdigit(): n = n+1
        if l.isalpha(): n = n+1
        #Exclude Greek letters
        if isGreek(l): n = n+1

    if n<len(word): return True
    else: return False

def hUpper(word):
    ## At least an upper letter
    for l in word:
        if l.isupper(): return True
    return False

def hLower(word):
    ## At least a lower letter
    for l in word:
        if l.islower(): return True
    return False 

def hGreek(word):
    ## At least a greek letter
    for l in word:
        if isGreek(l): return True
    return False    

#============================FEATURES===========================#

def word2features(sent, i, S1, S2, S3):
    ## Getting word features

    ## Saving CoreNLP annotations
    listElem = sent[i].split('|')
    word = listElem[0]
    lemma = listElem[1]
    postag = listElem[2]
    #ner = listElem[4]

    #=========================== G =============================#
    ## NAME LEVEL G
    ## FUTURE TYPE General features

    features = {
        ## basal features
        'lemma': lemma,
        'postag': postag
        }

    ## more tha one word in sentence
    if i > 0:        
        ## Anterior word
        listElem = sent[i - 1].split('|')

        ## Saving CoreNLP annotations
        lemma0 = listElem[1]
        postag0 = listElem[2]        
        features.update({            
            #LemaG anterior      
            '-1:lemma': lemma0,
            #Postag anterior 
            '-1:postag': postag0,
        })

    ## is not the last word
    if i < len(sent) - 1:
        ## Posterior word 
        listElem = sent[i + 1].split('|')
        ## Saving CoreNLP annotations       
        lemma2 = listElem[1]
        postag2 = listElem[2]

        features.update({           
            #LemaG  posterior      
            '+1:lemma': lemma2,
            #Postag posterior 
            '+1:postag': postag2,
        })    

    #=========================== S1 =============================#
    ## NAME LEVEL S1
    ## FEATURE TYPE Inner word features

    if S1:
        #Add features
        features['hUpper']=  hUpper(word)
        features['hLower']=  hUpper(word)
        features['hGreek']=  hGreek(word)        
        features['symb']= symb(word)
        #word firstChar
        features['word[:1]']= word[:1]    
        #word secondChar
        if len(word)>1:
            features['word[:2]']= word[:2]        
        '''
        #lemma and postag firstChar
        features['lemma[:1]']= lemma[:1]
        features['postag[:1]']= postag[:1]
        #lemma and postag secondChar
        if len(lemma)>1:
            features['lemma[:2]']= lemma[:2]
        if len(postag)>1:
            features['postag[:2]']= postag[:2]
        '''
    #=========================== S2 =============================#
    ## NAME LEVEL S2
    ## FEATURE TYPE Complete word features

    if S2:
        #Add features
        features['word']=  word
        features['isUpper']=  word.isupper()
        features['isLower']=  word.islower()
        features['isGreek']=  isGreek(word)       
        features['isNumber']= word.isdigit()
        ## more tha one word in sentence
        if i > 0:        
            ## Anterior word
            listElem = sent[i - 1].split('|')
            ## Saving CoreNLP annotations
            word0 = listElem[0]
            features['-1:word']=  word0

        ## is not the last word
        if i < len(sent)-1:
            listElem = sent[i + 1].split('|')
            ## Saving CoreNLP annotations       
            word2 = listElem[0]
            ## Posterior word 
            features['+1:word']=  word2
    
    #=========================== S3 =============================#
    ## NAME LEVEL S3
    ## FEATURE TYPE Extended context features
    if S3:            
        ## more tha two words in sentence
        if i > 1:        
            ## two anterior lemma and postag
            listElem = sent[i - 2].split('|')
            ## Saving CoreNLP annotations
            lemma01 = listElem[1]
            postag01 = listElem[2]
            features['-2:lemma']=  lemma01
            features['-2:postag']=  postag01

        ## is not the penultimate word
        if i < len(sent) - 2:
            listElem = sent[i + 2].split('|')
            ## Saving CoreNLP annotations       
            lemma02 = listElem[1]
            postag02 = listElem[2]
            ## two posterior lemma and postag
            features['+2:lemma']=  lemma02
            features['+2:postag']=  postag02

    return features

def sent2features(sent, S1, S2, S3):
    ## Itering in sentence for each word and saving its features
    return [word2features(sent, i, S1, S2, S3) for i in range(len(sent))]

def sent2labels(sent):
    ## 3rd position by word is the label
    return [elem.split('|')[3] for elem in sent]

def sent2tokens(sent):
    return [token for token, postag, label in sent]

def print_transitions(trans_features, f):
    for (label_from, label_to), weight in trans_features:
        f.write("{:6} -> {:7} {:0.6f}\n".format(label_from, label_to, weight))

def print_state_features(state_features, f):
    for (attr, label), weight in state_features:
        f.write("{:0.6f} {:8} {}\n".format(weight, label, attr.encode("utf-8")))


__author__ = 'egaytan'

##################################################################
#                            MAIN PROGRAM                        #
##################################################################

if __name__ == "__main__":
    ## Defining parameters
    parser = OptionParser()
    parser.add_option("--inputPath",        dest="inputPath",       help="Path of training data set",           metavar="PATH")
    parser.add_option("--outputPath",       dest="outputPath",      help="Output path to place output files",   metavar="PATH")
    parser.add_option("--trainingFile",     dest="trainingFile",    help="File with training data set",         metavar="FILE")
    parser.add_option("--testFile",         dest="testFile",        help="File with test data set",             metavar="FILE")
    parser.add_option("--Gridname",         dest="Gridname",        help="Report number run",                   metavar="FILE")
    parser.add_option("--version",          dest="version",         help="Report file",                         metavar="FILE")
    parser.add_option("--S1",               dest="S1",              help="Future Type",                 action="store_true", default=False)
    parser.add_option("--S2",               dest="S2",              help="Future Type",                 action="store_true", default=False)
    parser.add_option("--S3",               dest="S3",              help="Future Type",                 action="store_true", default=False)
    parser.add_option("--excludeStopWords", dest="excludeStopWords",help="Exclude stop words",          action="store_true", default=False)
    parser.add_option("--excludeSymbols",   dest="excludeSymbols",  help="Exclude punctuation marks",   action="store_true", default=False)
    parser.add_option("--nrules",	    dest="nrules",          help="Number of crf rules on report", type="int")



    (options, args) = parser.parse_args()
    if len(args) > 0:
        parser.error("Any parameter given.")
        sys.exit(1)

    print('-------------------------------- PARAMETERS --------------------------------')
    print("Path of training data set: " + options.inputPath)
    print("File with training data set: " + str(options.trainingFile))
    print("Path of test data set: " + options.inputPath)
    print("File with test data set: " + str(options.testFile))
    print("Exclude stop words: " + str(options.excludeStopWords))
    print("Levels: " + str(options.S1) + " " + str(options.S2))
    print("Report file: " + str(options.version))
    print("Number of rules on report file: " + str(options.nrules))

    
    symbols = ['.', ',', ':', ';', '?', '!', '\'', '"', '<', '>', '(', ')', '-', '_', '/', '\\', '¿', '¡', '+', '{',
               '}', '[', ']', '*', '%', '$', '#', '&', '°', '`', '...']   
    print("Exclude symbols: " + str(options.excludeSymbols))

    print('-------------------------------- PROCESSING --------------------------------')
    print('Reading corpus...')
    t0 = time()

    sentencesTrainingData = []
    sentencesTestData = []

    stopwords = [word for word in stopwords.words('english')]

    with open(os.path.join(options.inputPath, options.trainingFile), "r") as iFile:
        for line in iFile.readlines():
            listLine = []
            line = line.strip('\n')
            for token in line.split():
                if options.excludeStopWords:
                    listToken = token.split('|')
                    lemma = listToken[1]
                    if lemma in stopwords:
                        continue
                if options.excludeSymbols:
                    listToken = token.split('|')
                    lemma = listToken[1]
                    if lemma in symbols:
                        continue
                listLine.append(token)
            sentencesTrainingData.append(listLine)
        print("   Sentences training data: " + str(len(sentencesTrainingData)))

    with open(os.path.join(options.inputPath, options.testFile), "r") as iFile:
        for line in iFile.readlines():
            listLine = []
            line = line.strip('\n')
            for token in line.split():
                if options.excludeStopWords:
                    listToken = token.split('|')
                    lemma = listToken[1]
                    if lemma in stopwords:
                        continue
                if options.excludeSymbols:
                    listToken = token.split('|')
                    lemma = listToken[1]
                    if lemma in symbols:
                        continue
                listLine.append(token)
            sentencesTestData.append(listLine)
        print("   Sentences test data: " + str(len(sentencesTestData)))

    print("Reading corpus done in: %fs" % (time() - t0))

    print('-------------------------------- FEATURES --------------------------------')

    Dtraning = sent2features(sentencesTrainingData[0], options.S1, options.S2, options.S3)[2]
    Dtest = sent2features(sentencesTestData[0], options.S1, options.S2, options.S3)[2]
    print('--------------------------Features Training ---------------------------')
    print(DF(list(Dtraning.items())))
    print('--------------------------- FeaturesTest -----------------------------')
    print(DF(list(Dtest.items())))

    t0 = time()

    X_train = [sent2features(s, options.S1, options.S2, options.S3) for s in sentencesTrainingData]
    y_train = [sent2labels(s) for s in sentencesTrainingData]

    X_test = [sent2features(s, options.S1, options.S2, options.S3) for s in sentencesTestData]
    # print X_test
    y_test = [sent2labels(s) for s in sentencesTestData]

    # Fixed parameters
    # crf = sklearn_crfsuite.CRF(
    #     algorithm='lbfgs',
    #     c1=0.1,
    #     c2=0.1,
    #     max_iterations=100,
    #     all_pgossible_transitions=True
    # )

    # Hyperparameter Optimization
    crf = sklearn_crfsuite.CRF(
        algorithm='lbfgs',
        max_iterations=100,
        all_possible_transitions=True
    )
    params_space = {
        'c1': scipy.stats.expon(scale=0.5),
        'c2': scipy.stats.expon(scale=0.05),
    }

    # Original: labels = list(crf.classes_)
    # Original: labels.remove('O')
    labels = list(['Gtype', 'Gversion', 'Med', 'Phase', 'Strain', 'Substrain', 'Supp', 'Technique', 'Temp', 'OD', 'Anti', 'Agit', 'Air', 'Vess', 'pH'])

    # use the same metric for evaluation
    f1_scorer = make_scorer(metrics.flat_f1_score,
                            average='weighted', labels=labels)

    # search
    rs = RandomizedSearchCV(crf, params_space,
                            cv=10,
                            verbose=3,
                            n_jobs=-1,
                            n_iter=20,
                            # n_iter=50,
                            scoring=f1_scorer)
    rs.fit(X_train, y_train)

    # Fixed parameters
    # crf.fit(X_train, y_train)

    # Best hiperparameters
    # crf = rs.best_estimator_
  
    nameReport = str(options.Gridname) + str(options.version) + '.txt'
    with open(os.path.join(options.outputPath, "reports", "report_" + nameReport), mode="w") as oFile:
        oFile.write("********** TRAINING AND TESTING REPORT **********\n")
        oFile.write("Training file: " + options.trainingFile + '\n')
        oFile.write('\n')
        oFile.write('best params:' + str(rs.best_params_) + '\n')
        oFile.write('best CV score:' + str(rs.best_score_) + '\n')
        oFile.write('model size: {:0.2f}M\n'.format(rs.best_estimator_.size_ / 1000000))

    print("Training done in: %fs" % (time() - t0))
    t0 = time()

    # Update best crf
    crf = rs.best_estimator_

    # Saving model
    print("     Saving training model...")
    t1 = time()
    
    nameModel = 'model_S1_' + str(options.S1) + '_S2_' + str(options.S2) + '_S3_' + str(options.S3) + '_' + str(options.Gridname)   + str(options.version) + '.mod'
    joblib.dump(crf, os.path.join(options.outputPath, "models", nameModel))
    print("        Saving training model done in: %fs" % (time() - t1))

    # Evaluation against test data
    y_pred = crf.predict(X_test)
    print("*********************************")
    print("Prediction done in: %fs" % (time() - t0))

    # labels = list(crf.classes_)
    # labels.remove('O')

    with open(os.path.join(options.outputPath, "reports", "report_" + nameReport), mode="a") as oFile:
        oFile.write('\n')
        oFile.write("Flat F1: " + str(metrics.flat_f1_score(y_test, y_pred, average='weighted', labels=labels)))
        oFile.write('\n')
        # labels = list(crf.classes_)
        sorted_labels = sorted(
            labels,
            key=lambda name: (name[1:], name[0])
        )
        oFile.write(metrics.flat_classification_report(
            y_test, y_pred, labels=sorted_labels, digits=3
        ))
        oFile.write('\n')

        oFile.write("\nTop likely transitions:\n")
        print_transitions(Counter(crf.transition_features_).most_common(options.nrules()), oFile)
        oFile.write('\n')

        oFile.write("\nTop unlikely transitions:\n")
        print_transitions(Counter(crf.transition_features_).most_common()[-options.nrules():], oFile)
        oFile.write('\n')

        oFile.write("\nTop positive:\n")
        print_state_features(Counter(crf.state_features_).most_common(200), oFile)
        oFile.write('\n')

        oFile.write("\nTop negative:\n")
        print_state_features(Counter(crf.state_features_).most_common()[-200:], oFile)
        oFile.write('\n')