from scipy.stats import binom as _binom, poisson as _poisson, geom as _geom, expon as _expon
from typing import Union, Any, Iterable, Sequence
import random as _rd
import numpy as _np
import abc as _abc

from coota.sets import *


class Chooser(object):
    @_abc.abstractmethod
    def choice(self, x: Sequence) -> Any:
        """
        This method specifies how the chooser selects an item.
        :param x: The data source from which the chooser chooses.
        :return: A single item chosen from `x`.
        """
        pass

    @_abc.abstractmethod
    def choices(self, x: Sequence, n: int) -> Sequence:
        """
        This method specifies how the chooser selects a batch of items.
        :param x: The data source from which the chooser chooses.
        :param n: The number of choices.
        :return: A list of items chosen from `x`.
        """
        pass


class DefaultChooser(Chooser):
    def choice(self, x: Sequence) -> Any:
        return _rd.choice(x)

    def choices(self, x: Sequence, n: int) -> Sequence:
        return _rd.choices(x, k=n)


class GaussianChooser(Chooser):
    @staticmethod
    def _analyse(x: Sequence) -> [float, float, int]:
        s = 0.5 * len(x)
        return s, s / 3

    def choices(self, x: Sequence, n: int) -> Sequence:
        choices = _np.random.normal(*self._analyse(x), size=n)
        r = []
        for c in choices:
            r.append(x[int(c)])
        return r

    def choice(self, x: Sequence) -> Any:
        return x[int(_np.random.normal(*self._analyse(x), size=1))]


class DiscreteChooser(Chooser):
    @_abc.abstractmethod
    def get_weights(self, length: int) -> list[float]:
        pass

    def choice(self, x: Sequence) -> Any:
        return _rd.choices(population=x, weights=self.get_weights(len(x)), k=1)[0]

    def choices(self, x: Sequence, n: int) -> Sequence:
        return _rd.choices(population=x, weights=self.get_weights(len(x)), k=n)


class BinomialChooser(DiscreteChooser):
    def __init__(self, p: float):
        self.p: float = p
        self._weights: Union[list[float], None] = None

    def get_weights(self, length: int) -> list[float]:
        if self._weights is None:
            self._weights = _binom.pmf(_np.arange(0, length, 1), length, self.p)
        return self._weights


class PoissonChooser(DiscreteChooser):
    def __init__(self, lam: float):
        self.lam: float = lam
        self._weights: Union[list[float], None] = None

    def get_weights(self, length: int) -> list[float]:
        if self._weights is None:
            self._weights = _poisson.pmf(_np.arange(0, length, 1), self.lam)
        return self._weights


class GeomChooser(DiscreteChooser):
    def __init__(self, k: float, p: float):
        self.k: float = k
        self.p: float = p
        self._weights: Union[list[float], None] = None

    def get_weights(self, length: int) -> list[float]:
        if self._weights is None:
            self._weights = _geom.pmf(self.k, self.p)
        return self._weights


class ExponChooser(DiscreteChooser):
    def __init__(self, lam: float):
        self.lam: float = lam
        self._weights: Union[list[float], None] = None

    def get_weights(self, length: int) -> list[float]:
        if self._weights is None:
            self._weights = _expon.pmf(_np.arange(0, length, 1), 1 / self.lam)
        return self._weights


class Association(object):
    def __init__(self, the_other_generator: Any):
        """
        :param the_other_generator: The other generator associated with this generator.
            That generator must generate before this generator.
        """
        if not isinstance(the_other_generator, Generator):
            raise TypeError("'obj' must be a Generator.")
        self._other_generator = the_other_generator

    def get_the_other_generator(self) -> Any:
        """
        :return: The other generator.
        """
        return self._other_generator

    @_abc.abstractmethod
    def associate(self, g: Any, the_other_generator_output: Any) -> Any:
        """
        This method specifies the association between the output of two generators.
        :param g: The generator whose `set_association()` is called with the association given to.
        :param the_other_generator_output: The output generated by the other generator.
        :return: Anything. If not none, the return will be returned by `generate()`, or the process will continue.
        """
        pass


class TheSameAs(Association):
    def associate(self, g: Any, the_other_generator_output: Any) -> Any:
        return the_other_generator_output


class GeneratorOutput(object):
    def __init__(self, output: Any):
        """
        :param output: The output.
        """
        self._output = output

    def get_output(self) -> Any:
        """
        :return: The output.
        """
        return self._output


class Generator(object):
    def __init__(self, *default_args, **args):
        """
        :param default_args: Given to `make()` when no arguments are given to `generate()`.
            For example, in a **GeneratorSequence**, `generate()` is called with no arguments,
            then the `default_args` will be given to `make()`.
            Required arguments are listed in the specific generators.
        :param args: Global arguments for the generator. Required arguments are listed in the specific generators.
        """
        self._parseable: bool = True
        self._uc: bool = True
        self._args: dict[str: Any] = {} if args is None else args
        self._default_args: tuple = default_args
        self._chooser: Chooser = DefaultChooser()
        self._source_cache: Union[Sequence, None] = None
        self._last: Any = None
        self._association: Union[Association, None] = None

    def __str__(self):
        return f"[Generator({self.get_args()})]"

    def _get_source_cache(self) -> Union[Sequence, None]:
        """
        :return: The cache of the source.
        """
        return self._source_cache

    def _set_source_cache(self, source_cache: Sequence) -> None:
        """
        :param source_cache: The cache of the source.
        """
        self._source_cache = source_cache

    def _get_last(self) -> Any:
        """
        :return: Last generated data.
        """
        return self._last

    def _set_last(self, last: Any) -> None:
        """
        :param last: Last generated data.
        :return:
        """
        self._last = last

    def get_parseable(self) -> bool:
        """
        :return: Whether the generator can be parsed as an argument.
        """
        return self._parseable

    def set_parseable(self, parseable: bool) -> None:
        """
        :param parseable: Whether the generator can be parsed as an argument. If false,
            the generator will be recognized as an argument itself instead of being parsed into an actual output.
        :return:
        """
        self._parseable = parseable

    def get_uc(self) -> bool:
        """
        :return: Whether the generator uses the cache of the source.
        """
        return self._uc

    def set_uc(self, use_cache: bool) -> None:
        """
        :param use_cache: Whether the generator uses the cache of the source. If true,
            the generator will only call `source()` once and use the cache instead after that.
            It's true by default. Set it to false if your source is not always static.
        :return:
        """
        self._uc = use_cache

    def get_args(self) -> dict:
        """
        :return: The global arguments of the generator.
        """
        return self._args

    def get_arg(self, name: str, required_type: type = object) -> Any:
        """
        :param name: The argument's name.
        :param required_type: The type of argument you require. If any type of argument is acceptable,
            set it to `object` which is also by default.
        :return: The argument's value. It can be `None` if the argument does not exist
            or is not of the same type as required.
        """
        args = self.get_args()
        if name not in args:
            return None
        arg = args[name]
        return arg if isinstance(arg, required_type) else None

    def get_arg_or_default(self, name: str, default: Any, required_type: type = object) -> Any:
        """
        :param name: The argument's name.
        :param default: The default value.
        :param required_type: The type of argument you require. If any type of argument is acceptable,
            set it to `object` which is also by default.
        :return: The argument's value. The default value will be returned if the argument does not exist
            or is not of the same type as required.
        """
        arg = self.get_arg(name, required_type)
        return default if arg is None else arg

    def get_required_arg(self, name: str, required_type: type = object) -> Any:
        """
        :param name: The argument's name.
        :param required_type: The type of argument you require. If any type of argument is acceptable,
            set it to `object` which is also by default.
        :return: The argument's value.
        :exception AttributeError: The argument does not exist.
        :exception TypeError: The argument is not of the same type as required.
        """
        arg = self.get_arg(name)
        if arg is None:
            raise AttributeError(f"'Args' must contain keys '{name}'.")
        if not isinstance(arg, required_type):
            raise TypeError(f"Expecting '{name}' to be type of {required_type}, but got {type(arg)} instead.")
        return arg

    def get_default_args(self) -> tuple:
        """
        :return: The default arguments to be given to `make()`.
        """
        return self._default_args

    def get_chooser(self) -> Chooser:
        """
        :return: The generator's chooser.
        """
        return self._chooser

    def set_chooser(self, chooser: Chooser) -> None:
        """
        :param chooser: Set the generator's chooser. A **DefaultChooser** is used by default.
            If you want to change the behavior of choosing, such as making it fit a certain distribution,
            you can do so by changing the chooser object.
        :return:
        """
        self._chooser = chooser

    def get_association(self) -> Union[Association, None]:
        """
        :return: The association with the other generator.
        """
        return self._association

    def set_association(self, association: Association) -> None:
        """
        :param association: The association with the other generator.
        :return:
        """
        self._association = association

    def choice(self) -> Any:
        """
        :return: One single item chosen from the source by the chooser.
        """
        return self.get_chooser().choice(self.get_source())

    def choices(self, n: int) -> Sequence:
        """
        :param n: The number of items.
        :return: A batch of items chosen from the source by the chooser.
        """
        s = self.get_source()
        return self.get_chooser().choices(s, n)

    @_abc.abstractmethod
    def source(self) -> Sequence:
        """
        This method specifies what data the generator may generate.
        :return: The original dataset from which the generator selects.
        """
        pass

    def get_source(self) -> Sequence:
        """
        :return: The same as `source()` returns. If `use_cache` is true, returns the source cache instead.
        """
        cache_on = self.get_uc()
        sc = self._get_source_cache()
        if cache_on and sc is None:
            self._set_source_cache(self.source())
        return self._get_source_cache() if cache_on else self.source()

    @_abc.abstractmethod
    def make(self, *args) -> Any:
        """
        This method specifies how to generate data.
        :param args: Optional arguments.
        :return: Anything.
        """
        return self.choice()

    def generate(self, *args, parse: bool = True) -> Any:
        """
        :param args: Optional arguments given to `make()`.
        :param parse: Whether to resolve the generator in parameters and return values. True: return an output.
            False: return the generator itself.
        :return: Anything.
        :exception LookupError: The associated generator generated before the generator to which is associated has
            generated.
        """
        r = None
        try:
            if not parse or not self.get_parseable():
                r = self
                return r
            a = self.get_association()
            if a is not None:
                og: Generator = a.get_the_other_generator()
                ogl = og._get_last()
                if ogl is None:
                    raise LookupError("The associated generator must not generate until the generator to which is "
                                      "associated has generated.")
                r = a.associate(self, ogl)
                if r is not None:
                    return r
            if args == ():
                args = self.get_default_args()
            args = list(args)
            for i in range(len(args)):
                arg = args[i]
                if isinstance(arg, Generator):
                    args[i] = arg.generate(parse=parse)
            r = self.make(*args)
            if isinstance(r, Generator):
                r = r.generate(parse=parse)
            return r
        finally:
            self._set_last(r)

    def output(self, *args, parse: bool = False) -> GeneratorOutput:
        """
        :param args: Optional arguments given to `make()`.
        :param parse: Whether to resolve the generator in parameters and return values.
        :return: The return value of `generate()` wrapped as a **GeneratorOutput**.
        """
        return GeneratorOutput(self.generate(*args, parse=parse))


class ItertableGenerator(Generator):
    def __init__(self, *default_args, **args):
        super().__init__(*default_args, **args)
        self._pointer: Any = 0
        self.initialize()

    def __iter__(self) -> Iterable:
        self.initialize()
        return self

    def __next__(self) -> Any:
        return self.generate()

    def get_pointer(self) -> Any:
        """
        :return: The pointer.
        """
        return self._pointer

    def set_pointer(self, pointer: Any) -> None:
        """
        :param pointer: The pointer which is 0 in default.
        :return:
        """
        self._pointer = pointer

    def choice(self) -> Any:
        """
        :return: `self.get_source()[pointer]`. The pointer must be an integer when calling this method.
        """
        p = self.get_pointer()
        if not isinstance(p, int):
            raise TypeError("When this method is called, the pointer must be of type int.")
        return self.get_source()[p]

    def choices(self, n: int) -> Sequence:
        """
        :param n: the number
        :return: `self.get_source()[pointer: pointer + n]`. The pointer must be an integer when calling this method.
        """
        p = self.get_pointer()
        if not isinstance(p, int):
            raise TypeError("When this method is called, the pointer must be of type int.")
        return self.get_source()[p: p + n]

    @_abc.abstractmethod
    def initialize(self) -> None:
        """
        This is a callback used to set initialization operations such as pointers.
        :return:
        """
        pass

    @_abc.abstractmethod
    def step(self) -> bool:
        """
        This method specifies how the iterator iterates.
        :return: True: continue iteration. False: stop iteration.
        """
        return False

    def generate(self, *args, parse: bool = True) -> Any:
        try:
            return super(ItertableGenerator, self).generate(*args, parse=parse)
        finally:
            if not self.step():
                raise StopIteration


class LetterGenerator(Generator):
    """
    Generate a letter.
    """

    def source(self) -> Sequence:
        return LETTER_SET

    def make(self, *args) -> Any:
        return self.choice()


class StringGenerator(LetterGenerator):
    """
    Generate a string of letters.
    Arguments:
        (Required) length of string
    """

    def make(self, *args) -> Any:
        return "".join(self.choices(args[0]))


class NumberGenerator(StringGenerator):
    """
    Generate a string of numbers.
    Arguments:
        (Required) length of string
    """

    def source(self) -> Sequence:
        return NUMBER_SET

    def make(self, *args) -> Any:
        s = super(NumberGenerator, self).make(*args)
        if len(s) > 1 and s[0] == "0":
            s[0] = self.get_chooser().choice("123456789")
        return s


class LetterAndNumberGenerator(StringGenerator):
    """
    Generate a string of letters and numbers.
    Arguments:
        (Required) length of string
    """

    def source(self) -> Sequence:
        return NUMBER_SET + LETTER_SET


class IntGenerator(Generator):
    """
    Generate an integer x(a ≤ x < b).
    Global Arguments:
        (Required) start: a
        (Required) stop: b
    """

    def source(self) -> Sequence:
        return range(self.get_required_arg("start", required_type=int),
                     self.get_required_arg("stop", required_type=int))

    def make(self, *args) -> Any:
        return self.choice()


class IntIterable(ItertableGenerator):
    def source(self) -> Sequence:
        return self.get_pointer()

    def make(self, *args) -> Any:
        return self.get_source()

    def initialize(self) -> None:
        self.set_uc(False)
        self.set_pointer(self.get_required_arg("start", required_type=int))

    def step(self) -> bool:
        self._pointer += self.get_arg_or_default("step", 1, required_type=int)
        return self._pointer <= self.get_required_arg("stop", required_type=int)
