import numpy as np
import pandas as pd
import xarray as xr

from .._optional import jax, jnp
from ..compiled import jitmethod, reset_compiled_methods
from ..dataset import Dataset, DataTree
from ..exceptions import MissingDataError
from ..folding import dissolve_zero_variance, fold_dataset
from ..optimize import OptimizeMixin
from .basemodel import MANGLE_DATA
from .basemodel import BaseModel as _BaseModel
from .jaxmodel import PanelMixin, _get_jnp_array
from .mixtures import MixtureList


class LatentClass(_BaseModel, OptimizeMixin, PanelMixin):
    compute_engine = "jax"

    def __init__(
        self, classmodel, choicemodels, datatree=None, float_dtype=np.float32, **kwargs
    ):
        self._model_subtype = "latent-class"
        classmodel._model_subtype = "class-membership"
        for k, m in choicemodels.items():
            m.ident = k
        PanelMixin.__init__(self, **kwargs)
        self._is_mangled = True
        self._dataset = None
        choicemodels_keys = sorted(choicemodels.keys())
        if (
            classmodel.datatree is None
            and isinstance(datatree, Dataset)
            and self.groupid
        ):
            df = datatree.to_dataframe()
            df["ingroup"] = df.groupby(self.groupid).cumcount() + 1
            classdata = dissolve_zero_variance(
                df.set_index([self.groupid, "ingroup"], drop=True).to_xarray(),
                "ingroup",
            )
            classdata = classdata.dc.set_altids(choicemodels_keys).drop_dims("ingroup")
            classdata.dc.CASEID = self.groupid
            classmodel.datatree = classdata
        else:
            pass
        self.ident = "latent-class"
        super().__init__(
            datatree=datatree,
            submodels=choicemodels,
            named_submodels={"classmodel": classmodel},
        )
        # self.datatree = datatree
        if classmodel.dataset is not None:
            assert sorted(classmodel.dataset.dc.altids()) == choicemodels_keys
        if float_dtype is not None:
            for v in self._models.values():
                v.float_dtype = float_dtype
        for _k, m in self._models.items():
            if m._model_subtype == "latent-class":
                pass
            elif m._model_subtype == "class-membership":
                m.groupid = self.groupid
            else:
                m.datatree = self.datatree
                m.groupid = self.groupid

    @property
    def _models(self):
        return self._parameter_bucket._models

    def save(self, filename, format="yaml", overwrite=False):
        from .saving import save_model

        return save_model(self, filename, format=format, overwrite=overwrite)

    @classmethod
    def from_dict(cls, content):
        _models = content.get("_models")
        from .saving import load_model

        classmodel = None
        choicemodels = {}
        for k, v in _models.items():
            if v._model_subtype == "class-membership":
                classmodel = load_model(v)
            else:
                choicemodels[k] = load_model(v)
        self = cls(classmodel, choicemodels)

        def loadthis(attr, wrapper=None, injector=None):
            i = content.get(attr, None)
            if i is not None:
                try:
                    if wrapper is not None:
                        i = wrapper(i)
                except AttributeError:
                    pass
                else:
                    if injector is None:
                        setattr(self, attr, i)
                    else:
                        injector(i)

        loadthis("float_dtype", lambda i: getattr(np, i))
        loadthis("compute_engine")
        loadthis("index_name")
        loadthis("parameters", xr.Dataset.from_dict, self.update_parameters)
        loadthis("availability_any")
        loadthis("availability_ca_var")
        loadthis("availability_co_vars")
        loadthis("choice_any")
        loadthis("choice_ca_var")
        loadthis("choice_co_code")
        loadthis("choice_co_vars")
        loadthis("constraint_intensity")
        loadthis("constraint_sharpness")
        loadthis("constraints")
        from .tree import NestingTree

        loadthis("graph", NestingTree.from_dict)
        loadthis("groupid")
        loadthis("logsum_parameter")
        loadthis("quantity_ca")
        loadthis("quantity_scale")
        loadthis("title")
        loadthis("utility_ca")
        loadthis("utility_co")
        loadthis("weight_co_var")
        loadthis("weight_normalization")
        return self

    @property
    def pf(self):
        return self.parameters.to_dataframe()

    @jitmethod
    def jax_probability(self, params):
        classmodel = self.class_membership_model
        pr_parts = []
        for n, k in enumerate(classmodel.dataset.dc.altids()):
            class_pr = jnp.expand_dims(classmodel.jax_probability(params)[..., n], -1)
            inclass_pr = self._models[k].jax_probability(params)
            if inclass_pr.ndim < class_pr.ndim:
                inclass_pr = inclass_pr.reshape(
                    class_pr.shape[0],
                    inclass_pr.shape[0] // class_pr.shape[0],
                    *inclass_pr.shape[1:],
                )
            print(f"{class_pr.shape=}")
            print(f"{inclass_pr.shape=}")
            pr_parts.append(inclass_pr * class_pr)
        return sum(pr_parts)

    @jitmethod
    def jax_loglike_casewise(self, params):
        n_alts = self.dataset.dc.n_alts
        ch = jnp.asarray(self.dataset["ch"])
        if ch.ndim == 2:
            pr = self.jax_probability(params)
            masked_pr = jnp.where(ch[..., :n_alts] > 0, pr[..., :n_alts], 1.0)
            log_pr = jnp.log(masked_pr)
            return (log_pr[..., :n_alts] * ch[..., :n_alts]).sum()
        elif ch.ndim >= 3:
            classmodel = self.class_membership_model
            class_pr = classmodel.jax_probability(params)
            likely_parts = []
            for n, k in enumerate(classmodel.dataset.dc.altids()):
                _yo_data = self._models[k].dataset
                k_pr = self._models[k].jax_probability(params)  # .reshape(ch.shape)
                masked_k_pr = jnp.where(ch[..., :n_alts] > 0, k_pr[..., :n_alts], 1.0)
                k_likely = jnp.power(masked_k_pr, ch[..., :n_alts]).prod([-2, -1])
                likely_parts.append(k_likely * class_pr[..., 0, n])
            return jnp.log(sum(likely_parts))

    @jitmethod
    def jax_loglike(self, params):
        return self.jax_loglike_casewise(params).sum()

    def loglike(
        self, x=None, *, start_case=None, stop_case=None, step_case=None, **kwargs
    ):
        if self.compute_engine != "jax":
            raise NotImplementedError(f"latent class with engine={self.compute_engine}")
        self.unmangle()
        if start_case is not None and start_case != 0:
            raise NotImplementedError(f"{start_case=} with engine=jax")
        if stop_case is not None and stop_case != -1:
            raise NotImplementedError(f"{stop_case=} with engine=jax")
        if step_case is not None and step_case != 1:
            raise NotImplementedError(f"{step_case=} with engine=jax")
        if x is not None:
            self.pvals = x
        result = float(self.jax_loglike(self.pvals))
        # if start_case is None and stop_case is None and step_case is None:
        #     self._check_if_best(result)
        return result

    def neg_loglike(
        self,
        x=None,
        start_case=None,
        stop_case=None,
        step_case=None,
        leave_out=-1,
        keep_only=-1,
        subsample=-1,
    ):
        result = self.loglike(
            x,
            start_case=start_case,
            stop_case=stop_case,
            step_case=step_case,
            leave_out=leave_out,
            keep_only=keep_only,
            subsample=subsample,
        )
        return -result

    @jitmethod
    def jax_d_loglike(self, params):
        return jax.grad(self.jax_loglike)(params)

    def d_loglike(
        self,
        x=None,
        *,
        start_case=None,
        stop_case=None,
        step_case=None,
        return_series=False,
        **kwargs,
    ):
        if self.compute_engine != "jax":
            raise NotImplementedError(f"latent class with engine={self.compute_engine}")
        self.unmangle()
        if start_case is not None and start_case != 0:
            raise NotImplementedError(f"{start_case=} with engine=jax")
        if stop_case is not None and stop_case != -1:
            raise NotImplementedError(f"{stop_case=} with engine=jax")
        if step_case is not None and step_case != 1:
            raise NotImplementedError(f"{step_case=} with engine=jax")
        if x is not None:
            self.pvals = x
        result = self.jax_d_loglike(self.pvals)

        # print("converge?=", jnp.max(jnp.absolute(result)))

        if return_series:
            result = pd.Series(result, index=self.pnames)
        return result

    def neg_d_loglike(self, x=None, start_case=0, stop_case=-1, step_case=1, **kwargs):
        result = self.d_loglike(
            x, start_case=start_case, stop_case=stop_case, step_case=step_case, **kwargs
        )
        return -np.asarray(result)

    def maximize_loglike(
        self,
        *args,
        **kwargs,
    ):
        """
        Maximize the log likelihood.

        Parameters
        ----------
        method : str, optional
            The optimization method to use.  See scipy.optimize for
            most possibilities, or use 'BHHH'. Defaults to SLSQP if
            there are any constraints or finite parameter bounds,
            otherwise defaults to BHHH.
        quiet : bool, default False
            Whether to suppress the dashboard.
        options : dict, optional
            These options are passed through to the `scipy.optimize.minimize`
            function.
        maxiter : int, optional
            Maximum number of iterations.  This argument is just added to
            `options` for most methods.

        Returns
        -------
        dictx
            A dictionary of results, including final log likelihood,
            elapsed time, and other statistics.  The exact items
            included in output will vary by estimation method.

        """
        from .optimization import maximize_loglike

        return maximize_loglike(self, *args, **kwargs)

    def reflow_data_arrays(self):
        """Reload the internal data_arrays so they are consistent with the datatree."""
        datatree = self.datatree
        if datatree is None:
            raise ValueError("missing datatree")

        classmodel = self.class_membership_model
        request = classmodel.required_data()
        request.pop("avail_any", None)
        for _kname, kmodel in self._models.items():
            if kmodel._model_subtype in ("class-membership", "latent-class"):
                continue
            kreq = kmodel.required_data()
            for k, v in kreq.items():
                if k not in request:
                    request[k] = v
                else:
                    if isinstance(request[k], dict):
                        request[k].update(v)
                    elif isinstance(request[k], (list | tuple)):
                        request[k] = list(set(request[k]) | set(v))
                    else:
                        if request[k] != v:
                            raise ValueError("incompatible requests")

        if isinstance(self.groupid, str):
            request["group_co"] = self.groupid

        from .data_arrays import prepare_data

        dataset, self.dataflows = prepare_data(
            datasource=datatree,
            request=request,
            float_dtype=np.float32,
            cache_dir=datatree.cache_dir,
            flows=getattr(self, "dataflows", None),
        )
        if isinstance(self.groupid, str):
            dataset = fold_dataset(dataset, "group")
        elif self.groupid is not None:
            dataset = fold_dataset(dataset, self.groupid)
        self.dataset = dataset

        for _kname, kmodel in self._models.items():
            if kmodel._model_subtype in ("class-membership", "latent-class"):
                continue
            kmodel.dataset = self.dataset
            # kmodel.reflow_data_arrays() # not full reflow, just...
            kmodel._data_arrays = kmodel.dataset.dc.to_arrays(
                kmodel.graph,
                float_dtype=kmodel.float_dtype,
            )
            # if kmodel.work_arrays is not None:
            kmodel._rebuild_fixed_arrays()
            kmodel._rebuild_work_arrays()
        classmodel_ids = [
            kid
            for (kid, km) in self._models.items()
            if km._model_subtype not in ("latent-class", "class-membership")
        ]
        classmodel_data = dissolve_zero_variance(
            self.dataset.drop_dims(self.dataset.dc.ALTID).dc.set_altids(classmodel_ids),
            "ingroup",
        )
        classmodel.dataset = classmodel_data
        classmodel.reflow_data_arrays()
        # classmodel._data_arrays = classmodel.dataset.dc.to_arrays(
        #     classmodel.graph,
        #     float_dtype=classmodel.float_dtype,
        # )
        # if classmodel.work_arrays is not None:
        #     classmodel._rebuild_work_arrays()

    def mangle(self, data=True, structure=True):
        super().mangle(data=data, structure=structure)
        reset_compiled_methods(self)
        self._is_mangled = True

    def unmangle(self, force=False, structure_only=False):
        if self._is_mangled:
            super().unmangle(force=force, structure_only=structure_only)
            if not structure_only:
                self.reflow_data_arrays()
                self._is_mangled = False
            else:
                self._is_mangled = MANGLE_DATA

    @property
    def dataset(self):
        """larch.Dataset : Data arrays as loaded for model computation."""
        super().unmangle()
        if self._dataset is None:
            self.reflow_data_arrays()
        try:
            return self._dataset
        except AttributeError:
            return None

    @dataset.setter
    def dataset(self, dataset):
        if dataset is self._dataset:
            return
        from xarray import Dataset as _Dataset

        if isinstance(dataset, Dataset):
            self._dataset = dataset
            self._data_arrays = None
            self.mangle()
        elif isinstance(dataset, _Dataset):
            self._dataset = Dataset(dataset)
            self._data_arrays = None
            self.mangle()
        else:
            raise TypeError(f"dataset must be Dataset not {type(dataset)}")

    @property
    def datatree(self):
        """DataTree : A source for data for the model."""
        try:
            return self._datatree
        except AttributeError:
            return None

    @datatree.setter
    def datatree(self, tree):
        if tree is self.datatree:
            return
        if isinstance(tree, DataTree) or tree is None:
            self._datatree = tree
            self.mangle()
        elif isinstance(tree, Dataset):
            self._datatree = tree.dc.as_tree()
            self.mangle()
        else:
            try:
                self._datatree = DataTree(main=Dataset.construct(tree))
            except Exception as err:
                raise TypeError(f"datatree must be DataTree not {type(tree)}") from err
            else:
                self.mangle()

    @property
    def data_as_loaded(self):
        return self._dataset

    @property
    def class_membership_model(self):
        # TODO optimize this
        for m in self._models.values():
            if m._model_subtype == "class-membership":
                return m

    def total_weight(self):
        """
        Compute the total weight of cases in the loaded data.

        Returns
        -------
        float
        """
        if self.class_membership_model._data_arrays is not None:
            return self.class_membership_model._data_arrays.wt.sum()
        raise MissingDataError("no data_arrays are set")

    def logloss(
        self,
        x=None,
        start_case=None,
        stop_case=None,
        step_case=None,
        leave_out=-1,
        keep_only=-1,
        subsample=-1,
    ):
        result = self.loglike(
            x,
            start_case=start_case,
            stop_case=stop_case,
            step_case=step_case,
            leave_out=leave_out,
            keep_only=keep_only,
            subsample=subsample,
        )
        return -result / self.total_weight()

    def d_logloss(self, x=None, start_case=0, stop_case=-1, step_case=1, **kwargs):
        result = self.d_loglike(
            x, start_case=start_case, stop_case=stop_case, step_case=step_case, **kwargs
        )
        return -np.asarray(result) / self.total_weight()

    def simple_fit_bhhh(self, *args, **kwargs):
        raise NotImplementedError()

    def calculate_parameter_covariance(self, pvals=None, *, robust=False):
        if pvals is None:
            pvals = self.pvals
        locks = np.asarray(self.pholdfast.astype(bool))
        if self.compute_engine == "jax":
            se, hess, ihess = self.jax_param_cov(pvals)
        else:
            raise NotImplementedError(f"compute_engine={self.compute_engine}")
            # hess = -self.d2_loglike(pvals)
            # if self.parameters["holdfast"].sum():
            #     free = self.pholdfast == 0
            #     hess_ = hess[free][:, free]
            #     ihess_ = np.linalg.inv(hess_)
            #     ihess = _arr_inflate(ihess_, locks)
            # else:
            #     ihess = np.linalg.inv(hess)
            # se = np.sqrt(ihess.diagonal())
            # self.pstderr = se
        hess = np.asarray(hess).copy()
        hess[locks, :] = 0
        hess[:, locks] = 0
        ihess = np.asarray(ihess).copy()
        ihess[locks, :] = 0
        ihess[:, locks] = 0
        self.add_parameter_array("hess", hess)
        self.add_parameter_array("ihess", ihess)

        # constrained covariance
        if self.constraints:
            constraints = list(self.constraints)
        else:
            constraints = []
        try:
            constraints.extend(self._get_bounds_constraints())
        except AttributeError:
            pass

        if constraints:
            binding_constraints = list()
            self.add_parameter_array("unconstrained_std_err", self.pstderr)
            self.add_parameter_array("unconstrained_covariance_matrix", ihess)

            s = np.asarray(ihess)
            pvals = self.pvals
            for c in constraints:
                if np.absolute(c.fun(pvals)) < c.binding_tol:
                    binding_constraints.append(c)
                    b = c.jac(self.pf.value)
                    den = b @ s @ b
                    if den != 0:
                        s = s - (1 / den) * s @ b.reshape(-1, 1) @ b.reshape(1, -1) @ s
            self.add_parameter_array("covariance_matrix", s)
            self.pstderr = np.sqrt(s.diagonal())

            # Fix numerical issues on some constraints, add constrained notes
            if binding_constraints or any(self.pholdfast != 0):
                notes = {}
                for c in binding_constraints:
                    pa = c.get_parameters()
                    for p in pa:
                        # if self.pf.loc[p, 't_stat'] > 1e5:
                        #     self.pf.loc[p, 't_stat'] = np.inf
                        #     self.pf.loc[p, 'std_err'] = np.nan
                        # if self.pf.loc[p, 't_stat'] < -1e5:
                        #     self.pf.loc[p, 't_stat'] = -np.inf
                        #     self.pf.loc[p, 'std_err'] = np.nan
                        n = notes.get(p, [])
                        n.append(c.get_binding_note(pvals))
                        notes[p] = n
                constrained_note = (
                    pd.Series({k: "\n".join(v) for k, v in notes.items()}, dtype=object)
                    .reindex(self.pnames)
                    .fillna("")
                )
                constrained_note[self.pholdfast != 0] = "fixed value"
                self.add_parameter_array("constrained", constrained_note)

        if robust:
            self.robust_covariance()
            se = self.parameters["robust_std_err"]

        return se, hess, ihess

    def robust_covariance(self):
        raise NotImplementedError()


class MixedLatentClass(LatentClass):
    mixtures = MixtureList()

    def __init__(
        self, *args, n_draws=100, prerolled_draws=True, common_draws=False, **kwargs
    ):
        super().__init__(*args, **kwargs)
        self._n_draws = n_draws
        self._draws = None
        self.prerolled_draws = prerolled_draws
        self.common_draws = common_draws

    @classmethod
    def from_dict(cls, content):
        self = super().from_dict(content)

        def loadthis(attr, wrapper=None, injector=None):
            i = content.get(attr, None)
            if i is not None:
                try:
                    if wrapper is not None:
                        i = wrapper(i)
                except AttributeError:
                    pass
                else:
                    if injector is None:
                        setattr(self, attr, i)
                    else:
                        injector(i)

        loadthis("mixtures", self.mixtures.from_list)
        loadthis("n_draws")
        loadthis("prerolled_draws")
        loadthis("common_draws")

    def apply_random_draws(self, parameters, draws=None):
        # if draws is None:
        #     draws = self._draws
        parameters = jnp.broadcast_to(
            parameters, [*draws.shape[:-1], parameters.shape[0]]
        )
        for mix_n, mix in enumerate(self.mixtures):
            u = draws[..., mix_n]
            parameters = mix.roll(u, parameters)
        return parameters

    def _make_random_draws_out(self, n_draws, n_mixtures, random_key):
        draws = jax.random.uniform(
            random_key,
            [n_draws, n_mixtures],
        )

        def body(i, carry):
            draws, rnd_key = carry
            rnd_key, subkey = jax.random.split(rnd_key)
            draws = draws.at[:, i].add(jax.random.permutation(subkey, draws.shape[0]))
            return draws, rnd_key

        draws, random_key = jax.lax.fori_loop(
            0, draws.shape[1], body, (draws, random_key)
        )
        return draws * (1 / n_draws), random_key

    @jitmethod
    def _jax_probability_by_class(self, params, databundle):
        # classmodel = self['classmodel']
        pr_parts = []
        for _n, k in enumerate(self.class_membership_model.dataset.dc.altids()):
            pr_parts.append(
                self._models[k]._jax_probability(params, databundle)
                # * jnp.expand_dims(classmodel._jax_probability(params, databundle)[..., n], -1)
            )
        return jnp.stack(pr_parts)

    @jitmethod
    def _jax_likelihood_by_class(self, params, databundle):
        n_alts = self.dataset.dc.n_alts
        ch = databundle.get("ch", None)[..., :n_alts]
        pr_parts = self._jax_probability_by_class(params, databundle)[..., :n_alts]
        likely = jnp.where(
            jnp.expand_dims(ch, 0), pr_parts, 1.0
        )  # TODO make power if needed
        return likely

    @jitmethod(static_argnums=(3,), static_argnames="n_draws")
    def _jax_loglike_casewise_mixed(
        self, params, databundle, groupbundle=None, n_draws=100
    ):
        classmodel = self.class_membership_model
        if self.prerolled_draws:
            draws = groupbundle.get("draws", None)
        else:
            rk = groupbundle.get("rk", None)
            draws, _ = self._make_random_draws_out(n_draws, len(self.mixtures), rk)
        ch = databundle.get("ch", None)
        rand_params = self.apply_random_draws(params, draws)
        if ch.ndim == 2:  # PANEL DATA
            # vmap over ingroup
            likelihood_f = jax.vmap(
                self._jax_likelihood_by_class,
                in_axes=(None, 0),
            )
            # vmap over draws
            likelihood = jax.vmap(
                likelihood_f,
                in_axes=(0, None),
                out_axes=-1,
            )(rand_params, databundle)
            # collapse likelihood over all alternatives over all cases-within-panel
            likelihood = likelihood.prod([0, 2])
            # likelihood.shape is now (n_classes, n_draws)
            ################
            class_pr = jax.vmap(
                classmodel._jax_probability,
                in_axes=(0, None),
                out_axes=-1,
            )(
                rand_params,
                {
                    "co": databundle["co"][0],
                },
            )
            # class_pr.shape = (nclasses, ndraws)
            meta_likely = (likelihood * class_pr).sum(0).mean(0)
            return jnp.log(meta_likely)
        else:
            # vmap over draws
            likelihood = jax.vmap(
                self._jax_likelihood_by_class,
                in_axes=(0, None),
                out_axes=-1,
            )(rand_params, databundle)
            # collapse likelihood over all alternatives
            likelihood = likelihood.prod(0)
            # average over all draws
            likelihood = likelihood.mean(0)
            return jnp.log(likelihood)

    @jitmethod
    def jax_loglike_casewise(self, params):
        if len(self.mixtures) == 0:
            return super().jax_loglike_casewise(params)
        ca = _get_jnp_array(self.dataset, "ca")
        co = _get_jnp_array(self.dataset, "co")
        av = _get_jnp_array(self.dataset, "av")
        ch = _get_jnp_array(self.dataset, "ch")
        n_draws = self._n_draws
        seed = getattr(self, "seed", 42)
        if av is not None:
            depth = av.ndim - 1
            shape = av.shape[:-1]
        elif co is not None:
            depth = co.ndim - 1
            shape = co.shape[:-1]
        elif ca is not None:
            depth = ca.ndim - 2
            shape = ca.shape[:-2]
        elif ch is not None:
            depth = ch.ndim - 1
            shape = ch.shape[:-1]
        else:
            raise ValueError("missing data")
        random_key = jax.random.PRNGKey(seed)
        if self.groupid is not None:
            depth = depth - 1
            shape = shape[:-1]
        f = (
            self._jax_loglike_casewise_mixed
        )  # params, databundle, groupbundle=None, n_draws=100
        from .random import keysplit

        commons = None if self.common_draws else 0
        for _i in range(depth):
            f = jax.vmap(f, in_axes=(None, 0, commons, None))
            if not self.prerolled_draws:
                random_key, shape = keysplit(random_key, shape)
        if self.prerolled_draws:
            return f(
                params,
                {"ca": ca, "co": co, "av": av, "ch": ch},
                {"draws": self._draws},
                n_draws,
            )
        else:
            return f(
                params,
                {"ca": ca, "co": co, "av": av, "ch": ch},
                {"rk": random_key},
                n_draws,
            )

    def make_random_draws(self, engine="numpy"):
        self.unmangle()
        for i in self.mixtures:
            i.prep(self)
        n_panels = self.dataset.dc.n_panels
        n_mixtures = len(self.mixtures)
        n_draws = self._n_draws
        draws = None
        if self._draws is not None:
            if self.common_draws and self._draws.shape == (n_draws, n_mixtures):
                draws = self._draws
            if not self.common_draws and self._draws.shape == (
                n_panels,
                n_draws,
                n_mixtures,
            ):
                draws = self._draws
        if draws is None:
            if engine == "numpy":
                seed = getattr(self, "seed", 0)
                if self.common_draws:
                    if n_draws > 0 and n_mixtures > 0:
                        draws, seed = self._make_random_draws_numpy(
                            n_draws, n_mixtures, seed
                        )
                else:
                    if n_draws > 0 and n_mixtures > 0 and n_panels > 0:
                        draws, seed = self._make_random_draws_numpy_2(
                            n_draws, n_mixtures, n_panels, seed
                        )
                    else:
                        draws = None
            elif engine == "jax":
                seed = getattr(self, "seed", 0)
                rk = jax.random.PRNGKey(seed)
                if self.common_draws:
                    if n_draws > 0 and n_mixtures > 0:
                        draws = self._make_random_draws_out(n_draws, n_mixtures, rk)[0]
                else:
                    if n_draws > 0 and n_mixtures > 0 and n_panels > 0:
                        draws = self._make_random_draws_out_2(
                            n_draws, n_mixtures, n_panels, rk
                        )[0]
                    else:
                        draws = None
            else:
                raise ValueError(f"unknown random engine {engine!r}")
            if self.prerolled_draws:
                self._draws = draws
        return draws

    # @jitmethod(static_argnums=(0,1,2), static_argnames=('n_draws', 'n_mixtures', 'n_cases'))
    def _make_random_draws_out_2(self, n_draws, n_mixtures, n_cases, random_key):
        def body(carry, x):
            rkey = carry
            draws, rkey = self._make_random_draws_out(n_draws, n_mixtures, rkey)
            return rkey, draws

        random_key, draws = jax.lax.scan(body, random_key, None, n_cases)
        return draws, random_key
