from typing import Literal

import numpy as np


class Likelihood:
    """Define a :py:class:`Likelihood` object.

    This object is used to assess how similar your observed cluster is, stored in
    a :py:class:`asteca.Cluster` object,  compared to a given
    synthetic cluster, generated by the :py:meth:`asteca.Synthetic.generate` method.

    :param my_cluster: :py:class:`asteca.Cluster` object with the
        loaded data for the observed cluster
    :type my_cluster: Cluster
    :param lkl_name: Currently only the Poisson likelihood ratio (``plr``) defined in
        `Tremmel et al. (2013)
        <https://ui.adsabs.harvard.edu/abs/2013ApJ...766...19T/abstract>`__
        is accepted
    :type lkl_name: str
    :param bin_method: Bin method used to split the color-magnitude diagram into cells
        (`Hess diagram <https://en.wikipedia.org/wiki/Hess_diagram>`__); one of:
        ``knuth, blocks, scott, freedman or fixed``. See
        `Choosing Histogram Bins <https://docs.astropy.org/en/stable/visualization/histogram.html>`__
        in `astropy <https://www.astropy.org/>`__ documentation for details on the
        ``knuth, blocks, scott, freedman`` methods.  The method ``fixed``
        uses (15, 10) bins in magnitude and color(s) respectively
    :type bin_method: Literal["blocks", "knuth", "scott", "freedman", "fixed"]

    :raises ValueError: If any of the attributes is not recognized as a valid option
    """

    from .cluster import Cluster

    def __init__(
        self,
        my_cluster: Cluster,
        lkl_name: str = "plr",
        bin_method: Literal["blocks", "knuth", "scott", "freedman", "fixed"] = "knuth",
    ) -> None:
        self.lkl_name = lkl_name
        self.bin_method = bin_method

        likelihoods = ("plr", "bins_distance", "chisq")
        if self.lkl_name not in likelihoods:
            raise ValueError(
                f"'{self.lkl_name}' not recognized. Should be one of {likelihoods}"
            )

        bin_methods = ("blocks", "knuth", "scott", "freedman", "fixed")
        if self.bin_method not in bin_methods:
            raise ValueError(
                f"Binning '{self.bin_method}' not recognized. "
                + f"Should be one of {bin_methods}"
            )

        if my_cluster.mag is None or my_cluster.color is None:
            raise ValueError(
                "This method requires a photometric magnitude and at least one color\n"
                + "defined in the 'Cluster' object"
            )

        # Extract photometry from 'Cluster' object
        self.mag = my_cluster.mag
        self.color = my_cluster.color
        color2 = my_cluster.color2

        from .modules import likelihood_priv as lpriv

        # Data used by the `get()` method
        self.m_ini_idx = 2  # (0->mag, 1->color, 2->mass_ini)
        if color2 is not None:
            self.m_ini_idx = 3  # (0->mag, 1->color, 2->color2, 3->mass_ini)

        colors = [self.color]
        if color2 is not None:
            colors += [color2]

        # Obtain data used by the ``likelihood.get()`` method
        self.ranges, self.Nbins = lpriv.bin_edges_f(bin_method, self.mag, colors)
        self.cl_z_idx, self.cl_histo_f_z = lpriv.lkl_data(
            self.mag, colors, self.ranges, self.Nbins
        )
        self.obs_mag_median = np.nanmedian(self.mag).astype(float)
        self.obs_col_median = np.nanmedian(colors[0]).astype(float)

        # Pre-compute 'tremmel' parameters
        n_colors = len(colors)
        self.bin_sizes = [self.Nbins[0] * self.Nbins[i + 1] for i in range(n_colors)]
        self.offsets = np.array(
            [0] + list(np.cumsum(self.bin_sizes[:-1])), dtype=np.int32
        )
        self.total_bins = sum(self.bin_sizes)

        self.max_lkl = None
        if self.lkl_name == "plr":
            # Evaluate cluster against itself to obtain the maximum likelihood.
            self.max_lkl = self.get(np.array([self.mag, *colors]))
            # This is equal to evaluating
            # max_lkl = sum([gammaln(2*_ + 0.5) - gammaln(_ + 0.5) for _ in cl_histo_f_z])

        print("\nLikelihood object generated")

    def get(self, synth_clust: np.ndarray) -> float:
        """Evaluate the selected likelihood function.

        :param synth_clust: Array containing the synthetic cluster data. The
            shape of this array is assumed to be be: ``[magnitude, color1, (color2)]``,
            where ``magnitude`` and ``color`` are arrays with the synthetic magnitude
            and color photometric data (``color2`` is the optional second color defined)
            If the array contains any extra columns beyond these they will be ignored.
        :type synth_clust: np.ndarray

        :raise ValueError: If the likelihood function is not recognized

        :return: Likelihood value
        :rtype: float
        """
        from .modules import likelihood_priv as lpriv

        # Ignore any extra columns beyond the main photometry
        synth_clust = synth_clust[: self.m_ini_idx]

        if self.lkl_name == "plr":
            return lpriv.tremmel(
                self.ranges,
                self.Nbins,
                self.cl_z_idx,
                self.cl_histo_f_z,
                self.obs_mag_median,
                self.obs_col_median,
                self.offsets,
                self.bin_sizes,
                self.total_bins,
                self.max_lkl,
                synth_clust,
            )
        # if self.lkl_name == "visual":
        #     return lpriv.visual(self, synth_clust)
        # if self.lkl_name == "mean_dist":
        #     return lpriv.mean_dist(self, synth_clust)
        elif self.lkl_name == "bins_distance":
            return lpriv.bins_distance(self.mag, self.color, synth_clust)
        elif self.lkl_name == "chisq":
            return lpriv.chi_square(
                self.ranges, self.Nbins, self.cl_z_idx, self.cl_histo_f_z, synth_clust
            )
        else:
            raise ValueError(f"Likelihood '{self.lkl_name}' not recognized")
