import logging
import re
from collections import Counter
from dataclasses import dataclass
from statistics import median
from typing import Counter as CounterType
from typing import List, Tuple, Union

from ..core.data_types import DocResult, Document, Span
from ..core.registry import TaggerRegistry
from ..core.taggers import BaseTagger

REQUIRED_ENGLISH_WORDS = {"the", "be", "to", "of", "and", "that", "have", "with"}
SYMBOLS = {"#", "\u2026"}
BULLET_POINTS = {"*", "-"}


def robust_median(values: List[Union[int, float]]) -> float:
    if not values:
        return 0.0
    return float(median(values))


@dataclass
class GopherAttributes:
    fraction_of_characters_in_most_common_ngram: List[Tuple[int, float]]
    fraction_of_characters_in_duplicate_ngrams: List[Tuple[int, float]]
    character_count: int = 0
    word_count: int = 0
    median_word_length: float = 0.0
    symbol_to_word_ratio: float = 0.0
    fraction_of_words_with_alpha_character: float = 0.0
    required_word_count: int = 0
    fraction_of_lines_starting_with_bullet_point: float = 0.0
    fraction_of_lines_ending_with_ellipsis: float = 0.0
    fraction_of_duplicate_lines: float = 0.0
    fraction_of_characters_in_duplicate_lines: float = 0.0

    def as_spans(self) -> List[Span]:
        spans = []
        spans.extend(
            [
                Span(
                    0,
                    self.character_count,
                    f"fraction_of_characters_in_most_common_{n}grams",
                    v,
                )
                for n, v in self.fraction_of_characters_in_most_common_ngram
            ]
        )
        spans.extend(
            [
                Span(
                    0,
                    self.character_count,
                    f"fraction_of_characters_in_duplicate_{n}grams",
                    v,
                )
                for n, v in self.fraction_of_characters_in_duplicate_ngrams
            ]
        )
        spans.append(
            Span(
                0,
                self.character_count,
                type="character_count",
                score=self.character_count,
            )
        )
        spans.append(Span(0, self.character_count, type="word_count", score=self.word_count))
        spans.append(
            Span(
                0,
                self.character_count,
                type="median_word_length",
                score=self.median_word_length,
            )
        )
        spans.append(
            Span(
                0,
                self.character_count,
                type="symbol_to_word_ratio",
                score=self.symbol_to_word_ratio,
            )
        )
        spans.append(
            Span(
                0,
                self.character_count,
                type="fraction_of_words_with_alpha_character",
                score=self.fraction_of_words_with_alpha_character,
            )
        )
        spans.append(
            Span(
                0,
                self.character_count,
                type="required_word_count",
                score=self.required_word_count,
            )
        )
        spans.append(
            Span(
                0,
                self.character_count,
                type="fraction_of_lines_starting_with_bullet_point",
                score=self.fraction_of_lines_starting_with_bullet_point,
            )
        )
        spans.append(
            Span(
                0,
                self.character_count,
                type="fraction_of_lines_ending_with_ellipsis",
                score=self.fraction_of_lines_ending_with_ellipsis,
            )
        )
        spans.append(
            Span(
                0,
                self.character_count,
                type="fraction_of_duplicate_lines",
                score=self.fraction_of_duplicate_lines,
            )
        )
        spans.append(
            Span(
                0,
                self.character_count,
                type="fraction_of_characters_in_duplicate_lines",
                score=self.fraction_of_characters_in_duplicate_lines,
            )
        )
        return spans


def get_attributes(text: str, ignore_empty_lines: bool = False) -> GopherAttributes:
    attrs = GopherAttributes([], [])
    attrs.character_count = len(text)
    if attrs.character_count == 0:
        return attrs

    try:
        words = text.split()
        word_count = len(words)
        character_count = sum(len(word) for word in words)

        attrs.word_count = word_count
        attrs.median_word_length = robust_median([len(word) for word in words])
        attrs.symbol_to_word_ratio = sum(1 for word in words if any(s in word for s in SYMBOLS)) / max(
            word_count, 1
        )
        attrs.fraction_of_words_with_alpha_character = sum(
            1 for word in words if any(c.isalpha() for c in word)
        ) / max(word_count, 1)
        attrs.required_word_count = sum(1 for word in words if word in REQUIRED_ENGLISH_WORDS)

        all_counts = all_ngram_counts(words)

        count_most_common_ngrams = {2, 3, 4}
        for n, ngram_counts in all_counts:
            if not ngram_counts:
                continue
            if n in count_most_common_ngrams:
                most_common_ngram, count = ngram_counts.most_common(1)[0]
                value = count * sum(len(w) for w in most_common_ngram) / max(character_count, 1)
                attrs.fraction_of_characters_in_most_common_ngram.append((n, value))
            else:
                ng_char_count = sum(count * sum(len(w) for w in ng) for ng, count in ngram_counts.items())
                value = sum(
                    count * sum(len(w) for w in ng) for ng, count in ngram_counts.items() if count > 1
                ) / max(ng_char_count, 1)
                attrs.fraction_of_characters_in_duplicate_ngrams.append((n, value))

        if ignore_empty_lines:
            lines = re.split(r"\n+", text)
        else:
            lines = text.split("\n")

        line_count = len(lines)
        for line in lines:
            if any(line.startswith(s) for s in BULLET_POINTS):
                attrs.fraction_of_lines_starting_with_bullet_point += 1
            if line.endswith("\u2026"):
                attrs.fraction_of_lines_ending_with_ellipsis += 1
        attrs.fraction_of_lines_starting_with_bullet_point /= max(line_count, 1)
        attrs.fraction_of_lines_ending_with_ellipsis /= max(line_count, 1)

        line_counts = Counter(lines)
        attrs.fraction_of_duplicate_lines = sum(count for line, count in line_counts.items() if count > 1) / max(
            line_count, 1
        )
        attrs.fraction_of_characters_in_duplicate_lines = sum(
            len(line) * count for line, count in line_counts.items() if count > 1
        ) / max(character_count, 1)
    except Exception as e:
        logging.exception(f"Error processing text {e}: {text[:200]}")

    return attrs


def all_ngram_counts(words) -> List[Tuple[int, CounterType[Tuple[str, ...]]]]:
    return [(n, Counter(list(zip(*[words[i:] for i in range(n)])))) for n in range(2, 11)]


def all_ngram_counts_alt(words: List[str]) -> List[Tuple[int, CounterType[Tuple[str, ...]]]]:
    """Seems like it should be faster, but isn't"""
    ngram: List[Tuple[str, ...]] = list(zip(words, words[1:]))
    all_counts: List[Tuple[int, CounterType[Tuple[str, ...]]]] = [(2, Counter(ngram))]

    for n in range(3, 11):
        ngram = list(a + (b,) for a, b in zip(ngram, words[n - 1 :]))
        all_counts.append((n, Counter(ngram)))

    return all_counts


@TaggerRegistry.add("gopher_v1")
class GopherTagger(BaseTagger):
    def predict(self, doc: Document) -> DocResult:
        attrs = get_attributes(doc.text)
        result = DocResult(doc=doc, spans=attrs.as_spans())
        return result


@TaggerRegistry.add("gopher_v2")
class GopherTaggerV2(GopherTagger):
    def predict(self, doc: Document) -> DocResult:
        attrs = get_attributes(doc.text, ignore_empty_lines=True)
        result = DocResult(doc=doc, spans=attrs.as_spans())
        return result
