# coding: utf-8
__version__ = '3.1.1'

import threading
from collections import Counter
from typing import Callable, Dict, List, Optional

from simplebayes.categories import BayesCategories
from simplebayes.constants import CATEGORY_PATTERN
from simplebayes.errors import InvalidCategoryError
from simplebayes.models import CategorySummary, ClassificationResult
from simplebayes.persistence import (
    PERSISTED_MODEL_VERSION,
    dump_model_state,
    load_model_state,
    load_model_state_from_file,
    save_model_state_to_file,
    validate_model_state,
)
from simplebayes.tokenization import create_tokenizer, default_tokenize_text

__all__ = ['SimpleBayes']


class SimpleBayes:
    """A memory-based, optional-persistence naïve bayesian text classifier."""

    def __init__(
        self,
        tokenizer: Optional[Callable[[str], List[str]]] = None,
        alpha: float = 0.0,
        language: str = "english",
        remove_stop_words: bool = False,
    ) -> None:
        """
        :param tokenizer: A tokenizer override. When None, uses built-in tokenizer.
        :param alpha: Laplace smoothing parameter. Use > 0 (e.g. 0.01 or 1.0) to avoid
            zero probabilities for tokens unseen in a category. Default 0 preserves
            prior behavior.
        :param language: Language code for stemmer and stop words (e.g. "english",
            "spanish"). Default "english".
        :param remove_stop_words: If True, filter stop words. Default False (backwards compatible).
        """
        self.categories = BayesCategories()
        self.tokenizer = (
            tokenizer
            or create_tokenizer(language=language, remove_stop_words=remove_stop_words)
        )
        self.alpha = alpha
        self.probabilities = {}
        self._lock = threading.RLock()

    @classmethod
    def tokenize_text(cls, text: str) -> List[str]:
        """
        Default tokenize method; can be overridden

        :param text: the text we want to tokenize
        :type text: str
        :return: list of tokenized text
        :rtype: list
        """
        return default_tokenize_text(text)

    @classmethod
    def count_token_occurrences(cls, words: List[str]) -> Dict[str, int]:
        """
        Creates a key/value set of word/count for a given sample of text

        :param words: full list of all tokens, non-unique
        :type words: list
        :return: key/value pairs of words and their counts in the list
        :rtype: dict
        """
        return dict(Counter(words))

    def flush(self) -> None:
        """
        Deletes all tokens & categories
        """
        with self._lock:
            self.categories = BayesCategories()
            self.probabilities = {}

    def calculate_category_probability(self) -> None:
        """
        Caches the individual probabilities for each category
        """
        with self._lock:
            total_tally = 0.0
            probs = {}
            for category, bayes_category in \
                    self.categories.get_categories().items():
                count = bayes_category.get_tally()
                total_tally += count
                probs[category] = count

            # Calculating the probability
            for category, count in probs.items():
                if total_tally > 0:
                    probs[category] = float(count)/float(total_tally)
                else:
                    probs[category] = 0.0

            new_probabilities = {}
            for category, probability in probs.items():
                new_probabilities[category] = {
                    # Probability that any given token is of this category
                    'prc': probability,
                    # Probability that any given token is not of this category
                    'prnc': 1.0 - probability
                }
            self.probabilities = new_probabilities

    def train(self, category: str, text: str) -> None:
        """
        Trains a category with a sample of text

        :param category: the name of the category we want to train
        :type category: str
        :param text: the text we want to train the category with
        :type text: str
        """
        category = self.normalize_category(category)
        with self._lock:
            try:
                bayes_category = self.categories.get_category(category)
            except KeyError:
                bayes_category = self.categories.add_category(category)

            tokens = self.tokenizer(str(text))
            occurrence_counts = self.count_token_occurrences(tokens)

            for word, count in occurrence_counts.items():
                bayes_category.train_token(word, count)

            # Updating our per-category overall probabilities
            self.calculate_category_probability()

    def untrain(self, category: str, text: str) -> None:
        """
        Untrains a category with a sample of text

        :param category: the name of the category we want to train
        :type category: str
        :param text: the text we want to untrain the category with
        :type text: str
        """
        category = self.normalize_category(category)
        with self._lock:
            try:
                bayes_category = self.categories.get_category(category)
            except KeyError:
                return

            tokens = self.tokenizer(str(text))
            occurrence_counts = self.count_token_occurrences(tokens)

            for word, count in occurrence_counts.items():
                bayes_category.untrain_token(word, count)

            if bayes_category.get_tally() == 0:
                self.categories.delete_category(category)

            # Updating our per-category overall probabilities
            self.calculate_category_probability()

    def classify(self, text: str) -> Optional[str]:
        """
        Chooses the highest scoring category for a sample of text

        :param text: sample text to classify
        :type text: str
        :return: the "winning" category
        :rtype: str
        """
        with self._lock:
            score = self.score(text)
            highest_category, _ = self._find_highest_category(score)
            return highest_category

    def classify_result(self, text: str) -> ClassificationResult:
        """
        Returns structured classification output including score.
        """
        with self._lock:
            scores = self.score(text)
            highest_category, highest_score = self._find_highest_category(scores)
            return ClassificationResult(category=highest_category or None, score=highest_score)

    @classmethod
    def _find_highest_category(cls, scores: Dict[str, float]) -> tuple[Optional[str], float]:
        if not scores:
            return None, 0.0

        highest_category = None
        highest_score = 0.0
        for category in sorted(scores.keys()):
            category_score = float(scores[category])
            if category_score > highest_score:
                highest_score = category_score
                highest_category = category

        return highest_category, highest_score

    def score(self, text: str) -> Dict[str, float]:
        """
        Scores a sample of text

        :param text: sample text to score
        :type text: str
        :return: dict of scores per category
        :rtype: dict
        """
        with self._lock:
            occurs = self.count_token_occurrences(self.tokenizer(text))
            scores = {}
            for category in self.categories.get_categories():
                scores[category] = 0

            categories = self.categories.get_categories().items()

            for word, count in occurs.items():
                token_scores = {}

                # Adding up individual token scores
                for category, bayes_category in categories:
                    token_scores[category] = \
                        float(bayes_category.get_token_count(word))

                # We use this to get token-in-category probabilities
                token_tally = sum(token_scores.values())

                # If this token isn't found anywhere its probability is 0
                if token_tally == 0.0:
                    continue

                # Calculating bayes probability for this token
                # http://en.wikipedia.org/wiki/Naive_Bayes_spam_filtering
                for category, token_score in token_scores.items():
                    # Bayes probability * the number of occurrences of this token
                    scores[category] += count * \
                        self.calculate_bayesian_probability(
                            category,
                            token_score,
                            token_tally
                        )

            # Removing empty categories from the results
            final_scores = {}
            for category, score in scores.items():
                if score > 0:
                    final_scores[category] = score

            return final_scores

    def calculate_bayesian_probability(
        self, cat: str, token_score: float, token_tally: float
    ) -> float:
        """
        Calculates the bayesian probability for a given token/category

        :param cat: The category we're scoring for this token
        :type cat: str
        :param token_score: The tally of this token for this category
        :type token_score: float
        :param token_tally: The tally total for this token from all categories
        :type token_tally: float
        :return: bayesian probability
        :rtype: float
        """
        # P that any given token IS in this category
        prc = self.probabilities[cat]['prc']
        # P that any given token is NOT in this category
        prnc = self.probabilities[cat]['prnc']
        # Laplace smoothing: add alpha to avoid zero probabilities
        # (token_in_cat, token_not_in_cat) -> k=2 for binary view per token
        if self.alpha > 0:
            prtc = (token_score + self.alpha) / (token_tally + 2.0 * self.alpha)
            prtnc = (token_tally - token_score + self.alpha) / (
                token_tally + 2.0 * self.alpha
            )
        else:
            prtnc = (token_tally - token_score) / token_tally
            prtc = token_score / token_tally

        # Assembling the parts of the bayes equation
        numerator = prtc * prc
        denominator = numerator + (prtnc * prnc)

        # Returning the calculated bayes probability unless the denom. is 0
        return numerator / denominator if denominator != 0.0 else 0.0

    def tally(self, category: str) -> int:
        """
        Gets the tally for a requested category

        :param category: The category we want a tally for
        :type category: str
        :return: tally for a given category
        :rtype: int
        """
        with self._lock:
            try:
                bayes_category = self.categories.get_category(category)
            except KeyError:
                return 0

            return bayes_category.get_tally()

    def get_summaries(self) -> Dict[str, CategorySummary]:
        """
        Returns per-category summary details.
        """
        with self._lock:
            summaries: Dict[str, CategorySummary] = {}
            categories = self.categories.get_categories()

            for category_name, category in categories.items():
                category_probability = self.probabilities.get(
                    category_name,
                    {'prc': 0.0, 'prnc': 0.0},
                )
                summaries[category_name] = CategorySummary(
                    token_tally=category.get_tally(),
                    prob_in_cat=float(category_probability['prc']),
                    prob_not_in_cat=float(category_probability['prnc']),
                )

            return summaries

    def save(self, destination) -> None:
        """
        Saves classifier state to a text stream.
        """
        with self._lock:
            dump_model_state(destination, self._export_model_state())

    def load(self, source) -> None:
        """
        Loads classifier state from a text stream.
        """
        with self._lock:
            state = load_model_state(source)
            validate_model_state(state)
            self._apply_model_state(state)

    def save_to_file(self, absolute_path: str = "") -> None:
        """
        Saves classifier state to file using atomic replacement.
        """
        with self._lock:
            save_model_state_to_file(absolute_path, self._export_model_state())

    def load_from_file(self, absolute_path: str = "") -> None:
        """
        Loads classifier state from a persisted model file.
        """
        with self._lock:
            state = load_model_state_from_file(absolute_path)
            validate_model_state(state)
            self._apply_model_state(state)

    @classmethod
    def normalize_category(cls, category: str | None) -> str:
        """
        Validates and normalizes category input.
        """
        if category is None:
            raise InvalidCategoryError("category is required")

        normalized = str(category).strip()
        if not CATEGORY_PATTERN.match(normalized):
            raise InvalidCategoryError(
                "category must be 1-64 chars and only include letters, numbers, underscore, or hyphen",
            )

        return normalized

    def _export_model_state(self) -> Dict:
        categories = {}
        for category_name, category in self.categories.get_categories().items():
            category_tokens = {
                token: int(count)
                for token, count in category.tokens.items()
                if count > 0
            }
            categories[category_name] = {
                "tally": int(category.get_tally()),
                "tokens": category_tokens,
            }

        return {
            "version": PERSISTED_MODEL_VERSION,
            "categories": categories,
        }

    def _apply_model_state(self, state: Dict) -> None:
        self.categories = BayesCategories()
        categories = state.get("categories", {})
        for category_name, category_state in categories.items():
            category = self.categories.add_category(category_name)
            for token, count in category_state["tokens"].items():
                category.train_token(token, count)

        self.calculate_category_probability()
