"""Transformers that use stanza documents."""

import json
import logging
import os
import warnings

import nltk
import stanza
from sklearn.base import BaseEstimator, TransformerMixin
from stanza import Document
from stanza.models.common.doc import Sentence
from tqdm import tqdm

logger = logging.getLogger(__name__)

UNKNOWN_KEY = 'UNK'
WORD_FEATURES_JSON = os.path.join('resources', 'word_features.json')


def _get_word_attribute(word, attribute, fallback='_'):
    if fallback is None:
        raise ValueError("Fallback can't be None")
    if hasattr(word, attribute):
        a = getattr(word, attribute)
        if a is not None:
            return a
    if not hasattr(word, '_feats'):
        return fallback

    feats = word.feats
    if not feats or feats == '_':
        return fallback
    pairs = feats.split('|')
    features = {k: v for (k, v) in [x.split('=') for x in pairs]}
    if attribute not in features or features[attribute] is None:
        return fallback
    return features[attribute]


def _get_features_per_word(document):
    result = []
    words = []
    for sentence in document.sentences:
        words += sentence.words
    for word in words:
        features = word.feats
        if not features or features == '_':
            continue
        pairs = features.split('|')
        result.append({k: v for (k, v) in [x.split('=') for x in pairs]})
    return result


class StanzaWordFeatureFrequencyTransformer(BaseEstimator, TransformerMixin):
    """
    Returns a frequency matrix for all grammar features for each word of each
    document.

    This transformer looks for the file resources/word_features.json, where all
    possible values should be stored (unfortunately, stanfordnlp does not offer
    a comprehensive list of possible values for each model).

    If this transformer encounters an element not in that file, it will add one
    """

    def __init__(self):
        with open(WORD_FEATURES_JSON) as i_fh:
            self.features = json.load(i_fh)
            self.features.append(UNKNOWN_KEY)  # for unknown keys

    def fit(self, X, y=None, **fit_kwargs):
        return self

    def transform(self, X, y=None):
        result = []
        for document in X:
            doc = {}
            for feature in self.features:
                doc[feature] = 0
            for features_dict in _get_features_per_word(document):
                for k, v in features_dict.items():
                    key = f'{k}__{v}'
                    if key not in doc:
                        key = UNKNOWN_KEY
                        warnings.warn(
                            f'key {key} not found during fit, returning UNK')
                    doc[key] += 1
            result.append(doc)
        return result


class StanzaNlpToFieldTransformer(BaseEstimator, TransformerMixin):
    """
    Flattens a stanford document in the same order as the parsed text.

    Input (Document | Sentence):
        each document is expected to be a StanfordNLP document or sentence.
    Output: each document is returned as a single string, where each word is
        represented as the field that is to be extracted, separated by spaces.
        Sentences are separated by newlines.
    """

    def __init__(self, field=None):
        self.field = field

    def fit(self, x, y=None):
        return self

    def transform(self, x, y=None):
        result = []
        for document in x:
            document_result = []
            sentences: [Sentence] = []
            if type(document) == Sentence:
                sentences = [document]
            elif type(document) == Document:
                sentences = document.sentences
            for sentence in sentences:
                new_sentence = []
                for word in sentence.words:
                    label = _get_word_attribute(word, self.field, fallback='_')
                    if label is None:
                        raise ValueError(f'label is None for word: {word}')
                    new_sentence.append(label)
                document_result.append(' '.join(new_sentence))
            result.append('\n'.join(document_result))
        return result


class StanzaParserTransformer(BaseEstimator, TransformerMixin):
    """
    This transformer takes text and parses it using the stanza python
    package, which uses theano and various neural models to parse different
    features from natural text.
    """

    def __init__(self, language, silent=False, cpu=False):
        self.language = language
        self.silent = silent
        self.cpu = cpu
        if cpu:
            logger.info('using CPU for stanford parsing')
            os.environ['CUDA_VISIBLE_DEVICES'] = ''
        self.nlp = stanza.Pipeline(lang=self.language, use_gpu=not cpu)

    def fit(self, x, y=None):
        return self

    def transform(self, x, y=None):
        result = []
        if self.silent:
            pbar = x
        else:
            pbar = tqdm(x)
        for document in pbar:
            with open('/tmp/stanford_parsing', 'a') as log_fh:
                log_fh.write(document + '/n')
                result.append(self.nlp(document))
        return result


class StanzaToNltkTreesTransformer(BaseEstimator, TransformerMixin):
    """
    This transformer takes documents created by the stanza python package
    and transforms them into dependency trees in nltk format.

    Input: each document is either a stanza Document or Sentence.
    Output: each document is a list of NLTK Trees, one for each sentence.
        If the input document is a sentence, than the result is a list with
        a single entry.

    """

    def __init__(self, node_labels=None):
        """
        Initializes transformer.

        Args:
            node_labels (list of str): how each word is represented. Provide a
                list of ['lemma', 'upos', 'xpos', 'dependency'], e.g.:
                node_contents=['upos', 'dependency']
        """
        if type(node_labels) == str:
            node_labels = [node_labels]
        if not node_labels:
            node_labels = ['upos']
        self.node_labels = node_labels

    def fit(self, x, y=None):
        return self

    def transform(self, x, y=None):
        result = []
        for document in x:
            sentences = []
            if type(document) == Document:
                sentences = document.sentences
            elif type(document) == Sentence:
                sentences = [document]
            doc = []
            for sentence in sentences:
                dependency_tree = self.parse(sentence)
                labeled_tree = self.get_symbols(dependency_tree)
                doc.append(labeled_tree)
            if not doc:
                raise ValueError(
                    f'coud not parse anything from input {document}')
            result.append(doc)
        return result

    def get_symbols(self, tree):
        dependency = tree.label()
        if isinstance(dependency, str):
            return nltk.Tree(dependency, [self.get_symbols(c) for c in tree])
            # root case

        # the dependency consists of three parts:
        _, relationship, word_2 = dependency
        # the first word is already handled in other iterations of this method
        labels = []
        if 'dependency' in self.node_labels:
            labels.append(relationship)
        labels += [_get_word_attribute(word_2, label, fallback='_')
                   for label in self.node_labels
                   if label != 'dependency']
        label = '#'.join(labels)
        result = nltk.Tree(label, [self.get_symbols(c) for c in tree])
        return result

    def get_nodes(self, tree):
        leaves = []
        if isinstance(tree, nltk.Tree):
            if not isinstance(tree.label(), str):
                leaves.append(tree)
            for child in tree:
                leaves += self.get_nodes(child)
        return leaves

    def parse(self, sentence):
        roots = [d for d in sentence.dependencies if d[1] == 'root']
        unprocessed = [d for d in sentence.dependencies if d not in roots]
        # it is possible that multiple words have a root relationship
        tree = nltk.Tree('root', [nltk.Tree(x, []) for x in roots])

        while True:
            changed = False
            added_dependencies = []
            nodes = self.get_nodes(tree)
            for d in unprocessed:
                for leaf in nodes:
                    if hasattr(d[0], 'id'):
                        if leaf.label()[2].id == d[0].id:
                            leaf.append(nltk.Tree(d, []))
                            added_dependencies.append(d)
                            changed = True
                    elif hasattr(d[0], 'index'):
                        if leaf.label()[2].index == d[0].index:
                            leaf.append(nltk.Tree(d, []))
                            added_dependencies.append(d)
                            changed = True
            unprocessed = [x for x in unprocessed
                           if x not in added_dependencies]
            if len(unprocessed) == 0 or not changed:
                break

        return tree
