from __future__ import annotations

from typing import Any, Dict, List, Sequence

import cytoolz as tz
import numpy as np
from ahocorasick import Automaton

from robustnessgym.cachedops.spacy import Spacy
from robustnessgym.core.identifier import Identifier
from robustnessgym.slicebuilders.subpopulation import Subpopulation


class AhoCorasick:
    def __init__(self, *args, **kwargs):
        super(AhoCorasick, self).__init__(*args, **kwargs)

        # Use the Aho-Corasick search algorithm to speed up phrase lookups
        self.automaton = Automaton()

    @classmethod
    def from_phrases(cls, phrases: Dict[Any, str]) -> AhoCorasick:
        # Create a new automaton
        ahocorasick = cls()

        # Add all the phrases we want to search for
        for key, phrase in phrases.items():
            # As values, we add the key of the phrase
            ahocorasick.automaton.add_word(phrase, key)

        # Initialize Aho-Corasick
        ahocorasick.automaton.make_automaton()

        return ahocorasick


class HasPhrase(
    Subpopulation,
    # Spacy
):
    def __init__(
        self, phrases=None, identifiers: List[Identifier] = None, *args, **kwargs
    ):

        super(HasPhrase, self).__init__(
            # One slice per phrase
            identifiers=[
                Identifier(_name=self.__class__.__name__, phrase=phrase)
                for phrase in phrases
            ]
            if not identifiers
            else identifiers,
            *args,
            **kwargs
        )

        # This is the list of phrases that will be searched
        self.phrases = phrases
        if self.phrases is None:
            self.phrases = []

        # Create and populate Aho-Corasick automatons for words and phrases
        self.word_ahocorasick = AhoCorasick.from_phrases(
            {i: phrase for i, phrase in enumerate(self.phrases) if " " not in phrase}
        )
        self.phrase_ahocorasick = AhoCorasick.from_phrases(
            {i: phrase for i, phrase in enumerate(self.phrases) if " " in phrase}
        )

    @classmethod
    def from_file(cls, path: str) -> Subpopulation:
        """Load phrases from a file, one per line."""
        with open(path) as f:
            phrases = [line.strip() for line in f.readlines()]
        return cls(phrases=phrases)

    @classmethod
    def default(cls) -> Subpopulation:
        """A default vocabulary of phrases to search."""
        return cls(phrases=[])

    @classmethod
    def join(cls, *slicers: HasPhrase) -> Sequence[HasPhrase]:
        """Join to combine multiple HasPhrase slicers."""
        return [
            HasPhrase(phrases=list(tz.concat([slicer.phrases for slicer in slicers])))
        ]

    def apply(
        self,
        slice_membership: np.ndarray,
        batch: Dict[str, List],
        columns: List[str],
        *args,
        **kwargs
    ) -> np.ndarray:

        # Use the spacy cache to grab the tokens in each example (for each key)
        tokenized_batch = Spacy.retrieve(
            batch=batch,
            columns=[[key] for key in columns],
            proc_fns="tokens",
        )

        # Search for words
        if len(self.word_ahocorasick.automaton) > 0:
            for key, tokens_batch in tokenized_batch.items():
                for i, tokens in enumerate(tokens_batch):
                    # Get the values (indices) of all the matched tokens
                    matched_indices = [
                        self.word_ahocorasick.automaton.get(token)
                        for token in tokens
                        if self.word_ahocorasick.automaton.exists(token)
                    ]

                    # Fill in the slice labels for slices that are present
                    slice_membership[i, matched_indices] = 1

        # Search for phrases
        if len(self.phrase_ahocorasick.automaton) > 0:
            for key in columns:
                for i, example in enumerate(batch[key]):
                    # Get the values (indices) of all the matched phrases
                    matched_indices = [
                        index
                        for _, index in self.phrase_ahocorasick.automaton.iter(example)
                    ]

                    # Fill in the slice labels for slices that are present
                    slice_membership[i, matched_indices] = 1

        return slice_membership


# class HasAnyPhrase(Subpopulation):
#
#     def __init__(self,
#                  phrases: List[str] = None,
#                  identifier: Identifier = None,
#                  *args,
#                  **kwargs):
#         # Take the union of the phrases
#         subpopulation = Subpopulation.union(
#             HasPhrase(phrases=phrases),
#             identifier=Identifier(
#                 _name=self.__class__.__name__,
#                 phrases=set(phrases),
#             ) if not identifier else identifier
#         )
#
#         super(HasAnyPhrase, self).__init__(
#             identifiers=subpopulation.identifiers,
#             apply_fn=subpopulation.apply,
#             *args,
#             **kwargs
#         )


class HasAnyPhrase(Subpopulation):
    def __init__(
        self,
        phrase_groups: List[List[str]] = None,
        identifiers: List[Identifier] = None,
        *args,
        **kwargs
    ):

        # Keep track of the phrase groups
        self.phrase_groups = phrase_groups

        if identifiers:
            assert len(identifiers) == len(
                phrase_groups
            ), "Must have one identifier per phrase group."

        self.subpopulations = []
        # For every phrase group
        for i, phrases in enumerate(phrase_groups):
            # Take the union of the phrases
            self.subpopulations.append(
                Subpopulation.union(
                    HasPhrase(phrases=phrases),
                    identifier=Identifier(
                        _name=self.__class__.__name__,
                        phrases=set(phrases),
                    )
                    if not identifiers
                    else identifiers[i],
                )
            )

        super(HasAnyPhrase, self).__init__(
            identifiers=list(
                tz.concat(
                    [subpopulation.identifiers for subpopulation in self.subpopulations]
                )
            ),
            *args,
            **kwargs
        )

    def apply(
        self,
        slice_membership: np.ndarray,
        batch: Dict[str, List],
        columns: List[str],
        *args,
        **kwargs
    ) -> np.ndarray:

        # Run all the subpopulations in sequence to update the slice membership matrix
        for i, subpopulation in enumerate(self.subpopulations):
            slice_membership[:, i : i + 1] = subpopulation.apply(
                slice_membership=slice_membership[:, i : i + 1],
                batch=batch,
                columns=columns,
                *args,
                **kwargs
            )

        return slice_membership

    @classmethod
    def join(cls, *slicebuilders: HasAnyPhrase) -> Sequence[HasAnyPhrase]:
        # Join all the slicebuilders
        return [
            HasAnyPhrase(
                phrase_groups=[
                    phrases
                    for slicebuilder in slicebuilders
                    for phrases in slicebuilder.phrase_groups
                ],
                identifiers=[
                    identifier
                    for slicebuilder in slicebuilders
                    for identifier in slicebuilder.identifiers
                ],
            )
        ]


class HasAllPhrases(Subpopulation):
    def __init__(
        self,
        phrase_groups: List[List[str]] = None,
        identifiers: List[Identifier] = None,
        *args,
        **kwargs
    ):

        # Keep track of the phrase groups
        self.phrase_groups = phrase_groups

        if identifiers:
            assert len(identifiers) == len(
                phrase_groups
            ), "Must have one identifier per phrase group."

        self.subpopulations = []
        # For every phrase group
        for i, phrases in enumerate(phrase_groups):
            # Take the union of the phrases
            self.subpopulations.append(
                Subpopulation.intersection(
                    HasPhrase(phrases=phrases),
                    identifier=Identifier(
                        _name=self.__class__.__name__,
                        phrases=set(phrases),
                    )
                    if not identifiers
                    else identifiers[i],
                )
            )

        super(HasAllPhrases, self).__init__(
            identifiers=list(
                tz.concat(
                    [subpopulation.identifiers for subpopulation in self.subpopulations]
                )
            ),
            *args,
            **kwargs
        )

    def apply(
        self,
        slice_membership: np.ndarray,
        batch: Dict[str, List],
        columns: List[str],
        *args,
        **kwargs
    ) -> np.ndarray:

        # Run all the subpopulations in sequence to update the slice membership matrix
        for i, subpopulation in enumerate(self.subpopulations):
            slice_membership[:, i : i + 1] = subpopulation.apply(
                slice_membership=slice_membership[:, i : i + 1],
                batch=batch,
                columns=columns,
                *args,
                **kwargs
            )

        return slice_membership

    @classmethod
    def join(cls, *slicebuilders: HasAllPhrases) -> Sequence[HasAllPhrases]:
        # Join all the slicebuilders
        return [
            HasAllPhrases(
                phrase_groups=[
                    phrases
                    for slicebuilder in slicebuilders
                    for phrases in slicebuilder.phrase_groups
                ],
                identifiers=[
                    identifier
                    for slicebuilder in slicebuilders
                    for identifier in slicebuilder.identifiers
                ],
            )
        ]


# class HasAllPhrases(Subpopulation):
#
#     def __init__(self,
#                  phrases=None,
#                  identifier: Identifier = None,
#                  *args,
#                  **kwargs):
#         # Take the intersection of the phrases
#         subpopulation = Subpopulation.intersection(
#             HasPhrase(phrases=phrases),
#             identifier=Identifier(
#                 _name=self.__class__.__name__,
#                 phrases=set(phrases),
#             ) if not identifier else identifier,
#         )
#
#         super(HasAllPhrases, self).__init__(
#             identifiers=subpopulation.identifiers,
#             apply_fn=subpopulation.apply,
#             *args,
#             **kwargs
#         )


class HasIndefiniteArticle(HasAnyPhrase):
    def __init__(self):
        super(HasIndefiniteArticle, self).__init__(
            phrase_groups=[["a", "an"]],
            identifiers=[Identifier(_name=self.__class__.__name__)],
        )


class HasDefiniteArticle(HasAnyPhrase):
    def __init__(self):
        super(HasDefiniteArticle, self).__init__(
            phrase_groups=[["the"]],
            identifiers=[Identifier(_name=self.__class__.__name__)],
        )


class HasTemporalPreposition(HasAnyPhrase):
    def __init__(self):
        super(HasTemporalPreposition, self).__init__(
            phrase_groups=[["after", "before", "past"]],
            identifiers=[Identifier(_name=self.__class__.__name__)],
        )


class HasPosessivePreposition(HasAnyPhrase):
    def __init__(self):
        super(HasPosessivePreposition, self).__init__(
            phrase_groups=[["inside of", "with", "within"]],
            identifiers=[Identifier(_name=self.__class__.__name__)],
        )


class HasComparison(HasAnyPhrase):
    def __init__(self):
        super(HasComparison, self).__init__(
            phrase_groups=[["more", "less", "better", "worse", "bigger", "smaller"]],
            identifiers=[Identifier(_name=self.__class__.__name__)],
        )


class HasQuantifier(HasAnyPhrase):
    def __init__(self):
        super(HasQuantifier, self).__init__(
            phrase_groups=[["all", "some", "none"]],
            identifiers=[Identifier(_name=self.__class__.__name__)],
        )


class HasNegation(HasAnyPhrase):
    def __init__(self):
        super(HasNegation, self).__init__(
            phrase_groups=[
                [
                    "no",
                    "not",
                    "none",
                    "noone ",
                    "nobody",
                    "nothing",
                    "neither",
                    "nowhere",
                    "never",
                    "hardly",
                    "scarcely",
                    "barely",
                    "doesnt",
                    "isnt",
                    "wasnt",
                    "shouldnt",
                    "wouldnt",
                    "couldnt",
                    "wont",
                    "cant",
                    "dont",
                ]
            ],
            identifiers=[Identifier(_name=self.__class__.__name__)],
        )
