"""Miscellaneous functions."""

import gaussianprocessderivatives as gp
import matplotlib.pylab as plt
import numpy as np
from scipy import integrate
from scipy.interpolate import interp1d

import omniplate.clogger as clogger
import omniplate.omgenutils as gu
import omniplate.sunder as sunder
from omniplate.corrections import get_hypers, instantiate_GP
from omniplate.omstats import stats_err


@clogger.log
def average_over_expts(
    self,
    condition,
    strain,
    tvr="mean_OD",
    bd=False,
    add_noise=True,
    plot=False,
):
    """
    Average a time-dependent variable over all experiments.

    Uses a Matern Gaussian process.

    An alternative and best first choice is add_common_variable.

    Parameters
    ----------
    condition: string
        The condition of interest.
    strain: string
        The strain of interest.
    tvr: float
        The time-dependent variable to be averaged.
        For example, 'cGFPperOD' or 'mean_OD'.
    bd: dictionary, optional
        The limits on the hyperparameters for the Matern Gaussian process.
        For example, {0: (-5,5), 1: (-4,4), 2: (-5,2)}
        where the first element controls amplitude, setting the bounds to
        1e-5 and 1e5, the second controls flexibility, and the third
        determines the magnitude of the measurement error.
    add_noise: boolean
        If True, add the fitted magnitude of the measurement noise to the
        predicted standard deviation for better comparison with the spread
        of the data.

    Returns
    -------
    res: dictionary
        {'t' : time, tvr : time-dependent data, 'mn' : mean,
        'sd' : standard deviation}
        where 'mn' is the average found and 'sd' is its standard deviation.
        'tvr' is the data used to find the average.

    Examples
    --------
    >>> p.average_over_expts('1% Gal', 'GAL2', bd= {1: [-1,-1])})
    """
    # boundaries on hyperparameters
    if "OD" in tvr:
        bds = {0: (-4, 4), 1: (-1, 4), 2: (-6, 2)}
    else:
        bds = {0: (2, 12), 1: (-1, 4), 2: (4, 10)}
    if bd:
        bds = gu.merge_dicts(original=bds, update=bd)
    # extract data
    df = self.s[["experiment", "condition", "strain", "time", tvr]]
    ndf = df.query("condition == @condition and strain == @strain")
    # use GP to average over experiments
    x = ndf["time"].to_numpy()
    y = ndf[tvr].to_numpy()
    ys = y[np.argsort(x)]
    xs = np.sort(x)
    g = gp.maternGP(bds, xs, ys)
    print(f"Averaging over {tvr} experiments for {strain} in {condition}.")
    g.findhyperparameters(noruns=2, noinits=1000)
    g.results()
    g.predict(xs, addnoise=add_noise, derivs=1)
    if plot:
        plt.figure()
        g.sketch(".")
        plt.title("averaging " + strain + " in " + condition)
        plt.xlabel("time")
        plt.ylabel(tvr)
        plt.show(block=False)
    # return results as a dictionary
    res = {"t": xs, tvr: ys, "mn": g.f, "sd": np.sqrt(g.fvar)}
    return res


def _deduplicate_xy(x, y):
    """
    Return unique x values with the median y at each.

    Remove duplicate x values caused by OD plateauing by sorting,
    grouping, and taking the median y within each group.

    Parameters
    ----------
    x: numpy array
        The x values, potentially with duplicates.
    y: numpy array
        The corresponding y values.

    Returns
    -------
    ux: numpy array
        Sorted unique x values.
    uy: numpy array
        Median y value for each unique x.
    """
    order = np.argsort(x)
    xs, ys = x[order], y[order]
    ux, split_idx = np.unique(xs, return_index=True)
    groups = np.split(ys, split_idx[1:])
    uy = np.array([np.median(g) for g in groups])
    return ux, uy


def _get_fitness_samples(self, ref, com, yvar, no_samples):
    """
    Get sample data for the fitness penalty calculation.

    Either sample from Gaussian processes (when yvar is 'gr' and
    no_samples > 0) or extract data directly from the DataFrame.

    Parameters
    ----------
    self: omniplate.platereader instance
        An instance of platereader with data loaded.
    ref: list of strings
        A list of three strings: experiment, condition, and strain
        for the reference.
    com: list of strings
        A list of three strings: experiment, condition, and strain
        for the comparison.
    yvar: string
        The variable to be compared.
    no_samples: integer
        The number of bootstraps used to estimate the error.

    Returns
    -------
    xsref: numpy array
        x values (OD) for the reference, with shape (n, no_samples).
    ysref: numpy array
        y values for the reference, with shape (n, no_samples).
    xscom: numpy array
        x values (OD) for the comparison, with shape (n, no_samples).
    yscom: numpy array
        y values for the comparison, with shape (n, no_samples).
    """
    if no_samples and yvar == "gr":
        # instantiate GPs to estimate errors
        for ecs in [ref, com]:
            e, c, s = ecs
            hypers, cvfn = get_hypers(self, e, c, s, yvar)
            if hypers is None or cvfn is None:
                raise ValueError(
                    "\nYou must first run get_stats" " or set no_samples=0."
                )
            t, od = sunder.extract_wells(self.r, self.s, e, c, s, "OD")
            if ecs == ref:
                g_ref = instantiate_GP(hypers, cvfn, t, od, max_data_pts=None)
            else:
                g_com = instantiate_GP(hypers, cvfn, t, od, max_data_pts=None)
        # sample from GPs
        f0s, g0s = g_ref.sample(no_samples, derivs=1)
        f1s, g1s = g_com.sample(no_samples, derivs=1)
        return np.exp(f0s), g0s, np.exp(f1s), g1s
    # no estimates of errors
    if no_samples:
        print(
            "Cannot estimate errors - require y= 'gr'"
            " and a recently run get_stats"
        )
    ref_df = self.s.query(
        "experiment == @ref[0] and condition == @ref[1]"
        " and strain == @ref[2]"
    )
    com_df = self.s.query(
        "experiment == @com[0] and condition == @com[1]"
        " and strain == @com[2]"
    )
    xsref = ref_df["OD mean"].to_numpy()[:, None]
    ysref = ref_df[yvar].to_numpy()[:, None]
    xscom = com_df["OD mean"].to_numpy()[:, None]
    yscom = com_df[yvar].to_numpy()[:, None]
    if xsref.size == 0 or ysref.size == 0:
        print(f"{ref[0]}: Data missing for {ref[2]} in {ref[1]}")
        return None, None, None, None
    elif xscom.size == 0 or yscom.size == 0:
        print(f"{com[0]}: Data missing for {com[2]} in {com[1]}")
        return None, None, None, None
    return xsref, ysref, xscom, yscom


def _integrate_area(uxref, uyref, uxcom, uycom, use_abs, norm):
    """
    Integrate the area between two interpolated curves.

    Interpolate the reference and comparison curves, find their
    common OD range, and integrate the difference.

    Parameters
    ----------
    uxref: numpy array
        Unique x values for the reference.
    uyref: numpy array
        Corresponding y values for the reference.
    uxcom: numpy array
        Unique x values for the comparison.
    uycom: numpy array
        Corresponding y values for the comparison.
    use_abs: boolean
        If True, integrate the absolute difference.
    norm: boolean
        If True, also integrate the reference curve for
        normalisation.

    Returns
    -------
    fp: float
        The normalised area between the two curves.
    nrm: float or None
        The normalised area under the reference curve, or None
        if norm is False.
    iref: scipy.interpolate.interp1d
        Interpolation function for the reference.
    icom: scipy.interpolate.interp1d
        Interpolation function for the comparison.
    uxi: float
        The lower bound of the common OD range.
    uxf: float
        The upper bound of the common OD range.
    """
    iref = interp1d(uxref, uyref, fill_value="extrapolate", kind="slinear")
    icom = interp1d(uxcom, uycom, fill_value="extrapolate", kind="slinear")
    # find common range of x
    uxi = max(uxref[0], uxcom[0])
    uxf = min(uxref[-1], uxcom[-1])
    # integrate area between curves
    if use_abs:
        fp = integrate.quad(
            lambda x: np.abs(iref(x) - icom(x)),
            uxi,
            uxf,
            limit=100,
            full_output=1,
        )[0] / (uxf - uxi)
    else:
        fp = integrate.quad(
            lambda x: iref(x) - icom(x),
            uxi,
            uxf,
            limit=100,
            full_output=1,
        )[0] / (uxf - uxi)
    nrm_val = None
    if norm:
        nrm_val = integrate.quad(iref, uxi, uxf, limit=100, full_output=1)[
            0
        ] / (uxf - uxi)
    return fp, nrm_val, iref, icom, uxi, uxf


def _plot_fitness_penalty(
    uxref,
    uyref,
    uxcom,
    uycom,
    iref,
    icom,
    uxi,
    uxf,
    ref,
    com,
    yvar,
):
    """
    Plot the area between two fitness curves.

    Parameters
    ----------
    uxref: numpy array
        Unique x values for the reference.
    uyref: numpy array
        Corresponding y values for the reference.
    uxcom: numpy array
        Unique x values for the comparison.
    uycom: numpy array
        Corresponding y values for the comparison.
    iref: scipy.interpolate.interp1d
        Interpolation function for the reference.
    icom: scipy.interpolate.interp1d
        Interpolation function for the comparison.
    uxi: float
        The lower bound of the common OD range.
    uxf: float
        The upper bound of the common OD range.
    ref: list of strings
        Experiment, condition, and strain for the reference.
    com: list of strings
        Experiment, condition, and strain for the comparison.
    yvar: string
        The variable being compared.
    """
    plt.figure()
    plt.plot(uxref, uyref, "k-", uxcom, uycom, "b-")
    x = np.linspace(uxi, uxf, max(len(uxref), len(uxcom)))
    plt.fill_between(x, iref(x), icom(x), facecolor="magenta", alpha=0.1)
    plt.xlabel("OD")
    plt.ylabel(yvar)
    plt.legend(
        [
            f"{ref[0]}: {ref[2]} in {ref[1]}",
            f"{com[0]}: {com[2]} in {com[1]}",
        ],
        loc="upper left",
        bbox_to_anchor=(-0.05, 1.15),
    )
    plt.show(block=False)


def get_fitness_penalty(
    self,
    ref,
    com,
    yvar="gr",
    figs=True,
    no_samples=100,
    use_abs=False,
    norm=False,
):
    """
    Estimate the difference in fitness between two strains.

    Calculate the area between typically two growth rate versus OD
    curves, normalised by the length along the OD-axis where they
    overlap.

    Parameters
    -----------
    self: omniplate.platereader instance
        An instance of platereader with data loaded.
    ref: list of strings
        For only a single experiment, a list of two strings. The
        first string specifies the condition and the second
        specifies the strain to be used for the reference to which
        fitness is to be calculated.
        With multiple experiments, a list of three strings. The
        first string specifies the experiment, the second specifies
        the condition, and the third specifies the strain.
    com: list of strings
        For only a single experiment, a list of two strings. The
        first string specifies the condition and the second
        specifies the strain to be compared with the reference.
        With multiple experiments, a list of three strings. The
        first string specifies the experiment, the second specifies
        the condition, and the third specifies the strain.
    yvar: string, optional
        The variable to be compared.
    figs: boolean, optional
        If True, a plot of the area between the two growth rate
        versus OD curves is shown.
    no_samples: integer
        The number of bootstraps used to estimate the error.
    use_abs: boolean
        If True, integrate the absolute difference between the
        two curves.
    norm: boolean
        If True, return the mean and variance of the area under
        the reference strain for normalisation.

    Returns
    -------
    fp: float
        The area between the two curves.
    err: float
        An estimate of the error in the calculated error, found
        by bootstrapping.
    reffp: float, optional
        The area beneath the reference strain.
    referr: float, optional
        An estimate of the error in the calculated area for the
        reference strain.

    Example
    -------
    >>> from omniplate import PlateReader
    >>> p = PlateReader("ExampleData.xlsx",
            "ExampleDataContents.xlsx", datadir="data")
    >>> p.get_stats(strains="WT",
            conditions=["2% Mal", "2% Raf"])
    >>> fp, err = p.get_fitness_penalty(
            ["2% Mal", "WT"], ["2% Raf", "WT"])
    """
    if len(ref) == 2 and len(com) == 2:
        ref.insert(0, self.all_experiments[0])
        com.insert(0, self.all_experiments[0])
    xsref, ysref, xscom, yscom = _get_fitness_samples(
        self, ref, com, yvar, no_samples
    )
    if xsref is None:
        return np.nan, np.nan
    fps = np.zeros(xsref.shape[1])
    nrm = np.zeros(xsref.shape[1])
    for j, (xref, yref, xcom, ycom) in enumerate(
        zip(xsref.T, ysref.T, xscom.T, yscom.T)
    ):
        uxref, uyref = _deduplicate_xy(xref, yref)
        uxcom, uycom = _deduplicate_xy(xcom, ycom)
        fp, nrm_val, iref, icom, uxi, uxf = _integrate_area(
            uxref, uyref, uxcom, uycom, use_abs, norm
        )
        fps[j] = fp
        if norm:
            nrm[j] = nrm_val
        if figs and j == 0:
            _plot_fitness_penalty(
                uxref,
                uyref,
                uxcom,
                uycom,
                iref,
                icom,
                uxi,
                uxf,
                ref,
                com,
                yvar,
            )
    if norm:
        return (
            np.median(fps),
            stats_err(fps),
            np.median(nrm),
            stats_err(nrm),
        )
    return np.median(fps), stats_err(fps)
