# A Bayesian version of the method in Lichten ... Swain, BMC Biophys 2014
import numpy as np
from scipy.optimize import minimize

import omniplate.corrections as omcorr
import omniplate.omgenutils as gu
import omniplate.sunder as sunder
from omniplate.runfitderiv import runfitderiv

rng = np.random.default_rng()

# notation follows Lichten ... Swain
# GFP is denoted y; AutoFL is denoted z.
# The reference strain is denoted WT.


def de_nan(y, z):
    """Remove any replicates with NaN."""
    # NaNs are generated because experiments have different durations
    keep = ~np.any(np.isnan(y), axis=0)
    return y[:, keep], z[:, keep]


def sample_b(nosamples, bdata):
    """Sample background fluorescence."""
    s2 = np.var(bdata) / bdata.size
    u = np.mean(bdata)
    samples = u + s2 * rng.standard_normal(nosamples)
    if np.any(samples < 0):
        print("Warning: negative background fluorescence.")
    return samples


def get_background_samples(yn, zn, nosamples):
    """Get samples of background fluorescence for GFP and AutoFL."""
    by = sample_b(nosamples, yn)
    bz = sample_b(nosamples, zn)
    return by, bz


def update_stats_dict(sd, y, z, ywt, zwt, yn, zn, rg, by, bz):
    """Update with statistics needed for sampling."""
    sd["rg"] = rg
    sd["by"] = by
    sd["bz"] = bz
    # for wild type strain
    sd["Sy_wt"], sd["Ty_wt"], sd["Wy_wt"] = get_sufficient_stats(ywt)
    sd["Sz_wt"], sd["Tz_wt"], sd["Wz_wt"] = get_sufficient_stats(zwt)
    sd["sy_wt"] = np.std(ywt)
    sd["sz_wt"] = np.std(zwt)
    sd["n_wt"] = ywt.size
    sd["rguess"] = ra_guess(ywt, zwt)
    # for tagged strain
    sd["Sy"], sd["Ty"], sd["Wy"] = get_sufficient_stats(y)
    sd["Sz"], sd["Tz"], sd["Wz"] = get_sufficient_stats(z)
    sd["sy"] = np.std(y)
    sd["sz"] = np.std(z)
    sd["n"] = y.size
    sd["gguess"] = g_guess(y, z, ywt, zwt, rg)
    return sd


def get_sufficient_stats(y):
    """Get sufficient statistics generated by integrating autofluorescence."""
    S = np.sum(y**2)
    T = np.sum(y)
    W = np.sum(
        [y[i] * y[j] for i in range(y.size) for j in range(i + 1, y.size)]
    )
    return S, T, W


def ra_guess(ywt, zwt):
    """Get mean-field estimate of ra."""
    return np.mean(zwt / ywt)


def g_guess(y, z, ywt, zwt, rg):
    """Get mean-field estimate of GFP."""
    ra = ra_guess(ywt, zwt)
    gs = (ra * y - z) / (ra - rg)
    g_guess = np.nanmax([0.01, np.mean(gs[gs > 0])])
    return g_guess


def find_alpha(r, sy, sz, n, ty, tz, Sy, Ty, Wy, Sz, Tz, Wz):
    """Find exponent after integrating out autofluorescence."""
    r = np.asarray(r)
    if r.size > 1:
        r = r[:, None]
    alpha = (
        sz**4 * ((n - 1) * Sy - 2 * Wy)
        + r**2 * sy**4 * ((n - 1) * Sz - 2 * Wz)
        + sy**2
        * sz**2
        * (
            n * r**2 * (n * ty**2 + Sy - 2 * ty * Ty)
            + n * (n * tz**2 + Sz - 2 * tz * Tz)
            - 2 * r * (n * ty - Ty) * (n * tz - Tz)
        )
    )
    return alpha


def mlp(alpha, r, sy, sz, n):
    """Find minus log of probability after integrating autofluorescence."""
    r = np.asarray(r)
    if r.size > 1:
        r = r[:, None]
    mlp = (
        alpha / (2 * n * sy**2 * sz**2 * (r**2 * sy**2 + sz**2))
        + (n - 1) * np.log(sy)
        + (n - 1) * np.log(sz)
        + np.log(r**2 * sy**2 + sz**2) / 2
    )
    return mlp


def minus_log_int_prob_ra(ra, stats_dict):
    """Estimate interior probability of ra."""
    ra = np.asarray(ra)
    alpha = find_alpha(
        ra,
        stats_dict["sy_wt"],
        stats_dict["sz_wt"],
        stats_dict["n_wt"],
        stats_dict["by"],
        stats_dict["bz"],
        stats_dict["Sy_wt"],
        stats_dict["Ty_wt"],
        stats_dict["Wy_wt"],
        stats_dict["Sz_wt"],
        stats_dict["Tz_wt"],
        stats_dict["Wz_wt"],
    )
    minus_log_prob = mlp(
        alpha,
        ra,
        stats_dict["sy_wt"],
        stats_dict["sz_wt"],
        stats_dict["n_wt"],
    )
    return minus_log_prob


def minus_log_int_prob_g_ra(g, ra, stats_dict):
    """Estimate interior probability for a scalar g and ra."""
    alpha = find_alpha(
        ra,
        stats_dict["sy"],
        stats_dict["sz"],
        stats_dict["n"],
        g + stats_dict["by"],
        stats_dict["rg"] * g + stats_dict["bz"],
        stats_dict["Sy"],
        stats_dict["Ty"],
        stats_dict["Wy"],
        stats_dict["Sz"],
        stats_dict["Tz"],
        stats_dict["Wz"],
    )
    minus_log_prob = mlp(
        alpha,
        ra,
        stats_dict["sy"],
        stats_dict["sz"],
        stats_dict["n"],
    )
    return minus_log_prob


def minus_log_prob_g_ra_array(theta, stats_dict):
    """Find joint probability of a scalar g and ra for each by and bz."""
    g, ra = theta
    mlp = minus_log_int_prob_g_ra(g, ra, stats_dict) + minus_log_int_prob_ra(
        ra, stats_dict
    )
    return mlp


def minus_log_prob(theta, stats_dict):
    """Find joint probability of a scalar g and ra averaged over background."""
    mlp = minus_log_prob_g_ra_array(theta, stats_dict)
    minus_log_prob = average_background_fluorescence(mlp)
    return minus_log_prob


def average_background_fluorescence(mlp, deriv=None):
    """Average background fluorescence."""
    norm = np.min(mlp)
    # normalise to prevent underflows, exponentiate, then average
    norm_prob = np.exp(norm - mlp)
    if deriv is None:
        # revert normalisation and return -log(probability)
        minus_log_prob = norm - np.log(np.mean(norm_prob))
        return minus_log_prob
    else:
        # return weighted average of a derivative array
        deriv_average = np.sum(deriv * norm_prob) / np.sum(norm_prob)
        return deriv_average


def find_mode(stats_dict, bnd_factor=100):
    """Find most probable value of g and ra."""
    rmin = stats_dict["rguess"] / bnd_factor
    rmax = stats_dict["rguess"] * bnd_factor
    gmin = stats_dict["gguess"] / bnd_factor
    gmax = stats_dict["gguess"] * bnd_factor
    # aside
    # res = minimize(
    #     lambda x: average_background_fluorescence(
    #         minus_log_int_prob_ra(x, stats_dict)
    #     ),
    #     x0=stats_dict["rguess"],
    #     bounds=((rmin, rmax),),
    #     method="L-BFGS-B",
    # )
    # breakpoint()
    bounds = ((gmin, gmax), (rmin, rmax))
    if stats_dict["seq_prior_mu"] is None:
        x0 = [stats_dict["gguess"], stats_dict["rguess"]]
    else:
        x0 = stats_dict["seq_prior_mu"]
    min_mlp = np.inf
    mode = None
    no_attempts = 10
    rands = rng.standard_normal((no_attempts, 2))
    for i in range(no_attempts):
        x0 = rands[i, :] * np.sqrt(x0) + x0
        res = minimize(
            lambda x: minus_log_prob(x, stats_dict)
            + minus_log_prior(x, stats_dict),
            x0=x0,
            bounds=bounds,
            jac=lambda x: jac(x, stats_dict) + jac_prior(x, stats_dict),
            hess=lambda x: hess(x, stats_dict) + hess_prior(x, stats_dict),
            method="L-BFGS-B",
        )
        if res.success:
            if res.fun < min_mlp:
                mode = res.x
                for j in range(2):
                    if mode[j] == bounds[j][0] or mode[j] == bounds[j][1]:
                        print(f"parameter {j} at bound {mode[j]}")
                stats_dict["seq_prior_mu"] = mode
                stats_dict["seq_prior_hess"] = hess(
                    mode, stats_dict
                ) + hess_prior(mode, stats_dict)
                min_mlp = res.fun
    if mode is None:
        breakpoint()
    return mode


def correctauto_bayesian(
    self,
    f,
    refstrain,
    flcvfn,
    bd,
    max_data_pts,
    nosamples_for_fl,
    nosamples_for_ra,
    nosamples_for_bg,
    nosamples,
    experiments,
    experimentincludes,
    experimentexcludes,
    conditions,
    conditionincludes,
    conditionexcludes,
    strains,
    strainincludes,
    strainexcludes,
):
    """
    Correct fluorescence for auto- and background fluorescence.

    Use a Bayesian method to correct for autofluorescence from fluorescence
    measurements at two wavelengths and for background fluorescence.

    Implement demixing following Lichten ... Swain, BMC Biophys 2014.

    Integrate over autofluorescence exactly and other nuisance variable
    by sampling.
    """
    print("Using Bayesian approach for two fluorescence wavelengths.")
    print(f"Correcting autofluorescence using {f[0]} and {f[1]}.")
    rg = self._gamma
    bname = "bc" + f[0]
    bd_default = {0: (2, 8), 1: (-2, 4), 2: (5, 8)}
    if bd is not None:
        bdn = gu.mergedicts(original=bd_default, update=bd)
    else:
        bdn = bd_default
    for e in sunder.getset(
        self,
        experiments,
        experimentincludes,
        experimentexcludes,
        "experiment",
        nonull=True,
    ):
        for c in sunder.getset(
            self,
            conditions,
            conditionincludes,
            conditionexcludes,
            labeltype="condition",
            nonull=True,
            nomedia=True,
        ):
            # get data for reference strain
            # y for emission at 525; z for emission at 585
            _, (ywt, zwt) = sunder.extractwells(
                self.r, self.s, e, c, refstrain, f
            )
            ywt, zwt = de_nan(ywt, zwt)
            # get data for Null
            _, (yn, zn) = sunder.extractwells(self.r, self.s, e, c, "Null", f)
            yn, zn = de_nan(yn, zn)
            # check sufficient replicates
            if (ywt.shape[1] < 3) or (yn.shape[1] < 3):
                raise Exception(
                    f"There are less than three replicates for the {refstrain}"
                    " or Null strains."
                )
            for s in sunder.getset(
                self,
                strains,
                strainincludes,
                strainexcludes,
                labeltype="strain",
                nonull=True,
            ):
                if (
                    s != refstrain
                    and f"{s} in {c}" in self.allstrainsconditions[e]
                ):
                    # get data for tagged strain
                    t, (y, z, od) = sunder.extractwells(
                        self.r, self.s, e, c, s, f.copy() + ["OD"]
                    )
                    y, z = de_nan(y, z)
                    if y.size == 0 or z.size == 0:
                        print(f"Warning: No data found for {e}: {s} in {c}!!")
                        continue
                    if y.shape[1] < 3:
                        raise Exception(
                            "There are less than three replicates"
                            f" for {e}: {s} in {c}."
                        )
                    print(f" {e}: {s} in {c}")
                    # correct autofluorescence for each time point
                    predicted_fl = np.zeros(t.size)
                    # for sequential priors
                    stats_dict = {"seq_prior_mu": None, "seq_prior_hess": None}
                    for i in range(t.size):
                        by, bz = get_background_samples(
                            yn[i, :], zn[i, :], nosamples_for_bg
                        )
                        if i == 10:
                            breakpoint()
                        stats_dict = update_stats_dict(
                            stats_dict,
                            y[i, :],
                            z[i, :],
                            ywt[i, :],
                            zwt[i, :],
                            yn[i, :],
                            zn[i, :],
                            rg,
                            by,
                            bz,
                        )
                        predicted_fl[i] = find_mode(stats_dict)[0]
                    # smooth with GP and add to data frames
                    if False:
                        print("Smoothing...")
                        flgp, _ = runfitderiv(
                            self,
                            t,
                            predicted_fl,
                            f"{bname}",
                            f"d/dt_{bname}",
                            e,
                            c,
                            s,
                            bd=bdn,
                            cvfn=flcvfn,
                            logs=True,
                            figs=True,
                            max_data_pts=max_data_pts,
                        )
                    # sample ODs using GPs
                    lod_samples = omcorr.sample_ODs_with_GP(
                        self, e, c, s, t, od, nosamples
                    )
                    od_samples = np.exp(lod_samples)
                    # find samples of fluorescence per OD
                    smooth_fl_samples = predicted_fl[:, None]
                    # smooth_fl_samples = flgp.fitderivsample(nosamples)[0]
                    flperod = smooth_fl_samples / od_samples
                    # store results for fluorescence per OD
                    autofdict = {
                        "experiment": e,
                        "condition": c,
                        "strain": s,
                        "time": t,
                        f"{bname}perOD": np.nanmean(flperod, 1),
                        f"{bname}perOD_err": omcorr.nanstdzeros2nan(
                            flperod, 1
                        ),
                    }
                    # add to data frames
                    omcorr.addtodataframes(self, autofdict, bname)


def jac_arrays(theta, stats_dict):
    """Find Jacobian of minus_log_total_prob for each by and bz."""
    g, ra = theta
    rg = stats_dict["rg"]
    by = stats_dict["by"]
    bz = stats_dict["bz"]
    Sywt = stats_dict["Sy_wt"]
    Tywt = stats_dict["Ty_wt"]
    Wywt = stats_dict["Wy_wt"]
    Szwt = stats_dict["Sz_wt"]
    Tzwt = stats_dict["Tz_wt"]
    Wzwt = stats_dict["Wz_wt"]
    sywt = stats_dict["sy_wt"]
    szwt = stats_dict["sz_wt"]
    Sy = stats_dict["Sy"]
    Ty = stats_dict["Ty"]
    Wy = stats_dict["Wy"]
    Sz = stats_dict["Sz"]
    Tz = stats_dict["Tz"]
    Wz = stats_dict["Wz"]
    sy = stats_dict["sy"]
    sz = stats_dict["sz"]
    nwt = stats_dict["n_wt"]
    n = stats_dict["n"]
    # define jacobian over background fluorescence
    jac_v = np.empty(2, dtype="object")
    jac_v[0] = (
        (ra - rg)
        * (-(bz * n) + by * n * ra + g * n * ra - g * n * rg - ra * Ty + Tz)
    ) / (ra**2 * sy**2 + sz**2)
    jac_v[1] = (
        (ra * sy**2) / (ra**2 * sy**2 + sz**2)
        + (ra * sywt**2) / (ra**2 * sywt**2 + szwt**2)
        + (
            sy**2
            * sz**2
            * (
                2 * n * ra * ((by + g) ** 2 * n + Sy - 2 * (by + g) * Ty)
                - 2 * ((by + g) * n - Ty) * (bz * n + g * n * rg - Tz)
            )
            + 2 * ra * sy**4 * ((-1 + n) * Sz - 2 * Wz)
        )
        / (2.0 * n * sy**2 * sz**2 * (ra**2 * sy**2 + sz**2))
        - (
            ra
            * (
                sy**2
                * sz**2
                * (
                    n * ra**2 * ((by + g) ** 2 * n + Sy - 2 * (by + g) * Ty)
                    - 2 * ra * ((by + g) * n - Ty) * (bz * n + g * n * rg - Tz)
                    + n
                    * (n * (bz + g * rg) ** 2 + Sz - 2 * (bz + g * rg) * Tz)
                )
                + sz**4 * ((-1 + n) * Sy - 2 * Wy)
                + ra**2 * sy**4 * ((-1 + n) * Sz - 2 * Wz)
            )
        )
        / (n * (ra**2 * sy**2 * sz + sz**3) ** 2)
        + (
            sywt**2
            * szwt**2
            * (
                2 * nwt * ra * (by**2 * nwt + Sywt - 2 * by * Tywt)
                - 2 * (by * nwt - Tywt) * (bz * nwt - Tzwt)
            )
            + 2 * ra * sywt**4 * ((-1 + nwt) * Szwt - 2 * Wzwt)
        )
        / (2.0 * nwt * sywt**2 * szwt**2 * (ra**2 * sywt**2 + szwt**2))
        - (
            ra
            * (
                sywt**2
                * szwt**2
                * (
                    nwt * ra**2 * (by**2 * nwt + Sywt - 2 * by * Tywt)
                    - 2 * ra * (by * nwt - Tywt) * (bz * nwt - Tzwt)
                    + nwt * (bz**2 * nwt + Szwt - 2 * bz * Tzwt)
                )
                + szwt**4 * ((-1 + nwt) * Sywt - 2 * Wywt)
                + ra**2 * sywt**4 * ((-1 + nwt) * Szwt - 2 * Wzwt)
            )
        )
        / (nwt * (ra**2 * sywt**2 * szwt + szwt**3) ** 2)
    )
    return jac_v


def jac(theta, stats_dict, return_jac_v=False):
    """Get Jacobian averaged over background fluorescence."""
    jac_v = jac_arrays(theta, stats_dict)
    # average background fluorescence
    mlp = minus_log_prob_g_ra_array(theta, stats_dict)
    av_jac = np.empty(2)
    av_jac[0] = average_background_fluorescence(mlp, jac_v[0])
    av_jac[1] = average_background_fluorescence(mlp, jac_v[1])
    if return_jac_v:
        return av_jac, jac_v
    else:
        return av_jac


def hess_arrays(theta, stats_dict):
    """Find Hessian of minus_log_total_prob for each by and bz."""
    g, ra = theta
    rg = stats_dict["rg"]
    by = stats_dict["by"]
    bz = stats_dict["bz"]
    Sywt = stats_dict["Sy_wt"]
    Tywt = stats_dict["Ty_wt"]
    Wywt = stats_dict["Wy_wt"]
    Szwt = stats_dict["Sz_wt"]
    Tzwt = stats_dict["Tz_wt"]
    Wzwt = stats_dict["Wz_wt"]
    sywt = stats_dict["sy_wt"]
    szwt = stats_dict["sz_wt"]
    Sy = stats_dict["Sy"]
    Ty = stats_dict["Ty"]
    Wy = stats_dict["Wy"]
    Sz = stats_dict["Sz"]
    Tz = stats_dict["Tz"]
    Wz = stats_dict["Wz"]
    sy = stats_dict["sy"]
    sz = stats_dict["sz"]
    nwt = stats_dict["n_wt"]
    n = stats_dict["n"]
    # define hessian averaging over background fluorescence
    hess_v = np.empty((2, 2), dtype="object")
    hess_v[0, 0] = (n * (ra - rg) ** 2) / (ra**2 * sy**2 + sz**2)
    hess_v[0, 1] = (
        2 * g * n * ra**2 * rg * sy**2
        - 2 * g * n * ra * rg**2 * sy**2
        + 2 * g * n * ra * sz**2
        - 2 * g * n * rg * sz**2
        + bz * n * (ra**2 * sy**2 - 2 * ra * rg * sy**2 - sz**2)
        + by * n * (ra**2 * rg * sy**2 + 2 * ra * sz**2 - rg * sz**2)
        - ra**2 * rg * sy**2 * Ty
        - 2 * ra * sz**2 * Ty
        + rg * sz**2 * Ty
        - ra**2 * sy**2 * Tz
        + 2 * ra * rg * sy**2 * Tz
        + sz**2 * Tz
    ) / (ra**2 * sy**2 + sz**2) ** 2
    hess_v[1, 0] = hess_v[0, 1]
    hess_v[1, 1] = (
        (-2 * ra**2 * sy**4) / (ra**2 * sy**2 + sz**2) ** 2
        + sy**2 / (ra**2 * sy**2 + sz**2)
        - (2 * ra**2 * sywt**4) / (ra**2 * sywt**2 + szwt**2) ** 2
        + sywt**2 / (ra**2 * sywt**2 + szwt**2)
        + (
            (ra**2 * sy**2 + sz**2) ** 2
            * (
                2
                * n
                * sy**2
                * sz**2
                * ((by + g) ** 2 * n + Sy - 2 * (by + g) * Ty)
                + 2 * sy**4 * ((-1 + n) * Sz - 2 * Wz)
            )
            - 4
            * ra
            * sy**2
            * (ra**2 * sy**2 + sz**2)
            * (
                sy**2
                * sz**2
                * (
                    2 * n * ra * ((by + g) ** 2 * n + Sy - 2 * (by + g) * Ty)
                    - 2 * ((by + g) * n - Ty) * (bz * n + g * n * rg - Tz)
                )
                + 2 * ra * sy**4 * ((-1 + n) * Sz - 2 * Wz)
            )
            + 2
            * (3 * ra**2 * sy**4 - sy**2 * sz**2)
            * (
                sy**2
                * sz**2
                * (
                    n * ra**2 * ((by + g) ** 2 * n + Sy - 2 * (by + g) * Ty)
                    - 2 * ra * ((by + g) * n - Ty) * (bz * n + g * n * rg - Tz)
                    + n
                    * (n * (bz + g * rg) ** 2 + Sz - 2 * (bz + g * rg) * Tz)
                )
                + sz**4 * ((-1 + n) * Sy - 2 * Wy)
                + ra**2 * sy**4 * ((-1 + n) * Sz - 2 * Wz)
            )
        )
        / (2.0 * n * sy**2 * sz**2 * (ra**2 * sy**2 + sz**2) ** 3)
        + (
            -3 * ra**2 * sywt**2 * Sywt * szwt**2
            + Sywt * szwt**4
            + bz**2 * nwt**2 * sywt**2 * (3 * ra**2 * sywt**2 - szwt**2)
            + by**2 * nwt**2 * szwt**2 * (-3 * ra**2 * sywt**2 + szwt**2)
            + 3 * ra**2 * sywt**4 * Szwt
            - sywt**2 * szwt**2 * Szwt
            - 2 * ra**3 * sywt**4 * Tywt * Tzwt
            + 6 * ra * sywt**2 * szwt**2 * Tywt * Tzwt
            + 2
            * bz
            * nwt
            * sywt**2
            * (
                -(by * nwt * ra**3 * sywt**2)
                + 3 * by * nwt * ra * szwt**2
                + ra**3 * sywt**2 * Tywt
                - 3 * ra * szwt**2 * Tywt
                - 3 * ra**2 * sywt**2 * Tzwt
                + szwt**2 * Tzwt
            )
            + 2
            * by
            * nwt
            * (
                3 * ra**2 * sywt**2 * szwt**2 * Tywt
                - szwt**4 * Tywt
                + ra**3 * sywt**4 * Tzwt
                - 3 * ra * sywt**2 * szwt**2 * Tzwt
            )
            - 6 * ra**2 * sywt**2 * szwt**2 * Wywt
            + 2 * szwt**4 * Wywt
            + 6 * ra**2 * sywt**4 * Wzwt
            - 2 * sywt**2 * szwt**2 * Wzwt
        )
        / (nwt * (ra**2 * sywt**2 + szwt**2) ** 3)
    )
    return hess_v


def hess(theta, stats_dict):
    """Get Hessian averaged over background fluorescence."""
    mlp = minus_log_prob_g_ra_array(theta, stats_dict)
    av_jac, jac_v = jac(theta, stats_dict, return_jac_v=True)
    hess_v = hess_arrays(theta, stats_dict)
    av_hess = np.empty((2, 2))
    for i in range(2):
        for j in range(i, 2):
            av_hess[i, j] = (
                average_background_fluorescence(mlp, hess_v[i, j])
                - average_background_fluorescence(mlp, jac_v[0] * jac_v[1])
                - av_jac[0] * av_jac[1]
            )
    av_hess[1, 0] = av_hess[0, 1]
    return av_hess


def minus_log_prior(theta, stats_dict):
    """Get minus log of the prior probability."""
    g, ra = theta
    # mlp = np.log(g) + np.log(ra)
    mlp = 0
    if stats_dict["seq_prior_hess"] is not None:
        mu = stats_dict["seq_prior_mu"]
        h = stats_dict["seq_prior_hess"]
        mlp += (
            (g - mu[0]) ** 2 * h[0, 0]
            + 2 * (g - mu[0]) * (ra - mu[1]) * h[0, 1]
            + (ra - mu[1]) ** 2 * h[1, 1]
        ) / 2
    return mlp


def jac_prior(theta, stats_dict):
    """Get Jacobian for -log prior."""
    g, ra = theta
    jac = np.empty(2)
    jac[0] = 0  # 1 / g
    jac[1] = 0  # 1 / ra
    if stats_dict["seq_prior_mu"] is not None:
        mu = stats_dict["seq_prior_mu"]
        h = stats_dict["seq_prior_hess"]
        jac[0] += g * h[0, 0] - h[0, 0] * mu[0] + h[0, 1] * (-mu[1] + ra)
        jac[1] += (
            g * h[0, 1] - h[0, 1] * mu[0] - h[1, 1] * mu[1] + h[1, 1] * ra
        )
    return jac


def hess_prior(theta, stats_dict):
    """Get Hessian for -log prior."""
    g, ra = theta
    hess = np.empty((2, 2))
    hess[0, 0] = 0  # -1 / g**2
    hess[1, 1] = 0  # -1 / ra**2
    hess[0, 1] = 0
    if stats_dict["seq_prior_hess"] is not None:
        h = stats_dict["seq_prior_hess"]
        hess[0, 0] += h[0, 0]
        hess[1, 1] += h[1, 1]
        hess[0, 1] += h[0, 1]
    hess[1, 0] = hess[0, 1]
    return hess
