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

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.

max_no_attempts = 100


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 get_stats_dict(y, z, ywt, zwt, yn, zn, rg):
    """Define statistics needed for sampling."""
    sd = {}
    # 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.median(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.median(gs[gs > 0])])
    return g_guess


def prob_ra(ra, sd, by, bz):
    """Estimate probability of ra given reference strain and Null data."""
    ra = np.insert(np.asarray(ra), 0, sd["rguess"])
    ra = ra.reshape((ra.size, 1))
    alpha = find_alpha(
        ra,
        sd["sy_wt"],
        sd["sz_wt"],
        sd["n_wt"],
        by,
        bz,
        sd["Sy_wt"],
        sd["Ty_wt"],
        sd["Wy_wt"],
        sd["Sz_wt"],
        sd["Tz_wt"],
        sd["Wz_wt"],
    )
    minus_log_prob = mlp(alpha, ra, sd["sy_wt"], sd["sz_wt"], sd["n_wt"])
    # normalise by rguess and average over samples of background fluorescence
    raprobs = np.mean(np.exp(-minus_log_prob + minus_log_prob[0, :]), axis=1)
    # drop first element for rguess
    return raprobs[1:]


def find_alpha(r, sy, sz, n, ty, tz, Sy, Ty, Wy, Sz, Tz, Wz):
    """Find exponent after integrating out autofluorescence."""
    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."""
    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 prob_g(g, sd, by, bz, rg, ra_samples):
    """Estimate probability of GFP by averaging samples of ra."""
    g = np.insert(np.asarray(g), 0, sd["gguess"])
    if ra_samples.size * by.size * g.size < 1e8:
        # broadcast to 3D
        ra_samples = ra_samples.reshape((ra_samples.size, 1, 1))
        g = g.reshape((g.size, 1))
        alpha = find_alpha(
            ra_samples,
            sd["sy"],
            sd["sz"],
            sd["n"],
            g + by,
            rg * g + bz,
            sd["Sy"],
            sd["Ty"],
            sd["Wy"],
            sd["Sz"],
            sd["Tz"],
            sd["Wz"],
        )
        minus_log_prob = mlp(alpha, ra_samples, sd["sy"], sd["sz"], sd["n"])
        # normalise by prob for gguess
        norm = minus_log_prob[:, 0, :].reshape(ra_samples.size, 1, by.size)
        # average over samples for ra and background fluorescence
        gprobs = np.mean(np.exp(-minus_log_prob + norm), axis=(0, 2))
    else:
        # calculate probability of each g
        gprobs = np.empty(g.size)
        for i, gv in enumerate(g):
            alpha = find_alpha(
                ra_samples,
                sd["sy"],
                sd["sz"],
                sd["n"],
                gv + by,
                rg * gv + bz,
                sd["Sy"],
                sd["Ty"],
                sd["Wy"],
                sd["Sz"],
                sd["Tz"],
                sd["Wz"],
            )
            minus_log_prob = mlp(
                alpha, ra_samples, sd["sy"], sd["sz"], sd["n"]
            )
            if i == 0:
                norm = minus_log_prob
                gprobs[0] = 1
            else:
                # average over samples for ra and background fluorescence
                gprobs[i] = np.mean(np.exp(-minus_log_prob + norm))
    # drop first element for gguess
    return gprobs[1:]


def rejection_sample(nosamples, negprob, xmin, xmax):
    res = minimize_scalar(
        lambda x: negprob(x), bounds=(xmin, xmax), method="bounded"
    )
    pmax = -res.fun
    xs = np.empty(0)
    i = 0
    while (xs.size < nosamples) & (i < max_no_attempts):
        rands = rng.random((2 * nosamples, 2))
        x_try = xmin + (xmax - xmin) * rands[:, 0]
        probs = -negprob(x_try)
        xs = np.append(xs, x_try[pmax * rands[:, 1] < probs])
        i += 1
    xs = xs.reshape((xs.size, 1))
    return xs


def sample_g(nosamples_for_fl, sd, by, bz, rg, nosamples_for_ra):
    bnd_factor = 10
    # sample ra
    rmin = sd["rguess"] / bnd_factor
    rmax = bnd_factor * sd["rguess"]
    ra_samples = rejection_sample(
        nosamples_for_ra, lambda x: -prob_ra(x, sd, by, bz), rmin, rmax
    )
    # sample g given ra
    gmin = sd["gguess"] / bnd_factor
    gmax = bnd_factor * sd["gguess"]
    g_samples = rejection_sample(
        nosamples_for_fl,
        lambda x: -prob_g(x, sd, by, bz, rg, ra_samples),
        gmin,
        gmax,
    )
    g_samples = g_samples.reshape(g_samples.size)
    # return fixed number of samples
    if g_samples.size < nosamples_for_fl:
        with_nans = np.nan * np.ones(nosamples_for_fl)
        with_nans[np.where(g_samples)] = g_samples
        return with_nans
    else:
        return g_samples[:nosamples_for_fl]


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 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)}
    bdn = gu.mergedicts(original=bd_default, update=bd)
    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.empty((t.size, nosamples_for_fl))
                    for i in tqdm(range(t.size)):
                        sd = get_stats_dict(
                            y[i, :],
                            z[i, :],
                            ywt[i, :],
                            zwt[i, :],
                            yn[i, :],
                            zn[i, :],
                            rg,
                        )
                        by, bz = get_background_samples(
                            yn[i, :], zn[i, :], nosamples_for_bg
                        )
                        fl_samples = sample_g(
                            nosamples_for_fl, sd, by, bz, rg, nosamples_for_ra
                        )
                        predicted_fl[i, :] = fl_samples
                    # smooth with GP and add to data frames
                    print("Smoothing...")
                    breakpoint()
                    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 = 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)
