#
#   2025 Fabian Jankowski
#   Normal - normal mixture model.
#

import logging

import numpy as np
import pymc as pm

from fitpdf.models.model import Model


class NN(Model):
    name = "Normal - normal"

    def __init__(self):
        """
        Model distribution.
        """

        self.__log = logging.getLogger("fitpdf.models.nn")

        self.ncomp = 2

    def __repr__(self):
        """
        Representation of the object.
        """

        info_dict = {"bla": "XXX"}

        info_str = "{0}".format(info_dict)

        return info_str

    def __str__(self):
        """
        String representation of the object.
        """

        info_str = "{0}: {1}".format(self.name, repr(self))

        return info_str

    def get_model(self, t_data, t_offp, params):
        """
        Construct a mixture model.

        Parameters
        ----------
        t_data: ~np.array of float
            The input data to be fit.
        t_offp: ~np.array of float
            The off-pulse data.
        params: dict
            Additional parameters that influence the processing.

        Returns
        -------
        model: ~pm.Model
            A mixture model.
        """

        data = t_data.copy()
        offp = t_offp.copy()

        # on-pulse mean and std
        onp_mean = np.mean(data)
        onp_std = np.std(data)
        print(f"On-pulse mean: {onp_mean:.5f}")
        print(f"On-pulse std: {onp_std:.5f}")

        # off-pulse mean and std
        offp_mean = np.mean(offp)
        offp_std = np.std(offp)
        print(f"Off-pulse mean: {offp_mean:.5f}")
        print(f"Off-pulse std: {offp_std:.5f}")

        # mixture contributions
        if params["weights"] is None:
            _weights = np.array([0.3, 0.7])
        else:
            _weights = np.array(params["weights"])
            self.__log.info(f"Using custom component weights: {_weights}")

        _weights /= np.sum(_weights)
        self.__log.info(f"Mixture weights: {_weights}")

        coords = {"component": np.arange(2), "obs_id": np.arange(len(data))}

        with pm.Model(coords=coords) as model:
            x = pm.Data("x", data, dims="obs_id")

            # mixture weights
            w = pm.Dirichlet("w", a=_weights, dims="component")

            # priors
            mu = pm.Normal(
                "mu",
                mu=np.array([offp_mean, onp_mean]),
                sigma=np.array([0.1, 1.0]) * onp_std,
                dims="component",
            )
            sigma = pm.HalfNormal(
                "sigma",
                sigma=np.array([offp_std, onp_std]),
                dims="component",
            )

            # 1) normal distribution for nulling
            # 2) normal distribution for pulses
            components = pm.Normal.dist(
                mu=mu,
                sigma=sigma,
                shape=(2,),
            )

            pm.Mixture("obs", w=w, comp_dists=components, observed=x, dims="obs_id")

        return model

    def get_mode(self, posterior, icomp):
        """
        Compute the mode of the model component.

        Returns
        -------
        modes: ~np.array of float
            The mode of the model component.
        """

        mu = posterior["mu"]

        if icomp in [0, 1]:
            mode = mu
        else:
            raise NotImplementedError(f"Component not implemented: {icomp}")

        return mode
