#!/usr/bin/env python
# -*- coding: utf-8 -*

import configparser
import csv
import os
import gc
from utility.XMLProtocol import XMLProtocol
from collections import Counter
from tqdm import tqdm
from sklearn.feature_extraction.text import CountVectorizer
from itertools import groupby, chain
from operator import itemgetter
import locale
locale.setlocale(locale.LC_COLLATE, "C")  # Sets locale to portable "C" locale.


def n_grams(files, group_by_feature="year",
            input_type_name="lemmatized_without_stopwords"):
    """
    Clacluates 1 to 5 grams for given input protocols. Can either handel
    lemmatized or non lemmatized files. Writes the ngrams to a tab separated csv
    file. One row inclueds the ngram, the match count of it, the year or date,
    or rede_id or redner_id. One file per unigram, bigram, trigram etc. per
    group key will be created. (There wil be one file for unigrams starting with
    the letter 'A' one for unigrams starting with 'B' etc.)
    Third parameter is a string set by the user which will be added to
    the file names to help distinguish lemmatized and non lemmatized ngrams etc.
    The more protocols are used as input the more RAM the script needs.
    For all 4106 protocols 32GB of RAM with a 32GB swap file was used!
    """
    config = configparser.ConfigParser()
    config.read("config.ini")
    output_path = config["File paths"]["nlp_output"]
    output_path = os.path.join(output_path, "n-grams")
    if not os.path.exists(output_path):
        os.mkdir(output_path)
    for step in tqdm(range(6)[1:], desc="Current ngram calculating"):
        N_GRAMS = []
        file_name_prefix = str(step) + "_grams"
        counter_vectorizer = CountVectorizer(ngram_range=(step, step),
                                             lowercase=False)
        for file_path in tqdm(sorted(files), desc="File status"):
            xml = XMLProtocol()
            xml.read_xml(file_path)
            feature_year = xml.xml_tree.xpath("@sitzung-datum")[0][-4:]
            feature_mont_year = xml.xml_tree.xpath("@sitzung-datum")[0][-7:]
            speeches = xml.xml_tree.xpath(".//rede | .//sitzungsbeginn")
            for speech in speeches:
                # gets id of current speech
                feature_rede_id = speech.xpath("@id")
                if(len(feature_rede_id) == 0):
                    feature_rede_id = "sitzungsbeginn"
                else:
                    feature_rede_id = feature_rede_id[0]
                # gets id of current speaker
                feature_redner_id = speech.xpath(".//redner/@id")[0]
                # gets speech text from tokenized or lemmatized protocol
                speech_text = speech.xpath("node()[2]")[0]  # gets second child of speech
                if(speech_text.text is not None):
                    tmp_str = speech_text.text

                ngrams = counter_vectorizer.build_analyzer()
                ngrams_list = ngrams(tmp_str)

                if(group_by_feature == "year"):
                    pairs = [(pair,) + (feature_year,) for pair
                             in ngrams_list]
                elif(group_by_feature == "month_year"):
                    pairs = [(pair,) + (feature_mont_year,) for pair
                             in ngrams_list]
                elif(group_by_feature == "speaker"):
                    pairs = [(pair,) + (feature_redner_id,) for pair
                             in ngrams_list]
                elif(group_by_feature == "speech"):
                    pairs = [(pair,) + (feature_rede_id,) for pair
                             in ngrams_list]
                N_GRAMS.extend(pairs)
            speeches = None
        # puts uppercase ngram at first position in line to sort by this
        # will be delted later on
        print("Start counting ngrams.")
        N_GRAMS = Counter(N_GRAMS)
        print("Finished counting ngrams.")
        print("Start sorting ngrams")
        N_GRAMS = [item[0][0][0].upper()
                   + "||"
                   + item[0][0]
                   + "||"
                   + str(item[0][1])
                   + "||"
                   + str(item[1])
                   for item in N_GRAMS.items()]
        N_GRAMS = sorted(N_GRAMS, key=locale.strxfrm)
        print("Finished sorting ngrams")
        # sorts all ngrams into groups one group for each german uppercasse
        # letter except ß
        # Also one group for every decimal from 0 to 10
        # Other non ascii or non decimal ngrams will be sorted in own groups
        # These groups will be joined together later on into one non ascii group
        alphabetically = []
        tmp_list = []
        for letter, entries in tqdm(groupby(N_GRAMS, key=itemgetter(0)),
                                    desc="Grouping ngrams alphabetically"):
            if(letter):
                print(letter)
                for entry in entries:
                    tmp_list.append(entry)
            alphabetically.append(tmp_list)
            tmp_list = []
            N_GRAMS = None
            gc.collect() # frees RAM
        key_list = ([i for i in range(10)]
                    + "A B C D E F G H I J K L M N O P Q R S T U V W X Y Z".split()
                    + ["_Non_ASCII"])
        # groups all non ascii ngrams into one list to save them into one csv
        if(len(alphabetically) > 37):
            joined_tail = alphabetically[36:]
            joined_tail = chain.from_iterable(list(joined_tail))
            del alphabetically[36:]
            alphabetically.append(joined_tail)
        # save groups to individual files
        for group, key in tqdm(zip(alphabetically, key_list),
                               desc="Writing ngrams to files"):
            group_ngrams = [entry.split("||")[1:] for entry in group]
            file_name = (str(key)
                         + "_"
                         + file_name_prefix
                         + "_per_"
                         + group_by_feature
                         + "_"
                         + input_type_name
                         + ".csv")
            file_output_path = os.path.join(output_path, file_name)
            with open(file_output_path, "w", newline="", encoding="utf8") as file:
                writer = csv.writer(file, delimiter="\t")
                writer.writerows(group_ngrams)
        alphabetically = None


if __name__ == '__main__':
    n_grams()
