import numpy as np

from butools.ph import CheckPHRepresentation
from butools.ph import SamplesFromPH
from butools.ph import MomentsFromPH
from butools.ph import PdfFromPH
from butools.ph import CdfFromPH
from butools.fitting import PHFromTrace

from prolothar_queue_mining.model.distribution.distribution import Distribution
from prolothar_queue_mining.model.distribution.continuous_distribution import ContinuousDistribution

class PhaseTypeDistribution(ContinuousDistribution):
    """
    a phase type distribution
    http://webspn.hit.bme.hu/~telek/tools/butools/doc/ph.html
    """

    def __init__(self, initial_probabilities: np.ndarray, generator_matrix: np.ndarray,
                 nr_of_buffered_samples: int = 1000, seed: int|None = None,
                 offset: float = 0):
        """
        creates a new phase distribution with given initial probabilities and generator_matrix
        """
        self.__hash = hash((str(initial_probabilities), str(generator_matrix)))
        initial_probabilities = initial_probabilities.reshape((1, -1))
        if not CheckPHRepresentation(initial_probabilities, generator_matrix):
            raise ValueError(f'invalid PH: {initial_probabilities}, {generator_matrix}')
        self.__initial_probabilities = np.matrix(initial_probabilities)
        self.__generator_matrix = np.matrix(generator_matrix)
        self.__nr_of_buffered_samples = nr_of_buffered_samples
        self.set_seed(seed)
        self.__mean, self.__variance = MomentsFromPH(self.__initial_probabilities, self.__generator_matrix, K=2)
        self.__offset = offset

    def get_next_sample(self) -> float:
        try:
            return next(self.__sample_buffer_iterator) - self.__offset
        except StopIteration:
            self.__sample_buffer_iterator = iter(SamplesFromPH(
                self.__initial_probabilities, self.__generator_matrix, self.__nr_of_buffered_samples
            ))
            return next(self.__sample_buffer_iterator) - self.__offset

    def get_mean(self) -> float:
        return self.__mean - self.__offset

    def get_mode(self) -> float:
        raise NotImplementedError()

    def get_variance(self) -> float:
        return self.__variance

    def compute_pdf(self, x: float) -> float:
        return PdfFromPH(self.__initial_probabilities, self.__generator_matrix, [x - self.__offset])[0]

    def compute_cdf(self, x: float) -> float:
        return CdfFromPH(self.__initial_probabilities, self.__generator_matrix, [x - self.__offset])[0]

    def copy(self) -> ContinuousDistribution:
        return PhaseTypeDistribution(
            np.array(self.__initial_probabilities).flatten(),
            np.array(self.__generator_matrix),
            offset=self.__offset, seed=self.seed)

    def __repr__(self):
        return f'PhaseTypeDistribution({self.__initial_probabilities}, {self.__generator_matrix}, seed={self.seed})'

    def set_seed(self, seed: int|None):
        self.seed = seed
        self.__sample_buffer_iterator = iter([])

    def __hash__(self):
        return self.__hash

    def __eq__(self, other):
        return (
            isinstance(other, PhaseTypeDistribution) and
            self.__offset == other.__offset and
            np.array_equal(self.__initial_probabilities, other.__initial_probabilities) and
            np.array_equal(self.__generator_matrix, other.__generator_matrix)
        )

    @staticmethod
    def fit(data: list[float], seed: int = None, orders = 3) -> 'Distribution':
        #make min value large enough for numerical stability
        error = None
        for allowed_min_value in [10, 100, 1000, 10000]:
            try:
                offset = allowed_min_value - min(data)
                alpha, A, _ = PHFromTrace(np.array(data) + offset, orders)
                return PhaseTypeDistribution(alpha, A, offset=offset, seed=seed)
            except FloatingPointError as e:
                error = e
        raise error

    @staticmethod
    def fit_by_mean_and_variance(mean: float, variance: float, seed: float = None) -> 'Distribution':
        raise NotImplementedError()