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

"""
This class represents a simple Markov model to generate text.
"""
import json
import random
import re
import typing
import zlib
from pathlib import Path

import requests


class TextGenerator:
    """
    This class represents a simple Markov model to generate text.
    """

    #
    # CONSTRUCTOR
    #

    def __init__(self):
        self._token_ids: typing.Dict[str, int] = {}
        self._rev_token_ids: typing.Dict[int, str] = {}
        self._token_frequency: typing.Dict[int, int] = {}
        self._markov_model: typing.Dict[
            typing.Tuple[int, int], typing.Dict[int, int]
        ] = {}

    #
    # PRIVATE
    #

    def __len__(self):
        return len(self._token_ids)

    @staticmethod
    def _is_valid_token(s: str) -> bool:
        # forbid tokens of length 0
        if len(s) == 0:
            return False
        # forbid tokens that contain non-alphanumeric characters
        if any([not (x.isalnum() or x in ".?!,") for x in s]):
            return False
        # forbid tokens that consist solely of uppercase
        if s.upper() == s:
            return False
        # forbid tokens that consist solely of numbers
        if s.isnumeric():
            return False
        return True

    def _train_using_project_gutenberg(self, gutenberg_url: str) -> "TextGenerator":

        training_text: str = requests.get(gutenberg_url).text

        # trim header:
        header_text: str = "*** START OF THE PROJECT GUTENBERG EBOOK"
        if header_text in training_text:
            start_pos: int = training_text.find(
                "***", training_text.find(header_text) + len(header_text)
            ) + len("***")
            training_text = training_text[start_pos:]
        # trim footer
        footer_text: str = "*** END OF THE PROJECT GUTENBERG EBOOK"
        if footer_text in training_text:
            end_pos: int = training_text.find(footer_text)
            training_text = training_text[0:end_pos]
        # tokenize, build model
        tokens: typing.List[str] = re.split("[ \n\t\r]+", training_text)
        self._token_ids = {}
        self._rev_token_ids = {}
        self._token_frequency = {}
        self._markov_model = {}
        for i in range(0, len(tokens) - 2):

            # get token, set/create ID
            t0: str = tokens[i].strip(" \n\t")
            if not TextGenerator._is_valid_token(t0):
                continue
            t0_id: int = self._token_ids.get(t0, len(self._token_ids))
            self._token_ids[t0] = t0_id

            # get token, set/create ID
            t1: str = tokens[i + 1].strip(" \n\t")
            if not TextGenerator._is_valid_token(t1):
                continue
            t1_id: int = self._token_ids.get(t1, len(self._token_ids))
            self._token_ids[t1] = t1_id

            # get token, set/create ID
            t2: str = tokens[i + 2].strip(" \n\t")
            if not TextGenerator._is_valid_token(t2):
                continue
            t2_id: int = self._token_ids.get(t2, len(self._token_ids))
            self._token_ids[t2] = t2_id

            # update frequency
            self._token_frequency[t0_id] = self._token_frequency.get(t0_id, 0) + 1

            # update markov model
            mm_key: typing.Tuple[int, int] = (t0_id, t1_id)
            if mm_key not in self._markov_model:
                self._markov_model[mm_key] = {}
            self._markov_model[mm_key][t2_id] = (
                self._markov_model[mm_key].get(t2_id, 0) + 1
            )
        # reverse ID map
        self._rev_token_ids = {v: k for k, v in self._token_ids.items()}

        # return
        return self

    #
    # PUBLIC
    #

    def generate(self, min_sentence_length: int = 32) -> str:
        """
        This function generates a string using the statistical model
        stored in this TextGenerator
        :param min_sentence_length:     the minimal sentence length (in tokens)
        :return:                        a string
        """
        # build some text
        sentence_being_built: typing.List[str] = []

        # build text
        while len(
            sentence_being_built
        ) < min_sentence_length or not sentence_being_built[-1].endswith("."):

            # init
            if len(sentence_being_built) == 0:
                seed: typing.Tuple[int, int] = random.choice(
                    [
                        (t0, t1)
                        for (t0, t1) in self._markov_model.keys()
                        if len(self._rev_token_ids[t0]) > 0
                        and self._rev_token_ids[t0][0]
                        == self._rev_token_ids[t0][0].upper()
                    ]
                )
                sentence_being_built.append(self._rev_token_ids[seed[0]])
                sentence_being_built.append(self._rev_token_ids[seed[1]])
                continue

            # default
            t0_id: int = self._token_ids[sentence_being_built[-2]]
            t1_id: int = self._token_ids[sentence_being_built[-1]]

            # build key
            mm_key: typing.Tuple[int, int] = (t0_id, t1_id)
            if mm_key not in self._markov_model:
                sentence_being_built = []
                continue

            # select next element
            nexts: typing.Dict[int, int] = self._markov_model[mm_key]
            ops: typing.List[int] = []
            for k, v in nexts.items():
                for _ in range(0, v):
                    ops.append(k)
            if len(sentence_being_built) < min_sentence_length:
                sentence_being_built.append(self._rev_token_ids[random.choice(ops)])
            else:
                end_sentence_ops: typing.List[int] = [
                    x for x in ops if self._rev_token_ids[x].endswith(".")
                ]
                if len(end_sentence_ops) == 0:
                    sentence_being_built.append(self._rev_token_ids[random.choice(ops)])
                else:
                    sentence_being_built.append(
                        self._rev_token_ids[random.choice(end_sentence_ops)]
                    )

        # return
        return "".join([x + " " for x in sentence_being_built])

    def load(self, json_file: Path) -> "TextGenerator":
        """
        This function loads a TextGenerator from a (zipped) JSON file
        :param json_file:   the location where to load the (zipped) JSON
        :return:            self
        """
        with open(json_file, "rb") as json_file_handle:
            json_obj = json.loads(zlib.decompress(json_file_handle.read()))
        (
            self._token_ids,
            self._rev_token_ids,
            self._token_frequency,
            self._markov_model,
        ) = (
            {k: int(v) for k, v in json_obj["token_ids"].items()},
            {int(k): v for k, v in json_obj["rev_token_ids"].items()},
            {int(k): int(v) for k, v in json_obj["token_frequency"].items()},
            {
                (int(k.split("|")[0]), int(k.split("|")[1])): {
                    int(a): int(b) for a, b in v.items()
                }
                for k, v in json_obj["markov_model"].items()
            },
        )
        return self

    def store(self, file_name: str) -> None:
        """
        This function stores this TextGenerator in a (zipped) JSON format
        :param file_name:   the location where to store the (zipped) JSON
        :return:            None
        """
        with open(file_name, "wb") as json_file_handle:
            json_file_handle.write(
                zlib.compress(
                    bytes(
                        json.dumps(
                            {
                                "token_ids": self._token_ids,
                                "rev_token_ids": self._rev_token_ids,
                                "token_frequency": self._token_frequency,
                                "markov_model": {
                                    str(k[0]) + "|" + str(k[1]): v
                                    for k, v in self._markov_model.items()
                                },
                            }
                        ),
                        encoding="utf8",
                    ),
                    level=9,
                )
            )
