"""This module defines classes for Bayesian learning of probabilistic models
for forced-choice responses in phoneme-identification or -discrimination tests,
for which results are conventionally presented as confusion matrices.

Recorded experimental data are obtained from a cm_data.StimRespCountSet instance.

*** Class Overview:

ConfMatrixModel: Defines posterior distributions of all response probabilities,
    given all selected data from a phoneme-identification experiment.
    Learned from a single cm_data.StimRespCountSet instance

ResponseProbGroup: Container for all individual response probability models,
    for ONE group of test participants.

ResponseProbIndividual: Posterior distribution of phoneme response probabilities for
    ONE RESPONDENT in all selected test conditions,
    given response counts and a single ResponseProbPopulation instance,
    common for all participants recruited from the same population.

ResponseProbPopulation: Joint posterior distribution of
    response probabilities and concentration parameters
    (for random individuals)
    in the population from which ONE group of test subjects was recruited.
    ResponseProbPopulationPoint, ResponseProbPopulationSampled are two variants.

ResponseProbPopulationMean: posterior distribution of the mean response probabilities
    for the population.

PopulationConcPrior: Vague prior distribution of concentration parameters,
    common for ALL responses and ALL populations.


*** Model Theory: See
A. Leijon, G. E. Henter, and M. Dahlquist (2016).
Bayesian analysis of phoneme confusion matrices.
*IEEE Trans Audio, Speech, and Language Proc* 24(3):469–482.
doi: 10.1109/TASLP.2015.2512039.

Arne Leijon (2025): ConfMatrixCalc — Bayesian Analysis of
Phoneme Confusion Matrices. Unpublished doc report.


*** Version History:

* Version 0.8:
2026-02-24, skip rng_special, using std np.random.default_rng()
2026-02-21, New class PopulationModelMixture, merged subset of group_factor "dimensions".
            Modified ResponseProbModel.rvs() to allow arbitrary number of samples
2026-02-18, ConfMatrixModel.learn lets user choose sampled / point population mode
            lets user set random-generator seed for exactly reproducible results
2026-02-07, using module rng_special for Dirichlet sampling
2026-01-27, finally got ResponseProbPopulationSampled.learn working
2025-11-20, cleanup comments, new class ResponseProbPopulationMean,
            for the distribution of the mean response probabilities in the population.
2025-11-19, simplified Berger reference prior, same for all response elements
2025-11-15, defined first version of Berger reference prior for concentration params
2025-11-02, indexing of concentration parameters and response counts
            changed to allow different response sets for each stimulus.
2025-10-17, fix to agree with ConfMatrixFrame internal changes
2025-10-08, mutual_information_ds, prob_correct_ds -> pandas.Series result
2025-10-08, rvs_df -> pandas.DataFrame format for sampling results

* Version <= 0.7:
2020-04-02, ResponseProbPopulationSampled validated like IEEE ASLP paper
Works OK, but still slightly underestimated range of population MI distribution, like point-model.
Should test with even less informative population prior.
Minor code and comments cleanup.

2020-03-31,
Corrected ResponseProbPopulationSampled, NOT YET validated by simulations *******
ResponseProbPopulation = ResponseProbPopulationPoint by default, as before.
PopulationConcPrior.create() depending on point- vs sample-estimated population model.
Minor code and comments cleanup.

* PyPi Version 0.7.1:
Using the model structure described in Leijon et al. (2016):
Individual response probabilities are Dirichlet-distributed,
and the population distribution is also learned as a single Dirichlet distribution,
represented by class ResponseProbPopulationPoint.
The prior for the population concentration parameters is regularized
by element-wise independent (broad) gamma distributions.
Some additional code cleanup

2018-08-11, first functional version
"""
# *** Try FeatureSpace model again, like MatLab version 2013 ? ***

import numpy as np
import logging
# from itertools import product
from scipy.optimize import minimize  #, Bounds
from scipy.special import gammaln, psi

# from scipy.optimize import approx_fprime, check_grad  # only for test

import pandas as pd

from samppy import hamiltonian_sampler as ham

ham.VECTOR_AXIS = 1

logger = logging.getLogger(__name__)
# logger.setLevel(logging.DEBUG)  # **** TEST
# ham.logger.setLevel(logging.DEBUG)  # **** TEST


# ---------------------- Default settings for PopulationConcGamma
POPULATION_PRIOR_PSEUDO_COUNT = 1.
# = sum of prior mode across all response alternatives, i.e.,
#   POPULATION_PRIOR_MODE = POPULATION_PRIOR_PSEUDO_COUNT / n_resp
POPULATION_PRIOR_MEAN = 10.
# = mean of prior gamma distribution of concentration parameters
# used and validated in Leijon et al. (2016) for point-estimated population model.
# validated 2026 by simulations

# ---------------------- Settings for TEST of other PopulationConcGamma variants
POPULATION_PRIOR_MODE = 0.5
# used only if POPULATION_PRIOR_PSEUDO_COUNT is None
# = fixed mode of prior gamma distribution for EACH ELEMENT of concentration vectors
#   used in original version Leijon et al. (2016)
POPULATION_PRIOR_SHAPE = 0.5
# used only if both POPULATION_PRIOR_PSEUDO_COUNT and POPULATION_PRIOR_MODE are None

# ---------------------- Setting for PopulationConcBerger, NOT USED
# BERGER_REFERENCE_KNEE = 1.  # knee point in Berger density approximation

# ---------------------- Sampling and point-estimation parameters
SAMPLING_EPSILON = 0.1
# = initial stepsize epsilon in Hamiltonian sampler for ResponseProbPopulationSampled

DEFAULT_N_SAMPLES = 1000  # 1000 ?
# = default number of samples of random concentration parameter array
#   including one conc vector for each stim in each test condition
# *** allow user to set this easily?

CONC_MIN = 1.e-6  # lower limit of concentration values
# -----------------------------------------------------------------

__version__ = '2026-02-24'

# ------------------------------------------------------------------
class ConfMatrixModel:
    """Defines posterior distributions of response probabilities
    for all groups, subjects, test conditions, and test stimuli,
    given all included data from a phoneme-identification experiment.
    Class-method learn creates an instance
    from raw phoneme-identification data
    stored in a cm_data.StimRespCountSet instance.
    """
    # *** store conc and counts as dict (stim_group, values) like SimPopulation,
    #   avoiding cmf.gen_stim_resp_slices? NOT for now...

    #   Allow user to define seed for reproducible results ?

    def __init__(self, cmf, groups, rng):
        """
        :param cmf: single cm_data.ConfMatrixFrame instance,
            common for all subjects in all groups.
        :param groups: dict with group models, stored as
            groups[group_ID] = one ResponseProbGroup instance
            groups[group_ID].subjects[subject_ID] = a ResponseProbIndividual instance
            where
            group_ID is a unique tuple (of cmf.group_factor categories),
            subject_ID is a string-valued key,
                    MUST be unique within each group,
                    (but should usually be unique across all groups)
        :param rng: (subclass of) np.Random.Generator
        """
        self.cmf = cmf
        self.groups = groups
        self.rng = rng

    def __repr__(self):
        return self.__class__.__name__ + '(cmf=cmf, groups=groups)'

    @classmethod
    def learn(cls, ds,
              seed=None,
              sampled_conc=False,
              n_samples=DEFAULT_N_SAMPLES):
        """Create and learn ResponseProbGroup instances,
        one for each subject group,
        from a given StimRespCountSet instance
        :param ds: a single cm_data.StimRespCountSet instance,
            including observed response counts for all test subjects
            in all test conditions.
        :param seed: (optional) seed for common np.Random.Generator instance
        :param sampled_conc: (optional) boolean switch between
            PopulationResponseModelPoint vs. PopulationResponseModelSampled
        :param n_samples: (optional) integer n of concentration vector samples
            in PopulationResponseModelSampled
        :return: a single cls instance containing all response-probability models
        """
        global ResponseProbPopulation  # *** need global?
        rng = np.random.default_rng(seed)  # common Generator for all models
        ds.check_complete()
        # ensure that subjects have sufficiently complete data
        cmf = ds.cmf
        population_prior = PopulationConcPrior.create(ds)
        # = element-wise prior used for all groups, test conditions, and stimuli
        groups = dict()
        if sampled_conc:
            ResponseProbPopulation = ResponseProbPopulationSampled
        else:
            ResponseProbPopulation = ResponseProbPopulationPoint
        for (g, g_subjects) in ds.gen_response_counts():
            # g_subjects is generator yielding tuples (s_key, s_df),
            # where s_df is pd.DataFrame with counts in fixed format, same for all subjects
            s_cm = {s: s_df.to_numpy()
                    for (s, s_df) in g_subjects}
            if sampled_conc:
                delay_msg = 'Might take a few minutes...'
            else:
                delay_msg = 'Might take a few seconds...'
            logger.info(f'Learning group {repr(g)} with {len(s_cm)} participants. ' + delay_msg)
            groups[g] = ResponseProbGroup.learn(cmf, s_cm,
                                                population_prior,
                                                rng=rng,
                                                n_samples=n_samples)
        return cls(cmf, groups, rng)

    def gen_population_models(self, gf):
        """Generator of all population models that matches requested group_factor(s)
        :param gf: tuple with one or more group "dimension" labels, as defined in self.cmf.group_factors
        :return: iterator yielding tuples (g_label, g_model), where
            g_label = tuple of group category labels, one for each given "dimension"
                in the order as given in gf
            g_model = a PopulationModelMixture instance for the corresponding subpopulation,
                OR a ResponseProbGroup instance (incl. participants).

        Method: IF gf does not include ALL group "dimension" defined in self.cmf.group_factors,
        g_model is a weighed average across models in all OTHER "dimensions".
        """
        # *** new in v 0.8:;
        for g_label in self.cmf.gen_population_keys_in_factor(gf):
            # g_label is a tuple with zero, one or more group_factor dimensions
            # in arbitrary order, as requested for display.
            g_keys = list(self.cmf.gen_matching_group_keys(dict(zip(gf, g_label))                                                               ))
            # = all corresponding key tuple(s) with standard order as used in self.groups.keys()
            g_keys = list(filter(lambda gk: gk in self.groups.keys(), g_keys))
            # filtered, because some of the possible group categories may not exist
            if len(g_keys) == 1:
                yield g_label, self.groups[g_keys[0]]  # no merging needed, only perhaps reordered key
                # NOTE: complete ema_group.EmaGroupModel instance, including all Participants
            elif len(g_keys) > 1:
                g_weights = self.cmf.population_model_weights(g_keys)
                # len(g_weights) == len(g_keys)
                g_models = [self.groups[g_k].population for g_k in g_keys]
                # NOTE: only ema_group.PopulationMddel instances, without Participants
                yield g_label, PopulationModelMixture(g_models, g_weights,
                                                      cmf=self.cmf, rng=self.rng)
            # else:  len(g_keys) == 0: StopIteration automatically


# -------------------------------------------------------------------
class ResponseProbGroup:
    """Container of ResponseProbIndividual instances
        for all test subjects in ONE group,
    AND also including a ResponseProbPopulation instance adapted to the group data
    """
    def __init__(self, cmf, subjects, population, rng):
        """
        :param cmf: a ConfMatrixFrame instance.
        :param subjects: dict with (subject_id, ResponseProbIndividual) elements
        :param population: a single ResponseProbPopulation instance
        :param rng: common random Generator instance
        """
        self.cmf = cmf
        self.subjects = subjects
        self.population = population
        self.rng = rng

    def __repr__(self):  # ******** more informative ? ******'
        return (self.__class__.__name__ + '('+
                '\n\tcmf=cmf,' +
                '\n\tsubjects=subjects'
                f'\n\tpopulation={repr(self.population)})')

    @classmethod
    def learn(cls, cmf, s_cm, population_prior, rng, n_samples):
        """Create one ResponseProbIndividual instance for each subject,
        and learn a single ResponseProbPopulation prior for the population
        represented by the subjects, in ONE group of subjects.
        :param cmf: a cm_data.ConfMatrixFrame instance with experimental layout
        :param s_cm: dict with elements (s_key, cm), where cm = 2D array,
            cm[t, sr] = response count for each (stim, resp) pair in t-th test condition,
            stored with each row vector including all stimuli, indexed as
            cm[t, s_ind], for s_ind in cmf.gen_sr_slices(), where
                s_ind is a slice for ONE stimulus.
        :param population_prior: externally pre-calculated single prior,
            common for all population models
        :param rng: common random Generator instance
        :param n_samples: only needed for sampled ResponseProbPopulation
        :return: a ResponseProbGroup instance
        """
        population = ResponseProbPopulation.learn(r_count=[cm for cm in s_cm.values()],
                                                  prior=population_prior,
                                                  cmf=cmf,
                                                  rng=rng,
                                                  n_samples=n_samples,
                                                  )
        subjects = {s: ResponseProbIndividual(cmf, cm, population, rng)
                    for (s, cm) in s_cm.items()}
        return cls(cmf, subjects, population, rng)


# ------------------------------------------------------------------
class ResponseProbModel:
    """Superclass for models that can generate
    random samples of response-probability arrays
    from a posterior distribution of
    response probabilities in phoneme-identification tests,
    either for ONE respondent, or for ONE population
    """
    def __init__(self, cmf, rng):
        """
        :param cmf: a single cm_data.ConfMatrixFrame instance, common for all models
        :param rng: common random Generator instance
        """
        self.cmf = cmf
        self.rng = rng

    def conc(self):
        """Concentration parameters for Dirichlet distribution.
        Abstract method, must be implemented by subclass
        :return: 3D array a, stored as
            a[n, t, sr] = n-th sample of concentration parameter vector
            for cmf.stim[s] -> response cmf.resp[r]
            in t-th test condition, for each allowed pair (s, r),
            indexed by s_ind in self.cmf.gen_sr_slices() for each stim.
        """
        raise NotImplementedError

    def conc_df(self):
        """Concentration params in DataFrame format
        :return: a pd.DataFrame instance with
            index = (sample, test_cond), as defined by self.cmf.rvs_index()
            columns = (stim_group, stim, resp), as defined by self.cmf.stim_resp_index()
        """
        a = self.conc()
        return pd.DataFrame(a.reshape((-1, a.shape[-1])),
                            columns=self.cmf.stim_resp_index,
                            index=self.cmf.rvs_index(len(a)))

    def rvs(self, n_samples):
        """Generate random arrays of conditional response probabilities
        sampled from the Dirichlet distribution for each stimulus in each test condition.
        :param n_samples: desired number of sample arrays
        :return: u = 3D array of equally probable sample arrays, with
            u[n, t, sr] = n-th sample of conditional probability
            for r-th response, given s-th stimulus,
            in t-th test condition, for each allowed pair (s, r),
            indexed by sr in self.cmf.gen_sr_slices() for each stim.
            u.shape[0] == n_samples (approximately)
        """
        a = self.conc()
        # a[i, t, sr] = i-th sample of concentration vector
        # for responses to s-th stimulus in t-th test condition
        nc = a.shape[0]  # number of conc samples
        # There may be just ONE or MANY concentration-vector samples.
        if nc == 1:
            u = np.concatenate([np.stack([self.rng.dirichlet(a_tsr, size=n_samples)
                                          for a_tsr in a[0, :, s_ind]], axis=1)
                                for s_ind in self.cmf.gen_sr_slices()],
                               axis=-1)
            # NOTE: np.random.Generator.dirichlet can only use 1D conc vector
            return u
        ns = self.rng.multinomial(n_samples, pvals=np.ones(nc) / nc)
        # ns[i] = number of samples to be drawn with conc = a[i]
        u = [np.concatenate([np.stack([self.rng.dirichlet(a_tsr, size=n_i)
                                       for a_tsr in a_i[:, s_ind]], axis=1)
                             for s_ind in self.cmf.gen_sr_slices()],
                            axis=-1)
             for (a_i, n_i) in zip(a, ns) if n_i > 0]
        # each u[i][n, t, :] = n-th sample of all responses in t-th test condition for all ns[i] > 0
        # u[i].shape == (ns[i], *a.shape[1:]) for all i if ns[i] > 0
        return np.concatenate(u, axis=0)

    def rvs_df(self, n_samples):
        """Generate conditional response probabilities in pd.DataFrame Format,
        sampled from the Dirichlet distribution for each stimulus in each test condition.
        :param n_samples: desired number of sample arrays
        :return: pd.DataFrame with MultiIndex (_sample, *cmf.test_factors.keys(), Stim)
            and one column for each element in cmf.resp
        """
        u = self.rvs(n_samples)
        # = 3D array of equally probable sample arrays, with
        # u[ni, t, gsr]= ni-th sample of gsr-th (stim_group, stim, resp) in t-th test cond.
        df = pd.DataFrame(u.reshape((-1, u.shape[-1])),
                          columns=self.cmf.stim_resp_index,  # 3-level pd.MultiIndex for (sg, s, r)
                          index=self.cmf.rvs_index(len(u)))  # MultiIndex (_sample, test-conditions)
        return df


# -------------------------------------------------------------------
class PopulationModelMixture(ResponseProbModel):
    """A mixture including more than one population model,
    used only to display results merged across some group "dimension"
    """
    def __init__(self, comp, w, **kwargs):
        """
        :param comp: sequence of ResponseProbModel (subclass) instances
        :param w: array with weights; len(w) == len(comp)
        :param kwargs: dict for superclass
        """
        super().__init__(**kwargs)
        self.comp = comp
        w = np.asarray(w)
        w /= np.sum(w)  # normalized
        self.w = w

    @property
    def population(self):
        return self     # already just a population model, no subjects

    def random_ind(self):
        """Model for an UNSEEN RANDOM INDIVIDUAL in this merged population
        """
        return self

    def mean(self):
        """Distribution of the POPULATION MEAN (among all random individuals)
        IF and ONLY IF self.comp can generate (sampled) mean distribution.
        :return: ResponseProbPopulationMean instance, if possible
        """
        comp = [c.mean() for c in self.comp]
        # IFF they have mean(); else NotImplementedError
        w = self.w / np.max(self.w)
        a = [c.a if w_c == 1. else self.rng.choice(c.a, size=int(w_c * c.a.shape[0]),
                                                   replace=False)
             for (c, w_c) in zip(comp, w)]
        a = np.concatenate(a, axis=0)
        self.rng.shuffle(a, axis=0)
        return ResponseProbPopulationMean(cmf=self.cmf, a=a,
                                          prior=comp[0].prior,
                                          rng=self.rng)

    def rvs(self, n_samples):
        """Generate random arrays of conditional response probabilities
        sampled from each model in self.comp, with probabilities == self.w
        :param n_samples: desired number of sample arrays
        :return: u = 3D array of equally probable sample arrays, with
            u[n, t, sr] = n-th sample of conditional probability
            for r-th response, given s-th stimulus,
            in t-th test condition, for each allowed pair (s, r),
            indexed by sr in self.cmf.gen_sr_slices() for each stim.
            u.shape[0] == n_samples (perhaps approximately)
        """
        ns = self.rng.multinomial(n_samples, self.w)
        # ns[m] = n of samples to be generated from self.comp[m]
        u = [comp_m.rvs(n_samples=n_m)
             for (comp_m, n_m) in zip(self.comp, ns)]
        u = np.concatenate(u, axis=0)
        self.rng.shuffle(u, axis=0)
        return u


# -------------------------------------------------------------------
class ResponseProbIndividual(ResponseProbModel):
    """A ResponseProbIndividual instance represents the distribution of
    response probabilities in phoneme-identification tests
    for ONE listener in all test conditions.
    The vector of response probabilities for each presented phoneme category
    is a conditional Dirichlet RV,
    given observed counts and prior (=population) concentration param.

    The population distribution of concentration vectors may be
    EITHER a ResponseProbPopulationPoint instance,
        defining only a single MAP-estimated concentration vector for each stim and condition,
    OR a ResponseProbPopulationSampled instance,
        providing a large set of concentration-vector samples representing the distribution.

    Observed response counts are modeled by a multinomial distribution,
    P(x[t, s, r] propto U[t, s, r]**x[t, s, r],
    for the r-th response to the s-th stimulus in the t-th test condition,
    given a probability-mass row vector U[t, s],
    with elements U[t, s, r] = P(R=r | t, S=s)

    The probability density for U[t, s] is modelled as a Dirichlet distribution
    p_U(u) propto prod_r u_r**(a_r - 1), with posterior concentration parameters
    a_r = x_r + population.a_r, where x_r=x[t, s, r] is the observed response count.
    """
    def __init__(self, cmf, cm, population, rng):
        """
        :param cmf: reference to a ConfMatrixFrame instance,
            common for all model instances.
        :param cm: 2D array or array-like list of response counts, stored as
            cm[t, s_ind] = 1D vector with number of responses
            to s-th stim in t-th test condition, with
            s_ind in cmf.gen_sr_slices() = slice of indices for s-th stim.
        :param population: ResponseProbPopulation instance, common for all subjects,
            defining the probability distribution of concentration parameters
            in the population from which this subject was recruited.
        :param rng: common Generator instance
        """
        super().__init__(cmf=cmf, rng=rng)
        self.cm = cm
        self.population = population

    def __repr__(self):
        return (self.__class__.__name__ + '('+
                '\n\tcmf=cmf,' +
                '\n\tpopulation=population,' +
                '\n\tcm=cm')

    def conc(self):
        """Posterior distribution of concentration parameters,
        represented by an array of samples.
        Used by super().rvs
        :return: 3D array a, stored as
            a[n, t, s_ind] = n-th sample of concentration parameter
            for probability distribution of
            U[t, s_ind] = vector of response probabilities in t-th test condition
            for all allowed events cmf.stim[s] -> r_j for j in s_ind, with
            s_ind in self.cmf.gen_sr_slices() indicating allowed responses for s-th stim.
        """
        return self.population.conc() + self.cm


# ------------------------------------------ population models:

class ResponseProbPopulationPoint(ResponseProbModel):
    """Posterior distribution of concentration vectors in a population,
    used as concentration-vector prior for individual models,
    AND for predictive population response probabilities
    for random individuals in the population,
    for all stimulus phonemes, in all test conditions.

    This version uses a POINT estimate for each concentration vector,
    determined by MAP type II, given all observed response counts.

    This population model was described and validated in Leijon et al. (2016).
    """
    def __init__(self, cmf, a, prior, rng):
        """
        :param cmf: a ConfMatrixFrame instance, needed by super-class.
        :param a: 2D array or array-like list of concentration parameters
            a[t, s_ind] for the Dirichlet probability of all allowed responses
            to the s-th stim in the t-th test condition
        :param prior: a PopulationConcPrior instance.
        :param rng: common random Generator instance
        """
        super().__init__(cmf=cmf, rng=rng)
        self.a = np.asarray(a)
        self.prior = prior

    def __repr__(self):
        return (self.__class__.__name__ + '(' +
                f'\n\ta={repr(self.a)},' +
                f'\n\tcmf={repr(self.cmf)})')

    def random_ind(self):
        """Model for an UNSEEN RANDOM INDIVIDUAL in the population
        """
        return self

    def mean(self):
        """Distribution of the POPULATION MEAN (among all random individuals)
        :return: None, because only sampled subclass can do this
        """
        logger.warning('*** ' + self.__class__.__name__ + ' cannot show mean distribution')
        raise NotImplementedError  # return None #???

    def conc(self):
        """Posterior distribution of concentration parameters,
        represented by array self.a
        Used by super().rvs()
        :return: 3D array self.a, stored as
            a[0, t, s_ind] = ONE POINT estimate of concentration parameter vector
            for s-th stimulus phoneme in t-th test condition
            The singleton first axis is included only for
            compatibility with ResponseProbPopulationSampled class.
        """
        return self.a[np.newaxis, ...]

    @classmethod
    def learn(cls, cmf, r_count, prior, rng, n_samples=1):
        """Learn a MAP point estimate of concentration parameters,
        given an array of observed response counts.

        The probability of ONE observed response-count row vector x is
        P(x | u, a) propto prod_i u_i**(x_i+a_i), where
            a_i is the i-th element of the concentration vector a,
            x is the observed vector of response counts.
        By integrating over u (Leijon et al, 2016) we get
        prob(x, a) propto prior.pdf(a) * B(a+x) / B(a) propto pdf(a | x), where
        B(a) = multivariate beta function = normalization factor of Dirichlet distribution,
        defined by module function _log_multi_beta().
        This log_likelihood(a | x) is used for MAP type II estimation.

        Reference: Leijon et al.(2016).

        :param cmf: a ConfMatrixFrame instance
        :param r_count: 3D array-like list of response counts, stored as
            r_count[l, t, s_ind] = vector with number of responses to s-th stim
            in t-th test condition by l-th subject of a group, for all
            (s_ind, c_s) in cmf.gen_sr_slices()
        :param prior: single PopulationConcPrior instance, with methods conc_ll, d_conc_ll
        :param rng: common random Generator instance
        :param n_samples: NOT USED, using MAP point estimate
        :return: a single cls instance adapted to input data
        """
        def conc_one_tc(x_lr):
            """Learn a conc vector for all stim in ONE caller-selected test condition.
            :param x_lr: 2D array with response counts
                x_lr[l, s_ind] for l-th listener responses to s-th stim
            :return: a_t = conc vector for caller-selected test condition
                a_t.shape == x_lr.shape[-1:]
            """
            a_t = np.empty(x_lr.shape[-1])
            # for (s_ind, c_s) in cmf.gen_sr_slices():
            for s_ind in cmf.gen_sr_slices():
                a_t[s_ind] = cls.conc_point_estimate(x_lr[:, s_ind], prior)
            return a_t
        # ---------------------------------------------

        x = np.array(r_count).transpose((1, 0, 2))
        # x[t, l, s_ind] = response counts for l-th listener to s-th stim in t-th test condition
        a = [conc_one_tc(x_t)
             for x_t in x]
        # *** MAP type II done independently for all individual (t, s) cases
        # *** This is OK for the point estimate, but NOT for a sampled version.
        return cls(cmf, a, prior, rng)

    @classmethod
    def conc_point_estimate(cls, x_lr, prior):
        """Point estimate of ONE concentration row vector,
        for ONE stim category, in ONE test condition.
        common for all subjects.
        :param x_lr: 2D array of response counts,
            x_lr[l, j] = count for j-th allowed response of l-th listener,
            for ONE externally selected stimulus and test condition.
        :param prior: object defining ln p(a)
        :return: a = 1D vector of MAP-optimized concentration params
            len(a) == x_lr.shape[-1]
        """
        def neg_ll(a, x_lr):
            """
            negative log-likelihood for a single tentative concentration ROW vector
            :param a: 1D row vector, with concentration parameters for
                ONE stimulus in ONE test condition
                a[i] = concentration param vector for i-th response probability
            :param x_lr: 2D array with response counts
                x_lr[l, r] = response count for r-th response by l-th listener
                x_lr.shape[-1] = len(a)
            :return: nll = scalar neg. log-likelihood values
            """
            nll = - (_conc_ll(a, x_lr) +
                     np.sum(prior.conc_ll(a)))
            return nll

        def grad_neg_ll(a, x_lr):
            """gradient of neg_ll
            """
            d_nll = - (_grad_conc_ll(a, x_lr) +
                       prior.d_conc_ll(a))
            return d_nll
        # ----------------------------------------------------------

        n_resp = len(x_lr[0])  # = x_lr.shape[-1]
        bounds = [(CONC_MIN, None) for _ in range(n_resp)]
        # to prevent inf log-likelihood
        res = minimize(fun=neg_ll,
                       jac=grad_neg_ll,
                       args=(x_lr,),
                       bounds=bounds,
                       x0=np.ones(n_resp) * prior.mode)
        if res.success:
            return res.x
        else:
            logger.error(f'ResponseProbPopulation.learn minimize res= {res}')
            raise RuntimeError('MAP search did not converge. *** Should not happen ***')


# ------------------------------------------------------------------
class ResponseProbPopulationSampled(ResponseProbPopulationPoint):
    """Posterior distribution of concentration vectors for
    Dirichlet distribution of phoneme response probabilities,
    AND for response probabilities, given self.conc()
    for all stimulus phonemes, in all test conditions.

    This version uses a SAMPLED distribution of log-concentration parameters,
    calculated by Hamiltonian sampling, given all observed response counts.
    """
    def conc(self):
        """Posterior distribution of concentration parameters,
        represented by sample array self.a
        :return: 3D array self.a, stored as
            a[n, t, sr] = n-th sample of concentration parameter vector
            for sr-th stimulus-response pair in t-th test condition.
            The sr indexing is defined by self.cmf.stim_index()
        """
        return self.a

    def mean(self):
        """Create subclass instance representing distribution
        of population mean response probabilities
        :return: a ResponseProbPopulationMean instance
            with self properties, but overridden rvs_df method
        """
        return ResponseProbPopulationMean(self.cmf, self.a, self.prior, self.rng)

    @classmethod
    def learn(cls, cmf, r_count, prior, rng, n_samples=DEFAULT_N_SAMPLES):   # version 2026-01-27
        """Learn sampled distribution of concentration parameters,
        given an array of observed response counts.

        The posterior distribution of population concentration parameters
        is defined in doc report sec:PosteriorPopulationModel, eq:PosteriorAgivenX:

        p(a | x) propto prod_t prod_s p(a_ts) prod_l C(a_ts) / C(x_lts + a_ts), where
        C(z) = Gamma(sum_r z_r) / prod_r Gamma(z_r)
            is the Dirichlet normalization function for ONE concentration row vector z.
        p(a_ts) is the PRIOR density for concentration row vector a_ts for each t,s.

        Thus, disregarding the normalization constant,
        ln p(a | x) = sum_t sum_s ln p(a_ts) + L ln C(a_ts) - sum_l=1^L ln C(x_lts + a_ts),
        where L = number of listeners == len(x)

        The Hamiltonian sampling yields a set of samples of conc vectors
        with approximately equal probability mass for each sample.

        The sampling is done jointly for all test conditions and all (s, r) pairs.
        Therefore, the random concentration array A is temporarily stored as
        a single row vector, reshaped after sampling.

        Since some conc elements have very small magnitude, a < 0.0001,
        while other are much larger, a > 10 or so,
        it is numerically difficult to sample the concentration values as such.

        Instead, the sompling is done in the log domain, using the log-likelihood function
        L(ln_a) = ln p(exp(ln_a) | x), and its gradient w.r.t. ln_a.
        After sampling, all log-samples are transformed back to a = exp(ln_a)

        NOTE: This is NOT equivalent to sampling from the distribution of Y = ln A.
        We sample from the density p_A(a), only in a different representation.

        :param cmf: a ConfMatrixFrame instance
        :param r_count: 3D array-like with observed individual response counts, stored as
            x[l, t, sr] = number of responses cmf.resp[r] to cmf.stim[s]
            in t-th test condition by l-th listener recruited from current population.
            sr is an index list (or slice) for each allowed (s,r) tuple.
            This indexing method allows different response sets for each stim category.
        :param prior: a PopulationConcPrior instance with methods conc_ll and d_conc_ll,
            giving ln p(a[..., t, sr], c_s) ELEMENT-WISE, and gradient, for each row vector a_tsr,
            indexed with the same flat index sr for each allowed (s, r) tuple,
            as defined by cmf.stim_resp_index()
        :param n_samples: (optional) number of concentration array samples
        :return: a single cls instance adapted to input data, defined by
            a = a 3D array of concentration parameters, indexed as
            a[n, t, sr] for the n-th sample in the t-th test condition for the sr-th stim-resp pair.
            a.shape == (n_samples, *r_count.shape[1:])

            The t indexing is defined by cmf.test_cond_index().
            The sr indexing is defined by cmf.stim_resp_index,
            and help method cmf.gen_sr_slices()
        """
        def neg_ll(ln_a):
            """negative log-likelihood for an array of log-concentration ROW vectors
            Samples are represented by the LOG of concentration params.
            :param ln_a: 2D (or 1D) array, with tentative log-concentration-parameter samples for
                all stimuli in all test conditions
                ln_a[n, tsr] transformed to a[n, t, sr] = n-th sample concentration param,
                indexed by (t, sr) just like r_count
            :return: nll = 1D array with neg. log-likelihood values for each conc vector
                nll[n] = - ln p( exp(ln_a[n, :])| x]
                nll.shape == ln_a.shape[:-1]
            """
            a = np.exp(ln_a).reshape((-1, *r_count.shape[1:]))
            # a[n, t, sr] = n-th tentative sample for pair (s, r) in t-th test condition
            ll = np.sum([_conc_ll(a[..., s_ind], r_count[..., s_ind])  # already summed over l, sr
                         + np.sum(prior.conc_ll(a[..., s_ind]), axis=-1)  # = 2D array ln_p[n, t]
                         # for (s_ind, c_s) in cmf.gen_sr_slices()],
                         for s_ind in cmf.gen_sr_slices()],
                        axis=0)
            # ll[n, t] = log-likelihood for n-th sample in t-th test cond
            ll = np.sum(ll, axis=-1)  # sum over test conditions
            # ll[n] = ln p(exp(a[n] )| x) = log density for A[n], except normalization constant
            return - ll

        def grad_neg_ll(ln_a):
            """gradient of neg_ll w.r.t. ln_a elements
            :param ln_a: as for neg_ll
            :return: d_nll = negative gradient vectors
                d_nll[n, i] = - d ln p( ln_a[n, :] ) / d ln_a[n, i]
                d_nll.shape == ln_a.shape
            """
            a = np.exp(ln_a).reshape((-1, *r_count.shape[1:]))
            d_nll = np.zeros_like(a)
            # for (s_ind, c_s) in cmf.gen_sr_slices():
            for s_ind in cmf.gen_sr_slices():
                d_nll[..., s_ind] -= (_grad_conc_ll(a[..., s_ind],
                                                    r_count[..., s_ind])  # summed across l
                                      + prior.d_conc_ll(a[..., s_ind]))
            # = - d ln p(a | self) / d a
            d_nll *= a
            # = - d ln p(a | self) / d ln_a
            d_nll = d_nll.reshape(ln_a.shape)
            return d_nll
        # ----------------------------------------------------------------------
        r_count = np.asarray(r_count)
        a0 = super().learn(cmf, r_count, prior, rng).a
        # = start at MAP type II point estimate **** use own criteria, NOT super() ?
        ln_a0 = np.tile(np.log(a0).reshape((-1,)), (n_samples, 1))
        near_zero = np.log(CONC_MIN)
        near_inf = np.log(1000.)
        bounds = [(near_zero, near_inf) for _ in range(ln_a0.shape[-1])]
        # limit Hamiltonian samples to avoid numerical overflow
        ln_a0 += SAMPLING_EPSILON * rng.standard_normal(size=ln_a0.shape)  # initial dither
        ln_a0 = np.maximum(ln_a0, near_zero)
        ln_a0 = np.minimum(ln_a0, near_inf)
        # *** testing grad_neg_ll ************** :
        # test_ln_a = ln_a0[0]
        # logger.debug('approx gradient test_neg_ll = ' +
        #              np.array_str(approx_fprime(test_ln_a, neg_ll,
        #                                         epsilon=1e-6),
        #                           precision=3))
        # logger.debug('exact  gradient = ' +
        #              np.array_str(grad_neg_ll(test_ln_a),
        #                           precision=3))
        # err = check_grad(neg_ll, grad_neg_ll, test_ln_a, epsilon=1e-6)
        # logger.debug('check_grad err = ' + str(err))

        logger.debug(f'Sampling {ln_a0.shape[0]} samples of vectors with {ln_a0.shape[-1]} elements. Patience!')
        sampler = ham.HamiltonianBoundedSampler(x=ln_a0,
                                                fun=neg_ll,
                                                jac=grad_neg_ll,
                                                epsilon=SAMPLING_EPSILON,
                                                bounds=bounds,
                                                rng=rng,
                                                )
        a = np.exp(sampler.safe_sample(n_samples=n_samples, min_steps=5))
        # = back-transformed to concentration vectors, equally probable in this domain
        a = a.reshape(n_samples, *a0.shape)
        logger.debug(f'sampler accept_rate = {sampler.accept_rate:.1%}, ' +
                     f'n_steps = {sampler.n_steps:.0f}, ' +
                     f'epsilon = {sampler.epsilon:.4f}')
        logger.debug(f'conc.min= {np.amin(a0):.3g} -> {np.amin(a):.3g}. ' +
                     f'max= {np.amax(a0):.3f} -> {np.amax(a):.3f}. ' +
                     f'mean= {np.mean(a0):.3f} -> {np.mean(a):.3f}')
        return cls(cmf, a, prior, rng)


# ResponseProbPopulation = ResponseProbPopulationSampled
ResponseProbPopulation = ResponseProbPopulationPoint  # 2026-02-11: default
# = Default model class for population: but user may override with sampled_comc=True


class ResponseProbPopulationMean(ResponseProbPopulationSampled):
    """Posterior distribution of MEAN response probabilities, given self.conc()
    for all stimulus phonemes, in all test conditions.
    Method: E(U | a) = a / sum_k a_k, for Dirichlet U with concentration vector a
    Thus, the distribution of the mean is defined only for sampled population model.
    """
    def rvs_df(self, n_samples):
        """Generate conditional response probabilities in pd.DataFrame Format,
        sampled from the Dirichlet distribution for each stimulus in each test condition.
        :param n_samples: desired number of sample arrays  *** NOT USED ***
        :return: pd.DataFrame with MultiIndex (_sample, *cmf.test_factors.keys(), Stim)
            and one column for each element in cmf.resp
        """
        # u = np.maximum(self.conc(), np.finfo(float).tiny)  # cannot allow any zero conc
        # already assured by superclass learn method
        u = self.conc().copy()
        # u[i, t, sr] = i-th conc sample, not yet normalized
        # for responses to s-th stimulus in t-th test condition, for allowed (s, r) pairs
        # *** allow n_samples <> conc.shape[0] = module DEFAULT_N_SAMPLES ?
        for s_ind in self.cmf.gen_sr_slices():
            u[..., s_ind] /= np.sum(u[..., s_ind], axis=-1, keepdims=True)
            # = normalized response-prob across responses for each stim
        return pd.DataFrame(u.reshape((-1, u.shape[-1])),
                            columns=self.cmf.stim_resp_index,  # pd.MultiIndex for (stim_group, s, r)
                            index=self.cmf.rvs_index(len(u)))  # MultiIndex (_sample, test-conditions)


# ----------------------------------------- Help function for PopulationConcGamma
def prior_gamma_param(n_resp):
    """Gamma shape and rate parameters, for given number of response categories
    for testing various prior settings.
    :param n_resp: float or integer approx number of response categories
    :return: tuple (a, b)
    """
    # mainly to allow testing of varioius prior settings:
    if POPULATION_PRIOR_PSEUDO_COUNT is None:
        if POPULATION_PRIOR_MODE is None:
            a = POPULATION_PRIOR_SHAPE
            b = POPULATION_PRIOR_SHAPE / POPULATION_PRIOR_MEAN
            return a, b
        else:
            mode = POPULATION_PRIOR_MODE
            # = fixed mode for ALL elements regardless of n_resp
            # Default version validated in Leijon et al. (2016)
    else:  # default method:
        mode = POPULATION_PRIOR_PSEUDO_COUNT / n_resp
        # 2026: mode for ALL conc elements, depending on number of responses
    # mode = (a-1) / b
    # mean = a / b
    # 1./ b = mean - mode
    b = 1. / (POPULATION_PRIOR_MEAN - mode)
    a = POPULATION_PRIOR_MEAN * b
    return a, b


# ----------------------------------------- Population Prior
class PopulationConcGamma:
    """Broad prior distribution of concentration parameters,
    same for all groups, test-conditions and stimuli
    Implemented as a single scalar gamma(a, b) distribution
    with parameters determined by function prior_gamma_param
    """
    def __init__(self, a, b):
        """Scalar prior for all concentration params
        :param a: scalar gamma distribution shape >= 0
        :param b: scalar gamma distribution inverse scale  > 0
        """
        self.a = a
        self.b = b

    @classmethod
    def create(cls, ds):
        """Make gamma-distribution prior for each element of population concentration vector.
        :param ds: a cm_data.StimRespCountSet instance with raw count data, NOT USED
            allowed just in case future version might depend on number of observations
        :return: a cls object
        """
        n_resp = np.mean([sg.shape[1]
                         for sg in ds.cmf.stim_groups.values()])
        # = average number of response alternatives across stim groups
        # *** needed for 2026 default prior setting:
        (a, b) = prior_gamma_param(n_resp)
        logger.debug(f'Using gamma prior with a={a:.5f}, b={b:.6f}')
        return cls(a, b)

    @property
    def mode(self):
        return max(0., self.a - 1.) / self.b

    @property
    def mean(self):
        return self.a / self.b

    def __repr__(self):
        return (self.__class__.__name__ + '(' +
                ', '.join(f'{key}={repr(v)}'
                             for (key, v) in vars(self).items())
                + ')')

    def conc_ll(self, x):
        """NON-normalized *element-wise* logpdf of concentration parameters
        of Dirichlet distribution of response probabilities,
        for ONE stim category
        :param x: array or array-like list of tentative concentration values
            x.shape[-1] == number of all allowed responses to the given stim
        :return: ll = log-likelihood
            ll.shape == x.shape  **** should sum over axis=-1 already here ? *******
        """
        # *** params to be determined on the fly, depending on x.shape[-1] ? ***
        return (self.a - 1.) * np.log(x) - self.b * x

    def d_conc_ll(self, x):
        """Derivative of conc_ll() w.r.t. x
        :param x: array or array-like list of tentative concentration values
        :return: array d_ll, with
            d_ll[...] = d conc_ll(x_[...]) / d x[...]
            d_ll.shape == x.shape
        """
        return (self.a - 1.) / x - self.b


# class PopulationConcBerger:
#     """Broad prior distribution of concentration parameters,
#     common to all populations, test-conditions and stimuli.
#
#     Implemented as
#     p(a) = (1/2) r a^{-1/2} (a + r)^{-3/2}
#     This density is proper, but the mean is undefined.
#
#     p(ln_a) = a p(a) has a well-defined mode at
#     argmax_a p(ln_a) = r/2.
#
#     This is an approximation of the hierarchical reference prior
#     for the Dirichlet-multinomial application, as defined in
#     Proposition 4.2 of
#     Ref: J. O. Berger, J. M. Bernardo, and D. Sun. Overall objective priors.
#         Bayesian Analysis 10:189–246, 2015.
#     The derivation is examined in detail in
#     Leijon: Variational Bayesian Mixture Model Learning.
#         Doc report for ProbModels/Bayes/Mixture.
#
#     Version 2025-11-15: Tested version with non.uniform r values,
#         r = 1 for correct response, r = 1 /(K-1) otherwise.
#         But this led to strong over-estimation of MI result with sparse count data.
#     Version 2025-11-19: For now, we choose
#         r = 1 for ALL elements.
#         *** Best value might depend on number of observed data counts / n response categories?
#     """
#     @classmethod
#     def create(cls, ds):
#         """Make gamma-distribution prior for each element of population concentration vector.
#         :param ds: a cm_data.StimRespCountSet instance with raw count data
#             *** not used ***
#         :return: a cls object
#         """
#         return cls(r=BERGER_REFERENCE_KNEE)
#
#     def __init__(self, r):
#         """Scalar prior for all concentration params
#         :param r: scalar parameter for approximation to Berger's reference prior
#         """
#         self.r = r
#
#     @property
#     def mode(self):
#         """Mode = argmax_a p(ln_a)
#         for concentration element corresponding to CORRECT response
#         :return: scalar
#         """
#         return self.r / 2
#
#     @property
#     def mean(self):
#         return np.nan  # mean is undefined
#
#     def __repr__(self):  # ********** simplify
#         return self.__class__.__name__ + f'(mode p(log conc)= {self.mode}, mean={self.mean})'
#
#     def conc_ll(self, x):  # 2025-11-19
#         """NON-normalized ELEMENT-WISE logpdf of concentration parameters
#         of Dirichlet distribution of response probabilities,
#         for ONE (any) stim category
#         :param x: array or array-like list of tentative concentration values
#             x.shape[-1] == number of allowed responses to the given stim
#         :return: ll = log-likelihood ln p(x)
#             ll.shape == x.shape  **** should sum over axis=-1 already here ? *******
#         """
#         x = np.asarray(x)
#         return - np.log(2.) - (np.log(x) + 3. * np.log(x + self.r)) / 2
#
#     def d_conc_ll(self, x):  # 2025-11-19
#         """Derivative of conc_ll() w.r.t. x
#         :param x: array or array-like list of tentative concentration values
#         :return: array d_ll, with
#             d_ll[...] = d conc_ll(x_[...]) / d x[...]
#             d_ll.shape == x.shape
#         """
#         # r = self._conc_r_approx(x, correct_i)
#         return - (1. / x + 3. / (x + self.r)) / 2


# PopulationConcPrior = PopulationConcBerger    # tested 2025-11
PopulationConcPrior = PopulationConcGamma  # 2025-12-19 default


def version():
    if ResponseProbPopulation is ResponseProbPopulationPoint:
        return __version__ + ' with MAP point-estimated conc parameters'
    else:
        return __version__ + ' with sampled conc parameters'
        # if PopulationConcPrior is PopulationConcBerger:
        #     return v + 'with Berger reference prior'
        # else:
        #     return v + 'with Gamma prior'


# --------------------------------------------- module help functions:

def _log_multi_beta(a):
    """log-partition function, ln B(a), for Dirichlet distribution,
    given concentration vector
    Ref: B(a) = 1 / C(a), defined in Leijon et al, 2016, Eq. B.3
    :param a: tentative array of concentration parameters,
        a[..., :] = ...-th concentration row vector for ...-th Dirichlet distribution
    :return: ln_B = array,
        ln_B[...] = log-partition function for Dirichlet-distribution with conc a[..., :]
        ln_B.shape == a.shape[:-1]
    """
    return np.sum(gammaln(a), axis=-1) - gammaln(np.sum(a, axis=-1))


def _d_log_multi_beta(a):
    """Gradient of _log_multi_beta
    :param a: tentative array of concentration parameters,
        a[..., :] = ...-th concentration row vector
    :return: g = gradient with elements
        g[..., i] = d _log_multi_beta(a[..., :]) / d a[..., i]
        g.shape == a.shape
    """
    return psi(a) - psi(np.sum(a, axis=-1, keepdims=True))


def _conc_ll(a, x):
    """log-likelihood of Dirichlet concentration parameter vector
    of a Dirichlet-multinomial model of observed response counts
    for ONE stim category, for one or many samples / test conditions
    :param a: mD array of concentration parameters
        a[..., i] = concentration param for i-th response in ...-th sample / test condition
    :param x: 2D or 3D array with response counts
        x[l, ..., i] for i-th response by l-th listener in ...-th test condition, if any
        broadcast-compatible with a
        x.shape[-1] == a.shape[-1]; if x is 3D: x.shape[-2] == a.shape[-2]
    :return: LL = log-likelihood for a, given x, summed across listeners
        LL[..., t] = sum_l ln p(a[..., t, :] | x[l, t, :]), except normalization factor
        LL.shape == a.shape[:-1]

    Method: Listeners' responses are assumed conditionally independent,
    given a fixed but unknown response probability for each listener,
    modelled by one independent Dirichlet distribution for each listener.
    Therefore, Dirichlet-multinomial log-likelihood values are summed across listeners.
    Reference: Leijon et al., 2016, Eq. B.10, and
        doc report sec:PosteriorPopulationModel, eq:PosteriorAgivenX
    """
    return (np.sum([_log_multi_beta(a + x_l) for x_l in x], axis=0)
            - len(x) * _log_multi_beta(a))


def _grad_conc_ll(a, x):
    """gradient of _conc_ll w.r.t a, given same observed count x
    for ONE stim category
    :param a: array of concentration parameters
        a[..., t, i] = concentration param for i-th response in t-th test condition
    :param x: 2D or 3D array with response counts
        x[l, t, i] for i-th response by l-th listener in t-th test condition
        broadcast-compatible with a
        x.shape[-1] == a.shape[-1]; if x is 3D: x.shape[-2] == a.shape[-2]
    :return: d_ll = array with elements
        d_ll[..., t, i] = sum_l d _conc_ll(a[..., t :], x) / d a[l, t, i]
        d_ll.shape == a.shape
    """
    return (np.sum([_d_log_multi_beta(a + x_l) for x_l in x], axis=0)
            - len(x) * _d_log_multi_beta(a))


# ------------------------------------------------- TEST:
if __name__ == '__main__':
    from scipy.optimize import approx_fprime, check_grad
    # import matplotlib.pyplot as plt

    # -------------------------------- check _d_log_multi_beta
    n_r = 5
    a = 0.5 * np.ones(n_r)
    a = 1.+ np.arange(10)

    print('*** Checking _d_log_multi_beta:')

    print(f'_log_multi_beta({a}) = {_log_multi_beta(a)}')

    eps = 1.e-6

    print('approx gradient = ', approx_fprime(a, _log_multi_beta, epsilon=1e-6))
    print('exact  gradient = ', _d_log_multi_beta(a))

    err = check_grad(_log_multi_beta, _d_log_multi_beta, a, epsilon=1e-6)
    print('check_grad err = ', err)

    # -----------------------------------------------
    print('\n*** Checking _grad_conc_ll:')
    nx = 5
    x = np.random.default_rng().integers(low=0, high=10, size=(nx, len(a)))

    def test_conc_ll(a):
        return _conc_ll(a, x)

    def test_grad_conc_ll(a):
        return _grad_conc_ll(a, x)

    print('approx gradient _conc_ll = ', approx_fprime(a, test_conc_ll, epsilon=1e-6))
    print('exact  gradient = ', test_grad_conc_ll(a))
    err = check_grad(test_conc_ll, test_grad_conc_ll, a, epsilon=1e-6)
    print('check_grad err = ', err)


    # -----------------------------------------------
    # print('\n*** Checking Prior variant PopulationConcBerger.d_conc_ll:')
    # x = np.random.default_rng().integers(low=0, high=10, size=len(a))
    # pBerger = PopulationConcBerger(r=BERGER_REFERENCE_MODE * 2)
    # correct_i = 1
    #
    # def test_conc_ll(a):
    #     return np.sum(pBerger.conc_ll(a + x), axis=-1)
    #
    # def test_grad_conc_ll(a):
    #     return pBerger.d_conc_ll(a + x)
    #
    # print('approx gradient conc_ll = ', approx_fprime(a, test_conc_ll, epsilon=1e-6))
    # print('exact  gradient = ', test_grad_conc_ll(a))
    # err = check_grad(test_conc_ll, test_grad_conc_ll, a, epsilon=1e-6)
    # print('check_grad err = ', err)


    # -----------------------------------------------
    print('\n*** Checking Prior variant PopulationConcGamma.d_conc_ll:')
    x = np.random.default_rng().integers(low=0, high=10, size=len(a))
    pGamma = PopulationConcGamma(*prior_gamma_param(len(a)))

    def test_conc_ll(a):
        return np.sum(pGamma.conc_ll(a + x), axis=-1)

    def test_grad_conc_ll(a):
        return pGamma.d_conc_ll(a + x)

    print('approx gradient conc_ll = ', approx_fprime(a, test_conc_ll, epsilon=1e-6))
    print('exact  gradient = ', test_grad_conc_ll(a))
    err = check_grad(test_conc_ll, test_grad_conc_ll, a, epsilon=1e-6)
    print('check_grad err = ', err)
