"""General utility classes and functions for internal EmaCalc use.

Defines class
EmaParamBase --- container for common properties used by all model classes
    specifying the indexing of separate types of parameters,
    calculating logprob for observed data, given tentative model parameters.

NOTE: To make the model identifiable,
although both latent-variable locations and response thresholds are freely variable,
and the zero point on the attribute scale is arbitrary,
some restriction is necessary.

This behavior is user-controlled by initialization arguments
restrict_attribute and restrict_threshold.

The model is also slightly restricted by a weakly informative prior,
for numerical stability in case of extreme response patterns,
e.g., ALL ratings in the highest ordinal category.


*** Version history:
* Version 0.9.3:
2022-08-12, allowing ordinal scale tied to more than one attribute
2022-08-02, changed EmaParamBase properties specifying index slices
2022-07-21, xi parameters stored in order (alpha, beta, eta) parameter groups
2022-07-15, separate beta_slices, eta_slices;
            to allow threshold parameters tied to more than one attribute
2022-07-13, changed all names involving scenario -> situation

* Version 0.9.2:
2022-06-12, response-threshold mapping functions
    mapped_width(), d_mapped_width_d_eta(), and mapped_width_inv()
    slightly modified to be safe against numerical underflow.
    Global PRIOR_PSEUDO_RESPONDENT slightly modified, with no apparent effect on result.

* Version 0.9:
2022-03-17, use Pandas CategoricalDtype instances in EmaFrame situation_dtypes and attributes
2022-03-18, use Pandas DataFrame for EMA data in EmaDataSet, to allow many input file formats

* Version 0.8.1
2022-02-28, changed class name -> EmaParamBase
2022-02-23, GMM comp moved to ema_model.EmaGroupModel, separate GMM for each group

* Version 0.7.1
2022-01-11, _make_prior with log normalized uniform scenario prob

* Version 0.7
2022-01-02, minor cleanup _make_prior: ONLY scalar precision shape for GMM components

* Version 0.6
2021-12-03, user-definable prior inter-individual variance for reference parameters
2021-12-05, testing ways to set hyper-prior
2021-12-06, boolean properties restrict_attribute, restrict_threshold of EmaParamBase
            for user-specified model restriction
2021-12-08, minor checks against numerical overflow

* Version 0.5.1
2021-11-26, allow ZERO Attributes, i.e., empty emf.attribute_grades
2021-12-01, cleanup some doc comments
2021-12-02, Attribute sensory location fixed == 0. for first Scenario category
            regardless of regression_effect specification.

* Version 0.5
2021-11-24, first published beta version
"""
import numpy as np
from itertools import chain
from collections import namedtuple
from scipy.special import logit, expit
from scipy.special import logsumexp, softmax
import logging

from EmaCalc.gauss_gamma import GaussianRV

# ------------------------------------------------------
__version__ = "2022-08-12"

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

PRIOR_PSEUDO_RESPONDENT = 0.5  # seems to work OK
# = hyperprior total pseudo-count re ONE real respondent
# = prior GaussianRV.mean.learned_weight for all GMM components.
# This value is NOT critical for Occam's Razor behavior.

PRIOR_PARAM_SCALE = 1.
# = main hyperprior scale of most Gaussian model parameters,
# defined in Thurstone d-prime units for attribute parameters,
# rescaled if the Bradley (logistic) model is used for latent variables.
# This prior may have effect on the Occam's Razor function:
# small values may allow several active GMM components with small variance,
# but experiments suggest the prior value is not critical.

PRIOR_PREC_A = PRIOR_PSEUDO_RESPONDENT / 2
# = GaussianRV precision shape for ALL parameters

PRIOR_PREC_B = PRIOR_PARAM_SCALE**2 / 2
# = GaussianRV precision inv_scale for MOST EmaModel parameters
# -> allows very small precision <=> large inter-individual variance
# -> mode of prior component-element variance = PRIOR_PREC_B /(PRIOR_PREC_A + 1) approx= PRIOR_PREC_B

ETA_W_EPSILON = np.finfo(float).eps
# = additive constant to prevent too small mapped_width(eta)

# ------------------------------------------------------- Help classes
AttributeSlice = namedtuple('AttributeSlice', ['beta_slice', 'eta_slice'])


class Slicer:
    """Help class to create consecutive slices for parts of parameter vector
    """
    def __init__(self, start=0):
        self.start = start

    def slice(self, length):
        """Create a new slice object tight after previous slice
        :param length: desired slice size
        :return: a slice object
        """
        stop = self.start + length
        sl = slice(self.start, stop)
        self.start = stop  # for next call
        return sl


# ------------------------------------------------------------------
class EmaParamBase:
    """Container for common properties of ema_model.EmaModel and its sub-models.

    Each individual parameter distribution is represented by
    a large set of samples, stored as property EmaRespondentModel.xi, with
    xi[s, :] = s-th sample of the parameter vector for ONE respondent.

    All model classes share mapping properties defined here,
    for extracting parameter subsets of three different types
    from a xi array of parameter vectors.
    Parameters are stored consecutively in order (alpha, beta, eta)
    """
    def __init__(self, emf, effects, rv_class,
                 theta_map,
                 alpha_slices,
                 attribute_slices,
                 tied_scales,
                 comp_prior,
                 restrict_attribute=False,
                 restrict_threshold=False):
        """
        :param emf: single ema_data.EmaFrame instance,
        :param effects: iterable with regression-effect terms for attribute regression
            effects[i] = single key in emf.situation_dtypes.keys() or tuple of such keys
        :param rv_class: latent sensory variable class, defined in module ema_latent,
            defining the distribution as either logistic or normal (Gaussian).
        :param theta_map: fixed 2D array to extract latent-variable location samples from xi array
        :param alpha_slices: list of slice objects, such that
            xi[:, alpha_slices[t]] = alpha = log-prob for situations in t-th test phase
        :param attribute_slices: list of AttributeSlice objects, ordered like emf.attribute_dtypes,
            with attribute_slices[i] = (beta_slice, eta_slice), such that
            xi[:, beta_slice] = beta = regression-effect parameters for i-th Attribute
            xi[:, eta_slice] = eta = threshold parameters for i-th Attribute
        :param tied_scales: list of indices into attribute_slices,
            to allow more than one attribute to use SAME ordinal scale
        :param comp_prior: single GaussianRV instance, prior for ALL GMM components,
        :param restrict_attribute: (optional) boolean switch
            to force restriction on attribute sensory-variable locations
        :param restrict_threshold: (optional) boolean switch
            to force restriction on response-threshold locations
        """
        self.emf = emf
        self.effects = effects
        self.rv_class = rv_class
        self.theta_map = theta_map
        self.alpha_slices = alpha_slices
        self.attribute_slices = attribute_slices
        self.tied_scales = tied_scales
        self.comp_prior = comp_prior
        self.restrict_attribute = restrict_attribute
        self.restrict_threshold = restrict_threshold

    def __repr__(self):
        return (self.__class__.__name__ + '(' +
                '\n\t'.join(f'{k}= {repr(v)},'
                            for (k, v) in self.__dict__.items()) +
                '\n\t)')

    @classmethod
    def initialize(cls, emf, effects, rv_class,
                   restrict_attribute=False,
                   restrict_threshold=False):
        """Assign all parameter-extraction properties, and
        prior mixture components as GaussianRV objects with correct size.
        :param emf: single ema_data.EmaFrame instance,
            defining the experiment structure.
        :param effects: iterable with regression-effect terms for attribute regression
            effects[i] = single key in emf.situation_dtypes.keys() or tuple of such keys
        :param rv_class: latent variable class, defined in ema_latent,
            defining its distribution as either logistic or normal.
        :param restrict_attribute: (optional) boolean switch
            to force restriction on attribute sensory-variable locations
        :param restrict_threshold: (optional) boolean switch
            to force restriction on response-threshold locations
        :return: None

        Result: all properties initialized
        NOTE: xi parameter slices in order: all alpha, all beta, all eta
        """
        _check_effects(emf, effects)
        effects = [e_i if type(e_i) is tuple else (e_i,)
                   for e_i in effects]
        # = list of requested regression effects of situation categories
        # to be estimated for their influence on Attribute responses.
        # Each element MUST be a tuple with one or more situation key(s).

        # Define mapping indices from rows of xi array to its parts:
        n_phases = emf.situation_shape[0]
        n_situations = np.prod(emf.situation_shape[1:], dtype=int)  # within each phase
        slice_gen = Slicer()  # creating adjacent slice objects
        alpha_slices = [slice_gen.slice(n_situations)
                        for _ in range(n_phases)]
        # = list of slice objects, such that
        # xi[:, alpha_slices[t]] = alpha = log-prob for situations in t-th test phase
        # i.e., n_phases slices with equal lengths, for all test phases.

        theta_map = _theta_map(emf, effects)
        # = fixed 2D array to extract latent-variable location samples from xi array
        # by method attribute_theta(...)

        n_beta = theta_map.shape[0]
        # = number of regression-effect parameters, same for every Attribute
        n_attributes = len(emf.attribute_dtypes)
        beta_slices = [slice_gen.slice(n_beta)
                       for _ in range(n_attributes)]
        # = list of slice objects, such that
        # xi[:, beta_slices[i]] = beta = effect parameters for i-th attribute, such that
        # latent-variable location theta = np.dot(beta, theta_map)
        unique_slices = {scale_id: slice_gen.slice(len(scale.categories))
                         for (scale_id, scale) in emf.ordinal_scales.items()}
        eta_slices = [unique_slices[scale_id]
                      for scale_id in emf.attribute_scales.values()]
        # = list of slice objects, such that
        # xi[:, eta_slices[i]] = eta = threshold parameters for i-th attribute, such that
        # response thresholds tau = response_thresholds(eta)
        # Some thresholds possibly identical for more than one attribute
        attribute_slices = [AttributeSlice(beta, eta)
                            for (beta, eta) in zip(beta_slices,
                                                   eta_slices)]
        # *** attribute_slices as dict instead ?
        tied_scales = None
        if len(attribute_slices) > len(unique_slices):
            attribute_keys = list(emf.attribute_scales.keys())
            tied_scales = [[attribute_keys.index(a_key) for a_key in a_keys]
                           for a_keys in emf.scale_attributes.values()]
        # tied_scales[s] = list of attribute indices a_i, such that
        # attribute_slices[a_i].eta_slice is eta_slices[s],
        # allowing same eta parameters to be used for more than one attribute

        n_parameters = slice_gen.start  # no more slices needed
        comp_prior = _make_prior(n_parameters,
                                 alpha_slices,
                                 attribute_slices,
                                 rv_class)
        # = single GaussianRV instance, prior for ALL GMM components,
        # fixed throughout the VI learning procedure.
        logger.debug(f'PRIOR_PSEUDO_RESPONDENT = {PRIOR_PSEUDO_RESPONDENT:.3f}; ' +
                     f'PRIOR_PARAM_SCALE = {PRIOR_PARAM_SCALE:.3f}; ' +
                     f'PRIOR_PREC_A = {PRIOR_PREC_A:.3f}; ' +
                     f'PRIOR_PREC_B = {PRIOR_PREC_B:.3f}')
        logger.debug(f'comp_prior.prec.a = {comp_prior.prec.a:.3f}')
        logger.debug('comp_prior.prec.b = '
                     + np.array_str(comp_prior.prec.b,
                                    precision=5))
        logger.debug('mode{1 / comp_prior.prec}= '
                     + np.array_str(comp_prior.prec.mode_inv(),
                                    precision=5))
        return cls(emf, effects, rv_class,
                   theta_map,
                   alpha_slices,
                   attribute_slices,
                   tied_scales,
                   comp_prior,
                   restrict_attribute,
                   restrict_threshold)

    @property
    def n_parameters(self):  # *** needed ?
        if len(self.attribute_slices) > 0:
            return max(a_slice.eta_slice.stop
                       for a_slice in self.attribute_slices)
        else:
            return self.alpha_slices[-1].stop

    def beta_0_size(self):
        """number of effect parameters for FIRST effect term
        :return: n_beta = scalar

        NOTE: remaining effect terms are relative to first effect.
        """
        return np.prod([len(self.emf.situation_dtypes[sc_i].categories)
                        for sc_i in self.effects[0]],
                       dtype=int)

    def situation_prob(self, xi):
        """Extract probability-mass for situations, given parameters,
        used mainly by ema_display
        :param xi: 2D array with parameter samples
        :return: u = mD array with situation probability-mass within each phase,
            u[s, k0, k1, k2, ...] = s-th sample of P[(k1, k2,...)-th situation | phase k0]
            sum u[s, k0] == 1., for all s and k0

        2022-02-15, check for underflow
        """
        n_sc = self.alpha_slices[-1].stop
        alpha = xi[:, :n_sc].reshape((xi.shape[0], self.emf.n_phases, -1))
        alpha -= np.amax(alpha, axis=-1, keepdims=True)
        too_small = np.all(alpha < np.log(np.finfo(float).tiny), axis=-1)
        n_underflow = np.sum(too_small)
        if n_underflow > 0:
            logger.warning(f'situation_prob: {n_underflow} alpha samples too small. '
                           + 'Should not happen. Maybe too few responses?')
            logger.debug(f'alpha[too_small, :] = {alpha[too_small]}')
        u = np.exp(alpha)
        # avoid nan after normalization, should not be needed !
        u /= np.sum(u, axis=-1, keepdims=True)  # normalize within each Phase
        return u.reshape((-1, *self.emf.situation_shape))

    def attribute_theta(self, xi, a):
        """Extract location of latent sensory variable, for ONE given attribute
        used only by ema_display
        :param xi: 2D array with parameter sample vectors
            xi[s, :] = s-th parameter sample vector
        :param a: attribute key = one of self.emf.attribute_grades.keys()
        :return: theta = mD array, with
            theta[s, k0, k1, ...] = s-th sample of
                attribute location, given the (k0, k1, ...)-th situation.
        """
        a_index = list(self.emf.attribute_dtypes.keys()).index(a)
        beta_slice = self.attribute_slices[a_index].beta_slice
        # *** attribute_slices as dict instead ?
        beta = xi[..., beta_slice]
        return np.dot(beta, self.theta_map).reshape((-1, *self.emf.situation_shape))

    def attribute_tau(self, xi, a):
        """Extract response thresholds for ONE given attribute,
        EXCEPT fixed extreme limits at -inf, +inf.
        Used only by ema_display
        :param xi: 2D array with parameter sample vectors
            xi[s, :] = s-th parameter sample vector
        :param a: attribute key = one of self.emf.attribute_grades.keys()
        :return: tau = 2D array, with
            tau[s, l] = s-th sample of UPPER limit of l-th response interval,
                EXCEPT the last at +inf
                tau.shape[-1] == len(self.emf.attribute_grades[a]) - 1
        """
        a_index = list(self.emf.attribute_dtypes.keys()).index(a)
        eta_slice = self.attribute_slices[a_index].eta_slice
        eta = xi[..., eta_slice]
        return response_thresholds(eta)[..., 1:-1]

    # ---------------------- methods for logprob calculation:

    def situation_logprob(self, xi, situation_count):
        """Situation logprob given tentative parameter array
        Used by ema_respondent.EmaRespondentModel.adapt().
        :param xi: 1D or 2D array of tentative parameter samples
        :param situation_count: 2D array with response counts,
            situation_count[k0, k] = number of responses
            in k-th <=> (k1, k2, ...)-th situation category at k0-th test phase,
            using flattened index for situation dimensions 1,2,....
        :return: ll = scalar or 1D array
            ll[...] = ln P{situation_count | xi[..., :]}
            ll.shape == xi.shape[:-1]
        """
        def sit_lp(alpha):
            """
            :param alpha: non-normalized log-prob-mass for ONE test phase
            :return: lp = 1D or 2D array
                lp[..., k] = log P{visit k-th situation | alpha[..., :] }
                normalized as sum_k exp(lp[..., k]) == 1.
            """
            return alpha - logsumexp(alpha, axis=-1, keepdims=True)
        # -------------------------------------------------------
        return sum(np.dot(sit_lp(xi[..., alpha_ind]), z_t)
                   for (z_t, alpha_ind) in zip(situation_count,
                                               self.alpha_slices))

    def d_situation_logprob(self, xi, situation_count):
        """Gradient of situation_logprob,
        to be concatenated with d_rating_logprob(...)
        :param xi: 1D or 2D array of tentative parameter samples
        :param situation_count: 2D array with response counts
            situation_count[k0, k] = number of responses
            in k-th <=> (k1, k2, ...)-th situation category at k0-th test phase,
            using flattened index for situation dimensions 1,2,....
        :return: list of 1D or 2D arrays ll_t, with
            ll_t[..., j] = d ln P{situation_count[t] | xi[..., :]} / d xi[..., j_t]
                for j_t-th alpha parameter in t-th test phase
            ll_t.shape[-1] == situation_count.shape[-1]
        """
        def d_sit_lp(alpha):
            """Gradient of sit_lp(alpha)
           :param alpha: non-normalized log-prob-mass for ONE test phase
            :return: d_lp = 1D or 2D array with
                d_lp[..., k] = log P{visit k-th situation | alpha[..., :] }
                d_lp.shape == alpha.shape
            """
            return np.eye(alpha.shape[-1]) - softmax(alpha[..., None, :], axis=-1)
            # -------------------------------------------------------------------

        # return [np.einsum('k, ...kj -> ...j',  # ******* use dot or @?
        #                   z_t, d_sit_lp(xi[..., alpha_ind]))
        return [np.dot(z_t, d_sit_lp(xi[..., alpha_ind]))
                for (z_t, alpha_ind) in zip(situation_count,
                                            self.alpha_slices)]

    def rating_logprob(self, xi, rating_count):
        """Rating logprob given tentative parameter array
        Used by ema_respondent.EmaRespondentModel.adapt().
        :param xi: 1D or 2D array of tentative parameter samples
        :param rating_count: list of 2D arrays with response counts
            rating_count[i][l, k] = number of responses for i-th ATTRIBUTE,
            at l-th ordinal level, given the k-th <=> (k0, k1, k2, ...)-th situation
        :return: ll = scalar or 1D array
            ll[...] = log P{rating_count | xi[..., :]}
            ll.shape == xi.shape[:-1]
        """
        def rating_lp(beta, eta):
            """log-prob for ONE attribute
            :param beta: 1D or 2D array with effect params
            :param eta: 1D or 2D array with threshold params
            :return: lp = 2D or 3D array with
                lp[..., l, k] = ln P{l-th grade | k-th situation, beta, eta}
            """
            return self.rv_class.log_cdf_diff(*self.cdf_args(beta, eta))
        # --------------------------------------------------------------
        # ll_y = sum(np.einsum('lk, ...lk -> ...',
        #                      y_i, lp_i)
        return sum(np.sum(y_i * rating_lp(xi[..., beta_ind], xi[..., eta_ind]),
                          axis=(-2, -1))
                   for (y_i, (beta_ind, eta_ind)) in zip(rating_count,
                                                         self.attribute_slices)
                   )

    def d_rating_logprob(self, xi, rating_count):
        """Gradients of rating_logprob(...),
        for each separate group of parameters,
        to be concatenated with d_situation_logprob(...)
        :param xi: 1D or 2D array of tentative parameter samples
        :param rating_count: list of 2D arrays with response counts
            rating_count[i][l, k] = number of responses for i-th ATTRIBUTE,
            at l-th ordinal level, given the k-th <=> (k0, k1, k2, ...)-th situation
        :return: d_ll = list of 1D or 2D arrays d_ll_beta, and d_ll_eta, with elements
            d_ll_beta[a][..., j] = d ln P{rating_count[a] | xi[..., :]} / d beta_a[..., j]
                w.r.t. effect parameters beta_a for a-th attribute
            d_ll_eta[t][..., j] = d ln P{rating_count[a] | xi[..., :]} / d eta_t[..., j]
                w.r.t. threshold parameters eta for t-th attribute scale
                (each set of scale thresholds may be tied to one or several attributes)
            d_ll elements ordered as groups of parameters in parameter array xi
        """
        d_ll_beta = []
        d_ll_eta = []
        for (y_i, (beta_ind, eta_ind)) in zip(rating_count,
                                              self.attribute_slices):
            beta = xi[..., beta_ind]
            eta = xi[..., eta_ind]
            (dlp_low, dlp_high) = self.rv_class.d_log_cdf_diff(*self.cdf_args(beta, eta))
            d_ll_beta.append(np.einsum('...lk, lk, jk -> ...j',
                                       dlp_low + dlp_high, y_i, - self.theta_map))
            # tensordot 2 steps?

            d_tau_d_eta = d_response_thresholds(eta)
            d_ll_eta.append(np.einsum('...lk, lk, ...lj -> ...j',  # *** use dot , matmul ?
                                      dlp_low, y_i, d_tau_d_eta[..., :-1, :])
                            +
                            np.einsum('...lk, lk, ...lj -> ...j',
                                      dlp_high, y_i, d_tau_d_eta[..., 1:, :]))
        # all d_ll_beta, d_ll_eta elements have equal shape[:-1]
        # i.e., can be concatenated along axis=-1

        # if some eta scales tied to more than one attribute, sum those d_ll_eta elements:
        if self.tied_scales:
            d_ll_eta = [sum(d_ll_eta[i]
                            for i in tied_attribute_ind)
                        for tied_attribute_ind in self.tied_scales]
        # join into one list, elements to be concatenated by caller:
        return d_ll_beta + d_ll_eta

    def cdf_args(self, beta, eta):
        """Extract arguments for logprob calculation for ONE attribute
        :param beta: 1D or 2D array with beta parameter sample for given attribute
            beta[..., k] = ...-th sample of k-th regression-effect parameter
        :param eta: 1D or 2D array with threshold parameter sample for given attribute
            eta[..., l] = ...-th sample of l-th parameter
        :return: tuple (arg_low, arg_high) with 2D or 3D array args for probabilistic model,
            such that
            P[ l-th response | ...-th parameter sample, k-th situation ] =
            = rv_class.cdf(arg_high[..., l, k]) - rv_class.cdf(arg_low[..., l, k])
        """
        theta = np.dot(beta, self.theta_map)
        # theta[..., k] = ...-th sample of latent-variable location, given k-th situation
        tau = response_thresholds(eta)  # INCL extreme -Inf, +Inf limits
        a = tau[..., :-1, None] - theta[..., None, :]  # lower interval limits
        b = tau[..., 1:, None] - theta[..., None, :]  # upper interval limits
        return a, b

    def restrict(self, xi):
        """Adjust mean log situation-prob -> 0; -> no change of situation-prob.
        Force mean(theta) -> 0, OR median(tau) -> 0, after sampling,
        for all attributes.
        :param xi: ref to 2D array of sample parameter vectors
        :return: None
        Result: xi modified in place
            tau and associated theta samples modified by same amount
        """
        for sit_slice in self.alpha_slices:
            d = logsumexp(xi[:, sit_slice], axis=-1, keepdims=True)
            # logger.debug(f'xi situation restrict d: mean={np.mean(d):.3}; std={np.std(d):.3}; '
            #              + f'max={np.amax(d):.3}; min={np.amin(d):.3}')
            xi[:, sit_slice] -= d
        for (beta_slice, eta_slice) in self.attribute_slices:
            # beta = xi[..., beta_slice]
            # eta = xi[..., eta_slice]
            theta = np.dot(xi[..., beta_slice], self.theta_map)
            tau = response_thresholds(xi[..., eta_slice])[..., 1:-1]
            if self.restrict_attribute:
                d = np.mean(theta, axis=-1, keepdims=True)
            elif self.restrict_threshold:
                d = np.median(tau, axis=-1, keepdims=True)
            else:
                d = 0.
            # d = random offset to be zero-ed out with no effect on log-likelihood
            n_beta_0 = self.beta_0_size()  # **** only first regression-effect ! *****
            # = number of mean-location-controlling beta parameters = first part of beta_slice
            b0_slice = slice(beta_slice.start,
                             beta_slice.start + n_beta_0)
            # = beta slice only for FIRST regression effect term
            xi[:, b0_slice] -= d
            eta = tau_inv(tau - d)
            xi[:, eta_slice] = eta

    # ------------------------------------ Help methods for ema_respondent.initialize():
    def initialize_xi(self, z, y):
        """Initialize individual parameter vector, given observed counts
        :param z: 2D array with
            z[k0, k] = number of EMA records at k0-th test phase
            in k-th <=> (k1, k2, ...)-th situation, EXCL. k0= test phase
        :param y: list of 2D arrays,
            y[i][l, k] = number of responses at l-th ordinal level for i-th ATTRIBUTE question
            given k-th <=> (k0, k1, k2, ...)-th situation (INCL. k0 = test phase)
        :return: xi = 1D array with all parameters,
            ordered by parameter sub-types (alpha, beta, eta) as defined by self
        """
        alpha = list(_initialize_situation_param(z))
        eta = [_initialize_rating_eta(y_i)
               for y_i in y]
        beta = [_initialize_rating_beta(y_i, eta_i, self.theta_map)
                for (y_i, eta_i) in zip(y, eta)]
        if self.tied_scales:
            eta = [sum(eta[i] for i in tied_attribute_ind)
                   for tied_attribute_ind in self.tied_scales]
        for eta_scale in eta:
            eta_scale -= np.sum(eta_scale)
        # just in case _initial_rating_eta yields non-zero values
        # order as alpha[0], alpha[1],..., beta[0], beta[1], ..., eta[0], eta[1],...
        return np.concatenate([*chain(alpha, beta, eta)])


# ---------------------------------------- Other support functions:

def response_thresholds(eta):
    """Transform given log-category-width parameters to response thresholds.
    :param eta: = 1D or 2D array with
        eta[..., m] = ...-th sample of parameter defining
            non-normalized width of m-th interval in mapped domain [0, 1].
        eta.shape[-1] == M == number of response-scale intervals.
    :return: tau = 1D or 2D array, incl. all elements in [-inf, +inf]
        (tau[..., m], tau[..., m+1] = (LOWER, UPPER) limits for m-th ordinal response interval
        tau[..., 0] ==  - np.inf
        tau[..., -1] == + np.inf
        tau.ndim == eta.ndim; tau.shape[-1] == eta.shape[-1] + 1

    Method:
        Normalized widths and interval limits are defined in transformed domain [0, 1],
        using a logistic mapping function of actual thresholds tau:
        y = expit(tau), where y in [0, 1],
            y[..., 0] = 0.
            y[..., m+1] =  (w_0 +...+ w_m) / (w_0 + ... + w_{M-1};  0 <= m <= M-1
        Thus, thresholds are calculated with the inverse transform
        tau = logit(y)
        w_m = mapped_width(eta[..., m]) is a monotonic function of eta, always > 0
            = non-normalized width of transformed intervals
    """
    cum_w = np.cumsum(mapped_width(eta), axis=-1)
    z_shape = (*cum_w.shape[:-1], 1)
    cum_w = np.concatenate((np.zeros(z_shape), cum_w),
                           axis=-1)  # include cum_w[..., 0] = 0.
    # sum_w = cum_w[..., -1:]
    return logit(cum_w / cum_w[..., -1:])


def d_response_thresholds(eta):
    """Jacobian of cat_limits with respect to eta
    :param eta: 1D or 2D array (called eta in Leijon doc), with
        eta[..., m] = ...-th sample of log non-normalized width of m-th interval,
        eta.shape[-1] = M = number of response intervals
    :return: 2D or 3D array d_tau, with
        d_tau[..., m, i] = d tau[..., m] / d eta[..., i]; m = 0,..., M-2; i = 0, ..., M-1
            where (tau[s, m], tau[s, m+1] = (LOWER, UPPER) limits of m-th response interval
        d_tau[..., 0, :] = d_tau[..., -1, :] = 0., for extreme limits at +-inf
        d_tau.shape == (N, M+1, M); (N, M) = eta.shape
    """
    w = mapped_width(eta)
    # (n_samples, nw) = w.shape
    nw = w.shape[-1]
    cum_w = np.cumsum(w, axis=-1)
    cw = cum_w[..., :-1, np.newaxis]  # only inner limits
    sw = cum_w[..., -1:, np.newaxis]
    # tau[..., m+1] = ln cw[..., m]  - ln (sw[..., 0] - cw[..., m])
    # dcw_dw[..., m, i] = dcw[..., m] / dw[..., i]  = 1. if i <= m else 0.
    dcw_dw = np.tril(np.ones((nw - 1, nw), dtype=int))
    dtau_dw = dcw_dw / cw - (1 - dcw_dw) / (sw - cw)
    dtau_d_eta = dtau_dw * d_mapped_width_d_eta(w)[..., np.newaxis, :]
    # dtau_d_eta.shape = (n_samples, nw + 1, nw)  OR (nw + 1, nw)
    z_shape = list(dtau_d_eta.shape)
    z_shape[-2] = 1
    return np.concatenate((np.zeros(z_shape),  # *** use np.pad ???
                           dtau_d_eta,
                           np.zeros(z_shape)), axis=-2)


def tau_inv(tau):
    """Inverse of response_thresholds(eta)
    :param tau: 1D or 2D array with response thresholds, EXCEPT extremes at +-inf,
        i.e., all tau elements in (-inf, +inf),
        tau[..., m] = UPPER limit for m-th interval,
            = LOWER limit for the (m+1)-th interval
        tau.shape[-1] == M - 1 == number of response intervals - 1
    :return: eta: 1D or 2D array with
        eta[..., m] = ...-th sample of log non-normalized width of m-th interval.
        eta.shape[-1] == M == number of response intervals.

    Method:
        Normalized widths and interval limits are defined in transformed domain (0, 1.),
        using a logistic mapping function,
        y = expit(tau), where y in (0, 1]
            y[..., m] =  (w_0 +...+ w_m) / (w_0 + ... + w_{M-1};  0 <= m <= M-1
            w_m = _w(eta[..., m])
    """
    y = expit(tau)
    cat_shape = (*y.shape[:-1], 1)
    # include extreme limits at 0 and 1:
    y = np.concatenate((np.zeros(cat_shape),
                        y,
                        np.ones(cat_shape)), axis=-1)
    w = np.diff(y, axis=-1)
    return mapped_width_inv(w)


# ----------- Original up to version 0.9.1 -> numeric overflow in some extreme cases
# mapped_width = _w = np.exp
# _w_inv = np.log  # _w_inv(w) -> eta, such that _w(eta) == w
#
# ---------- *** piecewise (exp, linear ) variant, tested no good
# ---------- *** piecewise (inverted linear, linear ) variant, tested no good


def mapped_width(eta):
    """Mapping function from model eta param to
    non-normalized interval widths in (0, 1) range
    :param eta: = 1D or 2D array with
        eta[..., m] = ...-th sample of parameter defining
            non-normalized width of m-th interval in mapped domain [0, 1].
        eta.shape[-1] == M == number of response-scale intervals.
    :return: w = array with mapped widths
        w.shape == eta.shape
    """
    return np.exp(eta) + ETA_W_EPSILON


def d_mapped_width_d_eta(w):
    """Derivative of _w, , as a function of w, NOT eta.
    :param w: = mapped_width(eta) = 1D or 2D array
    :return: array dw, with
        dw[..., m] = d _w(eta)[..., m] / d eta[..., m]
    """
    return w - ETA_W_EPSILON


def mapped_width_inv(w):
    """Inverse of mapped_width(eta),
    used only after each VI iteration,
    so it is OK to be slightly inconsistent with mapped_width()
    :param w: = 1D or 2D array with
        w[..., m] = ...-th sample of
            non-normalized width of m-th interval in mapped domain [0, 1].
        sum_m w[..., m] approx == 1.
        w.shape[-1] == M == number of response-scale intervals.
    :return: eta = array with mapped widths mapped back to eta parameter
        eta.shape == w.shape
    """
    return np.log(w + np.finfo(float).tiny)


def _initialize_situation_param(z):
    """Crude initial estimate of situation logprob parameters
    :param z: 2D array with situation counts
        z[t, k] = n EMA records in k-th <-> (k1, k2,...)-th situation
            in t-th test phase
    :return: alpha = 2D array
        alpha[t, k] = log prob of k-th situation at t-th test phase
        alpha.shape == (z.shape[0], z[0].size)
    """
    p = z + PRIOR_PSEUDO_RESPONDENT
    p /= np.sum(p, axis=-1, keepdims=True)
    return np.log(p)


def _initialize_rating_eta(y):
    """Crude initial estimate of threshold-defining parameters for ONE attribute question
    :param y: 2D rating_count array,
        y[l, k] = number of responses at l-th ordinal level,
        given the k-th <-> (k0, k1, ...)-th situation category
    :return: eta = 1D array of eta parameters, defining thresholds as
        tau = ema_base.response_thresholds(eta)
        eta.shape == y.shape[:1]
    Method: set eta parameters to match observed relative freq of all rating grades
        as if theta = 0.
    OR just set all eta == 0.
    """
    # p = np.sum(y, axis=-1) + JEFFREYS_CONC
    # p /= np.sum(p)
    # return mapped_width_inv(p)
    # *** This method seems to improve the first LL,
    # but sometimes give worse or same long-term result
    return np.zeros(len(y))  # original simple method


def _initialize_rating_beta(y, eta, theta_map):
    """Crude initial estimate of regression-effect parameters for ONE attribute question
    :param y: 2D rating_count array,
        y[l, k] = number of responses at l-th ordinal level,
        given the k-th <-> (k0, k1, ...)-th situation category
    :param eta: 1D array with threshold-defining parameters
    :return: beta = array of sensory-variable locations theta, such that
        theta approx = np.dot(beta, theta_map),
        theta[k] = estimated location of latent variable,
        given the k-th <-> (k0, k1, ...)-th situation category
    """
    expit_tau = expit(response_thresholds(eta))
    typical_theta = logit((expit_tau[:-1] + expit_tau[1:]) / 2)
    # = back-transformed midpoints in each response interval
    # typical_theta[l] = theta given l-th response
    p = y + PRIOR_PSEUDO_RESPONDENT
    p /= np.sum(p, axis=0, keepdims=True)
    theta = np.dot(typical_theta, p)
    # = typical (midpoint) locations, given y
    beta = np.linalg.lstsq(theta_map.T, theta, rcond=None)[0]
    # -> theta approx = np.dot(beta, theta_map),
    return beta


def _make_prior(n_parameters,
                alpha_slices,
                attribute_slices,
                rv_class):
    """Create hyperprior for parameter distribution in the total population
    :param n_parameters: total number of parameters
    :param alpha_slices: list with index slice for each situation phase
    :param attribute_slices: list with one AttributeSlice tuple for each attribute
    :return: single GaussianRV instance
    """
    prec_a = PRIOR_PREC_A  # scalar, same for all elements
    prec_b = PRIOR_PREC_B * np.ones(n_parameters)
    for (b_slice, e_slice) in attribute_slices:
        prec_b[b_slice] *= rv_class.scale**2
    loc = np.zeros(n_parameters)
    for a_slice in alpha_slices:
        loc[a_slice] = - np.log(a_slice.stop - a_slice.start)
    return GaussianRV.initialize(loc=loc,
                                 learned_weight=PRIOR_PSEUDO_RESPONDENT,
                                 prec_a=prec_a,
                                 prec_b=prec_b)


def _check_effects(emf, effects):
    """Check that all effects refer to unique situation keys
    :param emf: ema_data.EmaFrame instance
    :param effects: iterable with regression-effect specifications
        effects[i] = single key in emf.situation_dtypes.keys(), OR tuple of such keys
    :return: None
    Result: raise RuntimeError if incorrect effects
    """
    effect_keys = sum((list(e_i) if type(e_i) is tuple else [e_i]
                      for e_i in effects), start=[])
    # = flat list with all situation keys from effects
    if len(set(effect_keys)) != len(effect_keys):
        raise RuntimeError('situation keys can occur only ONCE in regression effects')
    for sit_key in emf.situation_dtypes.keys():
        if len(emf.situation_dtypes[sit_key].categories) > 1 and sit_key not in effect_keys:
            logger.warning(f'situation dim {repr(sit_key)} to be disregarded in regression?')
    for e_i in effect_keys:
        if e_i not in emf.situation_dtypes.keys():
            raise RuntimeError(f'regression effect key {e_i} is not a situation key')


def _theta_map(emf, effects):
    """Create 2D array for extraction of attribute locations
    from parameter vector samples
    :param emf: ema_data.EmaFrame instance,
        defining the analysis model structure.
    :param effects: iterable with regression-effect specifications
        effects[i] = tuple with one or more key(s) from emf.situation_dtypes.keys()
    :return: th = 2D binary array, with
        th[j, k] = 1, IFF parameter
        beta[:, j] = xi[:, attr_slice][:, j]
        is the regression effect of k-th <=> (k0, k1, ...)-th situation on attribute location theta,
        and xi is a 2D array of row parameter-vector samples.
    """
    def theta_one(effect):
        """Make theta_map part for ONE effect term
        :param effect: tuple of one or more situation keys
        :return: mD array th, with
            th[j, k0, k1, ...] = j-th effect given (k0, k1, ...)-th situation
            th[j, 0, 0, ...] = 0, in 0-th reference situation.
            th.shape == (size_effect, *emf.situation_shape), where
            size_effect = size of category array defined by effect.
        """
        beta_shape = tuple(len(emf.situation_dtypes[sc_i].categories) for sc_i in effect)
        # = situation_shape, for this effect term
        beta_ndim = len(beta_shape)
        beta_size = int(np.prod(beta_shape))
        t = np.eye(beta_size).reshape((beta_size, *beta_shape))
        if effect != effects[0]:
            t = t[1:]
            # fixed effect beta[0] for all additional effect terms, EXCEPT FIRST
            # -> ONE overall theta intercept parameter only in FiRST effect term,
            # with same influence on theta in all effect terms
        # expand t.ndim to number of situation dimensions:
        t = t.reshape(t.shape + tuple(np.ones(len(emf.situation_dtypes) - beta_ndim,
                                              dtype=int)
                                      ))
        ax_0 = range(1, 1 + len(beta_shape))
        ax_new = tuple(list(emf.situation_dtypes).index(e_i) + 1
                       for e_i in effect)
        t = np.moveaxis(t, ax_0, ax_new)
        return t + np.zeros(emf.situation_shape)
    # ---------------------------------------------------------------
    th = np.concatenate([theta_one(e_tuple)
                         for e_tuple in effects],
                        axis=0)
    return th.reshape((th.shape[0], -1))  # keep as 2D for simplicity


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

    from ema_data import EmaFrame
    from ema_latent import Bradley

    print('*** Testing some ema_base module functions ')
    emf = EmaFrame.setup(situations={'CoSS': [f'{i}.' for i in range(1, 8)],
                                     'Viktigt': ('Lite',
                                                 'Ganska',
                                                 'Mycket'),
                                     # 'Test': (None, None),
                                     'HA': ('A', 'B')
                                     }, attributes={'Speech': ('Difficult',
                                                               'Normal',
                                                               'Easy'),
                                                    'Quality': ('Bad', 'Good')
                                                    })

    # print('emf=\n', emf)

    main_theta = {'Phase': [0.],
                  'CoSS': 0.1 * np.arange(len(emf.situation_dtypes['CoSS'].categories)),
                  'Viktigt': np.array([-1., 0., 1.]),
                  'HA': np.array([0., 1.])
                  }
    # = only main effects, additive
    regr_effects = ['HA', ('CoSS', 'Viktigt')]
    # regr_effects = [('CoSS', 'Viktigt'), 'HA']

    # effect_HA = np.array([-1., 1.])
    # effect_CoSS = 0.1 * np.arange(len(emf.situation_dtypes['CoSS']))
    # effect_Vikt = np.array([-1., 0., 1.])
    true_theta = main_theta['CoSS'][:, None] + main_theta['Viktigt']
    true_theta = true_theta[..., None] + main_theta['HA']
    true_theta = true_theta[None, ...]  # only one phase
    # stored like emf.situation_dtypes
    # true_theta[t, c, v, h] = theta in t-th phase, c-th CoSS, v-th Vikt, h-th HA

    theta_HA = main_theta['HA']
    beta_HA = theta_HA[1:] - theta_HA[0]  # rel. first HA
    theta_CoSS_Vikt = main_theta['CoSS'][:, None] + main_theta['Viktigt']
    theta_CoSS_Vikt = theta_CoSS_Vikt.reshape((-1,))
    beta_CoSS_Vikt = theta_CoSS_Vikt[1:] - theta_CoSS_Vikt[0]  # rel. first CoSS_Vikt case
    beta_all = np.concatenate((theta_HA + theta_CoSS_Vikt[0], beta_CoSS_Vikt))
    # for regr_effects = ['HA', ('CoSS', 'Viktigt')]
    # beta_all = np.concatenate((theta_CoSS_Vikt.reshape((-1)), beta_HA))
    # for regr_effects = [('CoSS', 'Viktigt'), 'HA']

    # = same for all attributes

    eta = [np.zeros(len(r_cat.categories))
           for r_cat in emf.attribute_dtypes.values()]

    lp_situations = np.zeros(emf.situation_shape).reshape((-1))
    # beta_eta = [np.concatenate((beta_all, eta_r))
    #             for eta_r in eta]
    # print(f'len(beta_eta): ', [len(be_i) for be_i in beta_eta])
    print(f'len(beta): ', [len(beta_all) for _ in range(len(eta))])
    print(f'len(eta): ', [len(e_i) for e_i in eta])

    # xi = np.concatenate((lp_situations, *beta_eta))
    xi = np.concatenate((lp_situations,
                         np.tile(beta_all, len(emf.attribute_dtypes)),
                         *eta))
    xi = xi[None, :]
    print(f'xi= {xi}')

    p_base = EmaParamBase.initialize(emf, regr_effects, rv_class=Bradley)
    print('p_base= ', p_base)

    print('\n*** Testing param extraction methods ***')

    for (phase, alpha_slice) in zip(emf.situation_dtypes[emf.phase_key].categories,
                                    p_base.alpha_slices):
        print(f'{emf.phase_key}= {phase}: xi[:, alpha_slice].shape= ', xi[:, alpha_slice].shape)
    # for (phase, alpha) in zip(emf.situation_dtypes[emf.phase_key], p_base.situation_logprob(xi)):
    #     print(f'{emf.phase_key}= {phase}: prob-mass=\n', np.exp(alpha))
    print(f'\nemf.situation_shape= {emf.situation_shape}')
    # for (a, beta_slice) in zip(emf.attribute_grades.keys(),
    #                            p_base.beta_slices):
    # for (a, a_slice) in p_base.attribute_slices.items():
    for (a, (beta_slice, eta_slice)) in zip(emf.attribute_dtypes.keys(),
                                            p_base.attribute_slices):
        print(f'\nAttribute {a}.beta_slice: {beta_slice}')
        print(f'\nAttribute {a}: xi[:, beta_slice]={xi[:, beta_slice]}')

    for a in emf.attribute_dtypes.keys():
        theta = p_base.attribute_theta(xi, a)
        print(f'\nAttribute {a}: theta.shape={theta.shape}.')
        sc_dim_0 = 'Phase'
        sc_dim_1 = 'CoSS'
        theta = theta.reshape((-1, *emf.situation_shape))  # *** just for testing
        for (sc_0, theta_0, true_theta_0) in zip(emf.situation_dtypes[sc_dim_0].categories,
                                                 theta[0], true_theta):
            print(f'\t{sc_dim_0}= {sc_0}:')
            for (sc_1, theta_1, true_theta_1) in zip(emf.situation_dtypes[sc_dim_1].categories,
                                                     theta_0, true_theta_0):
                print(f'\t\t{sc_dim_1}= {sc_1}: theta=\n', theta_1)
                print(f'\t\t{sc_dim_1}= {sc_1}: true_theta=\n', true_theta_1)

    for a in emf.attribute_dtypes.keys():
        tau = p_base.attribute_tau(xi, a)
        print(f'\nAttribute {a}: tau.shape={tau.shape}. tau=\n', tau)

    print('\n*** Testing response_thresholds() and tau_inv() ***')

    eta = np.arange(4) - 50.
    eta = np.array([0., 0., -1., -50., -50.])
    tau = response_thresholds(eta)
    print(f'eta = ', eta)
    print(f'_w(eta) = ', mapped_width(eta))
    print(f'tau({eta}) = ', tau)
    print(f'tau_inv(tau[1:-1] = ', tau_inv(tau[1:-1]))
    print(f'tau(tau_inv(tau[1:-1]) = ', response_thresholds(tau_inv(tau[1:-1])))

    # ----- **** check extreme tau, that might occur in restrict_xi ************
    tau = np.array([-np.inf, -100., 0., 0., 100., np.inf])
    print(f'tau = ', tau)
    print(f'tau_inv(tau[1:-1] = ', tau_inv(tau[1:-1]))
    print(f'tau(tau_inv(tau[1:-1]) = ', response_thresholds(tau_inv(tau[1:-1])))

    print('\n*** Testing d_response_thresholds gradient functions ***')

    eta = np.arange(4) - 50.
    eta = np.array([0., 0., -1., -50., -50.])
    limit = 1

    def fun(eta):
        return response_thresholds(eta.reshape((1, -1)))[0, limit]

    def jac(eta):
        return d_response_thresholds(eta.reshape((1, -1)))[0, limit]

    print(f'tau(eta) = {response_thresholds(eta.reshape((1, -1)))}')
    # print(f'tau(eta+d) = {response_thresholds(eta.reshape((1, -1)) + .1)} Should remain the same')

    print('approx gradient = ', approx_fprime(eta, fun, epsilon=1e-6))
    print('exact  gradient = ', jac(eta))
    err = check_grad(fun, jac, eta, epsilon=1e-6)
    print('check_grad err = ', err)

