import re
import sys
from itertools import chain
from contextlib import contextmanager
try:
   import cPickle as pickle
except:
   import pickle
from logging import debug, info, warn, error
from collections import Counter
from itertools import chain, product, groupby, permutations, combinations
try:
    import xml.etree.cElementTree as ET
except:
    import xml.etree.ElementTree as ET
import os
from os.path import exists, splitext

from project import resources
from project.utils.files import resolve
from project.utils.text import filename_to_id

################################################################################
# Generate CoreNLP output
################################################################################

USE_ANNOTATORS = ['tokenize','ssplit', 'pos', 'lemma', 'ner', 'parse', 'sentiment', 'relation']

CORENLP_TRAIN_DATA_CACHE = resolve(resources.CoreNLP_base, 'CoreNLP_train.bin')
CORENLP_TEST_DATA_CACHE = resolve(resources.CoreNLP_base, 'CoreNLP_test.bin')

def _create_CoreNLP_trainxml():
    """
    Generates CoreNLP output for the training data (run once)
    """
    input_files = resources.train_data_files()
    command.run_corenlp(resolve('~', 'corenlp')
                       ,input_files
                       ,resolve('..', 'data', 'CoreNLP', 'train_data')
                       ,annotators=USE_ANNOTATORS)


def _create_CoreNLP_test_xml():
    """
    Generates CoreNLP output for the training data (run once)
    """
    input_files = resources.test_data_files()
    command.run_corenlp(resolve('~', 'corenlp')
                       ,input_files
                       ,resolve('..', 'data', 'CoreNLP', 'test_data')
                       ,annotators=USE_ANNOTATORS)


def regenerate_cache():
    """
    Regenerates the CoreNLP data cache
    """
    if exists(CORENLP_TRAIN_DATA_CACHE):
        os.remove(CORENLP_TRAIN_DATA_CACHE)
    if exists(CORENLP_TEST_DATA_CACHE):
        os.remove(CORENLP_TEST_DATA_CACHE)

################################################################################

def tokens_with_key(CoreNLP_data, filenames=None):
    """
    Tokenizes the contents of filenames using CoreNLP returning a dict keyed
    by the filename

    Note: All files in filenames are assumed to have been processed with
    CoreNLP prior
    
    @param dict CoreNLP_data The dict returned from all_sentences()
    @param [str] filenames The filenames to return the tokens of. If omitted,
        all tokens for all files appearing in CoreNLP_data will be returned
    @returns {str:[str]}
    """
    names    = CoreNLP_data.keys()
    get_word = lambda token: token['word']
    return {name: map(get_word, s['tokens']) for name in names for s in CoreNLP_data[name]}


def all_sentences(for_data):
    """
    Returns a dict of all sentences data derived from CoreNLP. The key
    is the truncated filename (observation-ID), and the value is the output
    sentence data generated by parse_xml() for that particular file.

    @returns: {str: <sentence-data>}
    """
    assert(for_data in ('train', 'test'))

    if for_data == 'train':
        data_cache_file = CORENLP_TRAIN_DATA_CACHE
    else:
        data_cache_file = CORENLP_TEST_DATA_CACHE

    # If there's cached data, load it:
    if exists(data_cache_file):

        debug('> Loading cached CoreNLP data from {}'.format(data_cache_file))

        with open(data_cache_file, 'r') as f:
            return pickle.load(f)

    # Otherwise, generate the output from parse_xml()
    debug('> CoreNLP data {} not found; caching...'.format(data_cache_file))

    if for_data == 'train':
        filenames = resources.train_data_files('CoreNLP')
    else:
        filenames = resources.test_data_files('CoreNLP')

    #if include_test:
    #    filenames += resources.test_data_files('CoreNLP')

    # parse_xml(filename)[1] means to only keep the actual sentence data,
    # not the file name/observation identifier. Also, lops off the ".xml" part
    # from the CoreNLP output filename preserving the original filename
    data = {splitext(filename_to_id(filename))[0]: parse_xml(filename)[1] for filename in filenames}

    with open(data_cache_file, 'w') as f:
        pickle.dump(data, f)

    debug('> CoreNLP data cached to {}'.format(data_cache_file))

    return data


################################################################################
# CoreNLP related parsing
################################################################################

def parse_xml(filename):
    """
    Parses a CoreNLP XML output, returning a tuple of 
    (str:observation-id, [dict:<sentence-data>]) where sentences is a list of 
    dicts, where each <sentence-data> dict has the following keys:

    - 'tokens' : [{'word':str, 'lemma':str, 'POS':str, 'NER':str|None}]
    - 'dependencies': (str, str, str)
    - 'sentiment' : str
    - 'parse' : str

    @returns (str:observation-id, [dict:<sentence-data>])
    """
    sentences = []
    tree = ET.parse(filename)

    for s in tree.findall('.//sentence'):

        sentence = {
             'tokens': []
            ,'dependencies': []
            ,'sentiment': s.get('sentiment').lower()
            ,'parse': s.find('.//parse').text
        }

        # Tokens in the sentence:
        for t in s.findall('.//token'):

            word    = t.find('word').text.lower()
            lemma   = t.find('lemma').text.lower()
            pos_tag = t.find('POS').text
            ner_tag = t.find('NER').text
            if ner_tag == "O":
                ner_tag = None

            data = {'word':word, 'lemma':lemma, 'pos':pos_tag, 'ner':ner_tag}
            sentence['tokens'].append(data)

        # Dependencies in the sentence:
        for dep in s.findall(".//dependencies[@type='basic-dependencies']/dep"):
            dep_type  = dep.get('type').lower()
            governor  = dep.find('governor').text.lower()
            dependent = dep.find('dependent').text.lower()

            data = (dep_type, governor, dependent)
            sentence['dependencies'].append(data)

        sentences.append(sentence)

    return (filename_to_id(filename), sentences)
