from lib.CoNLL_Annotation import *
from collections import Counter
import pandas as pd
import numpy as np
from sklearn.metrics import precision_recall_fscore_support as eval_f1
from tabulate import tabulate
import logging, argparse, sys
from datetime import datetime


def eval_lemma(sys, gld):
    match, err, symbol = 0, 0, []
    mistakes = []
    for i, gld_tok in enumerate(gld.tokens):
        if gld_tok.lemma == sys.tokens[i].lemma:
            match += 1
        elif not sys.tokens[i].lemma.isalnum(): # This was added because Turku does not lemmatize symbols (it only copies them) => ERR ((',', '--', ','), 43642)
            symbol.append(sys.tokens[i].lemma)
            if sys.tokens[i].word == sys.tokens[i].lemma:
                match += 1
            else:
                err += 1
        else:
            err += 1
            mistakes.append((gld_tok.word, gld_tok.lemma, sys.tokens[i].lemma))
    return match, err, symbol, mistakes
    

def eval_pos(sys, gld):
    match, mistakes = 0, []
    y_gld, y_pred = [], []
    for i, gld_tok in enumerate(gld.tokens):
        y_gld.append(gld_tok.pos_tag)
        y_pred.append(sys.tokens[i].pos_tag)
        all_pos_labels.add(gld_tok.pos_tag)
        if gld_tok.pos_tag == sys.tokens[i].pos_tag:
            match += 1
        else:
            mistakes.append((gld_tok.word, gld_tok.pos_tag, sys.tokens[i].pos_tag))
    return y_gld, y_pred, match, mistakes


if __name__ == "__main__":
    """
        EVALUATIONS: 
            python TIGER/evaluate.py -t Turku\
                --sys_file /home/daza/datasets/TIGER_conll/tiger_turku_parsed.conllu \
                --gld_file /home/daza/datasets/TIGER_conll/tiger_release_aug07.corrected.16012013.conll09
                
            python TIGER/evaluate.py -t SpaCy\
                --sys_file /home/daza/datasets/TIGER_conll/tiger_spacy_parsed.conllu \
                --gld_file /home/daza/datasets/TIGER_conll/tiger_release_aug07.corrected.16012013.conll09
                
            python TIGER/evaluate.py -t RNNTagger\
                --sys_file /home/daza/datasets/TIGER_conll/tiger_all.parsed.RNNTagger.conll \
                --gld_file /home/daza/datasets/TIGER_conll/tiger_release_aug07.corrected.16012013.conll09
            
            python TIGER/evaluate.py -t TreeTagger\
                --sys_file /home/daza/datasets/TIGER_conll/tiger_all.parsed.TreeTagger.conll \
                --gld_file /home/daza/datasets/TIGER_conll/tiger_release_aug07.corrected.16012013.conll09
    """
    
    # =====================================================================================
    #                    INPUT PARAMS 
    # =====================================================================================
    parser = argparse.ArgumentParser()
    parser.add_argument("-s", "--sys_file", help="System output in CoNLL-U Format", required=True)
    parser.add_argument("-g", "--gld_file", help="Gold Labels to evaluate in CoNLL-U Format", required=True)
    parser.add_argument("-t", "--type_sys", help="Which system produced the outputs", default="system")
    args = parser.parse_args()
    
    # =====================================================================================
    #                    LOGGING INFO ...
    # =====================================================================================
    logger = logging.getLogger(__name__)
    console_hdlr = logging.StreamHandler(sys.stdout)
    file_hdlr = logging.FileHandler(filename=f"logs/Eval_Tiger.{args.type_sys}.log")
    logging.basicConfig(level=logging.INFO, handlers=[console_hdlr, file_hdlr])
    now_is = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    logger.info(f"\n\nEvaluating TIGER Corpus {now_is}")

    # Read the Original TiGeR Annotations
    gld_generator = read_conll_generator(args.gld_file, token_class=CoNLL09_Token)
    # Read the Annotations Generated by the Automatic Parser [Turku, SpaCy, RNNTagger] 
    if args.type_sys == "RNNTagger":
        sys_generator = read_conll_generator(args.sys_file, token_class=RNNTagger_Token)
    elif args.type_sys == "TreeTagger":
        sys_generator = read_conll_generator(args.sys_file, token_class=RNNTagger_Token, sent_sep="</S>")
    else:
        sys_generator = read_conll_generator(args.sys_file, token_class=CoNLLUP_Token)
    
    lemma_all_match, lemma_all_err, lemma_all_mistakes = 0, 0, []
    lemma_all_symbols = []
    pos_all_match, pos_all_err, pos_all_mistakes = 0, 0, []
    pos_all_pred, pos_all_gld = [], []
    all_pos_labels = set()
    
    for i, (s,g) in enumerate(zip(sys_generator, gld_generator)):
        # print([x.word for x in s.tokens])
        # print([x.word for x in g.tokens])
        assert len(s.tokens) == len(g.tokens), f"Token Mismatch! S={len(s.tokens)} G={len(g.tokens)} IX={i+1}"
        # Lemmas ...
        lemma_match, lemma_err, lemma_sym, mistakes = eval_lemma(s,g)
        lemma_all_match += lemma_match
        lemma_all_err += lemma_err
        lemma_all_mistakes += mistakes
        lemma_all_symbols += lemma_sym
        # POS Tags ...
        pos_gld, pos_pred, pos_match, pos_mistakes = eval_pos(s, g)
        pos_all_pred += pos_pred 
        pos_all_gld += pos_gld
        pos_all_match +=  pos_match
        pos_all_err += len(pos_mistakes)
        pos_all_mistakes += pos_mistakes
    
    # Lemmas ...
    logger.info(f"Lemma Matches = {lemma_all_match} || Errors = {lemma_all_err} || Symbol Chars = {len(lemma_all_symbols)}")
    logger.info(f"Lemma Accuracy = {(lemma_all_match*100/(lemma_all_match + lemma_all_err)):.2f}%\n")
    lemma_miss_df = pd.DataFrame(lemma_all_mistakes, columns =['Gold_Word', 'Gold_Lemma', 'Sys_Lemma']).value_counts()
    lemma_miss_df.to_csv(path_or_buf=f"outputs/LemmaErrors.{args.type_sys}.tsv", sep="\t")
    
    # POS Tags ...
    logger.info(f"POS Matches = {pos_all_match} || Errors = {pos_all_err}")
    logger.info(f"POS Tagging Accuracy = {(pos_all_match*100/(pos_all_match + pos_all_err)):.2f}%\n")
    pos_miss_df = pd.DataFrame(pos_all_mistakes, columns =['Gold_Word', 'Gold_POS', 'Sys_POS']).value_counts()
    pos_miss_df.to_csv(path_or_buf=f"outputs/POS-Errors.{args.type_sys}.tsv", sep="\t")
    
    ordered_labels = sorted(all_pos_labels)
    p_labels, r_labels, f_labels, support = eval_f1(y_true=pos_all_gld, y_pred=pos_all_pred, labels=ordered_labels , average=None)
    scores_per_label = zip(ordered_labels, [x*100 for x in p_labels], [x*100 for x in r_labels], [x*100 for x in f_labels])
    logger.info("\n\n")
    logger.info(tabulate(scores_per_label, headers=["POS Tag","Precision", "Recall", "F1"], floatfmt=".2f"))
    p_labels, r_labels, f_labels, support = eval_f1(y_true=np.array(pos_all_gld), y_pred=np.array(pos_all_pred), average='micro', zero_division=0)
    logger.info(f"Total Prec = {p_labels*100}\tRec = {r_labels*100}\tF1 = {f_labels*100}")