from nltk.util import ngrams

import datetime
import dateparser
import pytz
from datetime import datetime as dt
import dateutil.relativedelta
import gensim

from royston.tc_overlap_strategy import OverlapStrategy
from royston.util import normalise

utc = pytz.UTC

DEFAULT_OPTIONS = {
    # a threshold for the minimum number of times a phrase has to occur
    # in a single day before it can even be considered a trend for a given
    # subject.
    # @todo: work out a logical way of calculating this per category.
    "min_trend_freq": 4,
    # the context of the number of days to consider for the history
    "history_days": 90,
    # the number of days over which to check for trends
    "trend_days": 1,
    # the maximum size of the n-gram window
    "max_n": 6,
    # really not sure why I added this...assume it is to handle words that
    # just didn't get mentioned in the history period.
    "history_frequency_tolerance": 1.6,
    # the maximum number of results to return.
    "trends_top_n": 8,
}

DAY_IN_MS = 86400000

DATE_FIELDS = ["start", "end", "history_start", "history_end"]


# @todo: put all of these helpers in a separate file with the trend strategy...
def set_doc_phrases(doc_phrases, docs, phrases):
    """
    helper function for populating doc_phrases, such that doc_phrases[doc] =
    an array of
    phrases that are trending
    """
    for doc in docs:
        if doc not in doc_phrases:
            doc_phrases[doc] = []
        doc_phrases[doc] = doc_phrases[doc] + phrases


class Royston:
    def __init__(self, options={}):
        self.options = {}
        self.set_options(options)
        self.docs = {}
        # initialise the multi-dimensional ngram array storage
        self.ngrams = [[] for _ in range(self.options["max_n"] + 1)]
        # track the usage of the ngrams
        self.ngram_history = {}
        self.last_ingest_id = None
        self.doc2vec_docs = []
        self.doc2vec_tokens = {}

    def clean_date(self, d):
        if isinstance(d, datetime.datetime):
            return d.replace(tzinfo=pytz.UTC)
        return dateparser.parse(d).replace(tzinfo=pytz.UTC)

    def get_history_period(
        self, history_days=None, trend_days=None, start=None
    ):

        if history_days is None:
            history_days = self.options["history_days"]
        if trend_days is None:
            trend_days = self.options["trend_days"]
        if start is None:
            start = (
                dt.now()
                if "start" not in self.options
                else self.options["start"]
            )

        history_end = self.clean_date(start)
        history_start = self.clean_date(
            history_end
            - dateutil.relativedelta.relativedelta(days=history_days)
        )
        return [history_start, history_end]

    # returns the date range for the trend period
    # @todo: why is there so much manky code with this?!
    def get_trend_period(self, trend_days=None):
        if trend_days is None:
            trend_days = self.options["trend_days"]
        return [
            self.clean_date(
                dt.now()
                - dateutil.relativedelta.relativedelta(days=trend_days)
            ),
            self.clean_date(dt.now()),
        ]

    def set_periods(self, history_days=None, trend_days=None):
        [start, end] = self.get_trend_period(trend_days)
        [history_start, history_end] = self.get_history_period(trend_days)

        self.options = {
            **self.options,
            "start": self.clean_date(start),
            "end": self.clean_date(end),
            "history_start": self.clean_date(history_start),
            "history_end": self.clean_date(history_end),
        }

    def set_options(self, options):

        # time periods being reconfigured, but ambigious
        if (
            "start" in options
            and "history_start" not in options
            and "start" in self.options
            and "history_start" in self.options
        ):
            raise Exception(
                "start needs history_start when these have already been set"
            )

        # ensure consistent formatting for incoming dates
        for field in DATE_FIELDS:
            if field in options:
                options[field] = self.clean_date(options[field])

        self.options = {**DEFAULT_OPTIONS, **self.options, **options}

        # none specified - so recalculate @TODO: don't store time, if it's
        # always to be live...
        if "start" not in options:
            self.set_periods()

        if "history_start" not in self.options:
            [history_start, history_end] = self.get_history_period()
            self.options = {
                **self.options,
                **{
                    "history_start": history_start,
                    "history_end": history_end,
                },
            }

    def ingest_ngram(self, ngram, doc, n):
        """
        Add a new ngram into the ramekin.
        """
        # construct the storable ngram object
        # self.ngrams[n].append({
        #    'date': doc['date'], # store this so it can be pruned when old
        #    'ngram': ngram,
        #    'subject': doc['subject'] if 'subject' in doc else None
        # })
        # initialised hash element
        if ngram not in self.ngram_history:
            self.ngram_history[ngram] = {"occurances": []}

        self.ngram_history[ngram]["occurances"].append(
            {"date": doc["date"], "doc_id": doc["id"]}
        )

        # @todo: add this to a queue to look for trends...this.isNGramTrending
        # (ngram, doc);

    def ingest(self, raw_doc):
        """
        Ingest a single document into the collection.

        :param raw_doc: document to ingest, in this format:
        {
            id: <Unique ID - can be any format>,
            body: "Text",
            date: <ISO Date format string, or JavaScript date object>,
            subject: <Any object>
        }
        """
        # preprocess the date to check it's in the right format.
        if "date" not in raw_doc:
            raise Exception("No 'date' field set for document")
        date = self.clean_date(raw_doc["date"])

        if date < self.options["history_start"]:
            return

        # ensure there is an id set
        if "id" not in raw_doc:
            raise Exception("No 'id' field set for document")

        # throw error if the document already exists in the ramekin
        if raw_doc["id"] in self.docs:
            raise Exception(
                "Document %s has already been added." % (raw_doc["id"])
            )
        doc = {**raw_doc, "date": date}

        # we may need to revisit what doc data we store
        self.docs[doc["id"]] = doc

        # generate all the [1...n]-grams for the document
        tokens = normalise(doc["body"])
        for n in range(1, self.options["max_n"] + 1):
            doc_ngrams = ngrams(tokens, n)
            for ngram in doc_ngrams:
                self.ingest_ngram(ngram, doc, n)

        # record the id of the last ingest document
        self.last_ingest_id = doc["id"]

        self.ingest_doc2vec(tokens, doc["id"])
        self.doc2vec_tokens[doc["id"]] = tokens

    def ingest_all(self, docs):
        """
        ingests a set of documents into the current Royston.
        :param {docs}: a set of documents in the format expected format
        """
        for doc in docs:
            self.ingest(doc)

    # @todo: support subjects
    def ingest_doc2vec(self, tokens, doc_id):

        # For training data, add tags
        self.doc2vec_docs.append(
            gensim.models.doc2vec.TaggedDocument(tokens, [doc_id])
        )

    def train_doc2vec(self):
        epochs_count = len(self.doc2vec_docs)
        self.doc2vec_model = gensim.models.doc2vec.Doc2Vec(
            vector_size=500, min_count=2, epochs=epochs_count
        )
        self.doc2vec_model.build_vocab(self.doc2vec_docs)
        self.doc2vec_model.train(
            self.doc2vec_docs,
            total_examples=self.doc2vec_model.corpus_count,
            epochs=epochs_count,
        )

    def used_phrases(self, start, end):
        """
        Finds the phrases used between a particular date range.
        @todo: error handling.
        @todo: pre compute to store used phrases per day/hour
        @todo: this may be the main bottle neck - if a hashmap is created,
        it reduces the searches and just sets the value each time.
        returning just the values (or keys) would be quick??
        """
        used = set()
        for ngram in self.ngram_history:
            for occurance in self.ngram_history[ngram]["occurances"]:
                if occurance["date"] >= start and occurance["date"] < end:
                    used.add(ngram)
        return list(used)

    def prune(self):
        """
        Super simple method that prunes in the following conditions:

        1) before the start of the history (self.options['history_start'])
        2) the phrase was used only in the history period once
        """

        def is_in_range(occurance):
            return occurance["date"] >= self.options["history_start"]

        ngrams = list(self.ngram_history.keys())
        for ngram in ngrams:
            # remove anything before the start of our considered history (i.e.
            # stuff that is two old for us to care about)
            self.ngram_history[ngram]["occurances"] = list(
                filter(is_in_range, self.ngram_history[ngram]["occurances"])
            )
            # @TODO: BUG: this needs to count number of times it happened in
            # the past (OR cron to run only daily - but that's a bit lazy)
            if len(self.ngram_history[ngram]["occurances"]) < 2:
                del self.ngram_history[ngram]

    def find_docs(self, ngram, options):
        """
        Find the documents that contain the specified ngram
        """
        # sanitise input
        options["start"] = self.clean_date(options["start"])
        options["end"] = self.clean_date(options["end"])

        # not found - return nothing
        if ngram not in self.ngram_history:
            return []

        history = self.ngram_history[ngram]

        def matcher(doc):
            full = self.docs[doc["doc_id"]]
            return (
                doc["date"] >= options["start"]
                and doc["date"] < options["end"]
                and (
                    ("subject" not in options)
                    or (
                        "subject" in full
                        and options["subject"] == full["subject"]
                    )
                )
            )

        history_in_range = list(
            filter(lambda doc: matcher(doc), history["occurances"])
        )
        # return just the ids
        return list(map(lambda history: history["doc_id"], history_in_range))

    def count(self, ngram, options):
        """
        Count the number of times that an ngrams has occurred within the
        conditions of the options.

        :param ngram:
        :param options:
        :return int:
        """
        return len(self.find_docs(ngram, options))

    def count_history(self, ngram):

        return self.count(
            ngram,
            {
                "start": self.options["history_start"],
                "end": self.options["history_end"],
            },
        )

    def count_trend_period(self, ngram):
        return self.count(
            ngram, {"start": self.options["start"], "end": self.options["end"]}
        )

    #  change start and end time to be part of options early on...
    def get_ngram_trend(self, ngram, doc_phrases, combined_options):

        """
        Does the trend analysis related to an ngram.

        this needs a proper name and explaination.

        @todo: decouple doc_phrases?
        """

        # score if the phrase has trended in the last 24 hours

        # const trendDocs = self.findDocs(ngram, { start: self.options.start,
        # end: self.options.end })
        trend_docs = self.find_docs(ngram, combined_options)
        trend_range_count = len(trend_docs)
        history_range_count = self.count(
            ngram,
            {
                "start": self.options["history_start"],
                "end": self.options["history_end"],
            },
        )
        history_day_average = (
            self.options["history_frequency_tolerance"]
            * history_range_count
            / self.options["history_days"]
        )

        trend_day_average = trend_range_count / combined_options["trend_days"]
        history_trend_range_ratio = trend_day_average / (
            0.000001 if history_range_count == 0 else history_day_average
        )

        change = trend_day_average - history_day_average
        change_percent = (
            change
            / (0.000001 if history_range_count == 0 else history_day_average)
        ) * 100

        # add in the tolerance

        # if it's above the average
        # history_day_average
        if (trend_range_count > self.options["min_trend_freq"]) and (
            change_percent > 0
        ):
            phrase = {
                "phrases": ngram,
                "score": change_percent
                * len(ngram),  # history_trend_range_ratio * len(ngram),
                "history_range_count": history_range_count,
                "trend_range_count": trend_range_count,
                "trend_day_average": trend_day_average,
                "history_day_average": history_day_average,
                "history_trend_range_ratio": history_trend_range_ratio,
                "docs": trend_docs,
            }
            set_doc_phrases(doc_phrases, trend_docs, [ngram])
            return phrase

        return None

    def rank_trends(self, trends, doc_phrases, top_n):

        # WHAT DOES THIS DO?
        # rank the documents in each cluster, based on the docs etc.
        for trend in trends:
            docs = []
            # for each document in that trend, count the number of phrases
            # that match
            for doc in trend["docs"]:
                # count the number of phrases from the cluster that are
                # that doc
                inter = set(doc_phrases[doc]).intersection(
                    set(trend["phrases"])
                )
                matches = len(inter)
                docs.append({"doc": doc, "matches": matches})

            # sort based on the number of matches
            docs = sorted(docs, key=lambda x: x["matches"], reverse=True)
            # remove unnecessary sort data now it is sorted
            trend["docs"] = [doc["doc"] for doc in docs]

        # trim to just options.trendsTopN
        return trends[0:top_n]

    """
    Validate the trending options, setting defaults where necessary.
    @todo: this whole block is manky and needs a refactor - setup, search and
    cluster
    @todo: Create TrendingStrategy - this is then passed as a param for doing
    the trending this way.
    """

    def trend_phrases(self, options):
        """
        Finds the phrases that have been trending.

        This method does not consider dealing with related phrases and
        overlapping documents.

        """

        # if times ranges aren't specified, calcuate and stash here....
        # brittle - if 'start' is set, assumes all 4 dates are set.
        if "start" not in options:
            self.set_periods()
        combined_options = {**self.options, **options}

        """
        This is the really manky bit of code, that needs separating into a
        helper class just for the trending
        """

        # end of setup

        # start of trending:search
        combined_options = {**self.options, **options}
        start = combined_options["start"]
        end = combined_options["end"]

        # find all the common phrases used in respective subject, over the
        # past day this doesn't filter by subject, but it therefore should
        # get everything, and not be a problem?
        used_phrases = self.used_phrases(start, end)

        # duplicated data used later for sorting
        doc_phrases = {}

        # score each phrase from the trend period compared to it's historic use
        trend_phrases = list(
            map(
                lambda phrase: self.get_ngram_trend(
                    phrase, doc_phrases, combined_options
                ),
                used_phrases,
            )
        )

        # filter out Nones
        trend_phrases = list(
            filter(lambda phrase: phrase is not None, trend_phrases)
        )

        # rank results on their score
        trend_phrases = sorted(
            trend_phrases,
            key=lambda phrase: (-phrase["score"], phrase["phrases"]),
        )

        return [trend_phrases, doc_phrases]

    def trending(self, options={}):

        [trend_phrases, doc_phrases] = self.trend_phrases(options)
        # map to ngram trends
        if trend_phrases is None or len(trend_phrases) == 0:
            return []

        self.train_doc2vec()

        # add in the tokens
        # hash all the doc tokens that we care about
        def doc2vec_distance(trend_i, trend_j):
            distances = []
            for doc_i_id in trend_i["docs"]:
                doc_i_tokens = self.doc2vec_tokens[doc_i_id]
                for doc_j_id in trend_j["docs"]:
                    doc_j_tokens = self.doc2vec_tokens[doc_j_id]
                    sim = self.doc2vec_model.similarity_unseen_docs(
                        doc_i_tokens, doc_j_tokens
                    )
                    distances.append(sim)
            return 1 - (sum(distances) / len(distances))

        # end of trending:search

        # start of trending:cluster

        # run the clustering - find the phrase that is most similar to so many
        # others (i.e. i, where sum(i) = max( sum() )

        cluster_strategy = OverlapStrategy()
        trends = cluster_strategy.cluster(trend_phrases)

        return self.rank_trends(
            trends, doc_phrases, self.options["trends_top_n"]
        )
