# functions for correcting for non-linearity in the OD, for
# the fluorescence of the media, and for autofluorescence
import importlib.resources as import_files
import re

import gaussianprocessderivatives as gp
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from nunchaku import Nunchaku
from scipy.interpolate import interp1d
from scipy.optimize import minimize_scalar
from statsmodels.nonparametric.smoothers_lowess import lowess

import omniplate.admin as admin
import omniplate.omerrors as errors
import omniplate.omgenutils as gu
import omniplate.sunder as sunder
from omniplate.runfitderiv import runfitderiv
import omniplate.clogger as clogger


@clogger.log
def correctOD(
    self,
    figs=True,
    bd=None,
    gp_results=False,
    ODfname=None,
    odmatch_min=0.1,
    experiments="all",
    experimentincludes=False,
    experimentexcludes=False,
    conditions="all",
    conditionincludes=False,
    conditionexcludes=False,
):
    """
    Correct for the non-linear relationship between OD and cell number.

    Requires a set of dilution data set, with the default being haploid
    yeast growing in glucose.

    An alternative can be loaded from a file - a txt file of two columns
    with OD specified in the first column and the dilution factor specified
    in the second.

    Parameters
    ---------
    figs: boolean, optional
        If True, a plot of the fit to the dilution data is produced.
    bd: dictionary, optional
        Specifies the limits on the hyperparameters for the Gaussian
        process.
        For example, bd= {0: [-1, 4], 2: [2, 6]})
        sets confines the first hyperparameter to be between 1e-1 and 1e^4
        and confines the third hyperparameter between 1e2 and 1e6.
    gp_results: boolean, optional
        If True, show the results of fitting the Gaussian process
    ODfname: string, optional
        The name of the file with the dilution data used to correct OD for
        its non-linear dependence on numbers of cells. If unspecified, data
        for haploid budding yeast growing in glucose is used.
    odmatch_min: float, optional
        An expected minimal value of the OD up to which there is a linear
        scaling of the OD with cell numbers.
    experiments: string or list of strings
        The experiments to include.
    conditions: string or list of strings
        The conditions to include.
    experimentincludes: string, optional
        Selects only experiments that include the specified string in their
        name.
    experimentexcludes: string, optional
        Ignores experiments that include the specified string in their
        name.
    conditionincludes: string, optional
        Selects only conditions that include the specified string in their
        name.
    conditionexcludes: string, optional
        Ignores conditions that include the specified string in their name.

    Examples
    -------
    >>> p.correctOD()
    >>> p.correctOD(figs= False)
    """
    exps = sunder.getset(
        self,
        experiments,
        experimentincludes,
        experimentexcludes,
        "experiment",
        nonull=True,
    )
    cons = sunder.getset(
        self,
        conditions,
        conditionincludes,
        conditionexcludes,
        "condition",
        nonull=True,
        nomedia=False,
    )
    # fit dilution data
    gc, odmatch = findODcorrection(
        self.wdirpath, ODfname, figs, bd, gp_results, odmatch_min
    )
    # correct ODs
    for exp in exps:
        for c in cons:
            if self.sc[(self.sc.experiment == exp) & (self.sc.condition == c)][
                "OD_corrected"
            ].any():
                print(f"{exp}: OD is already corrected for {c}.")
            else:
                # correct all wells
                r_data = self.r.query(
                    "experiment == @exp and condition == @c"
                )["OD"].to_numpy()
                gc.batchpredict(r_data)
                # leave small ODs unchanged
                new_r = gc.f
                new_r[r_data < odmatch] = r_data[r_data < odmatch]
                # update r data frame
                self.r.loc[
                    (self.r.experiment == exp) & (self.r.condition == c),
                    "OD",
                ] = new_r
                # flag corrections in summary stats dataframe
                self.sc.loc[
                    (self.sc.experiment == exp) & (self.sc.condition == c),
                    "OD_corrected",
                ] = True
    # update s dataframe
    admin.update_s(self)


@clogger.log
def correctmedia(
    self,
    datatypes="all",
    commonmedia=False,
    experiments="all",
    experimentincludes=False,
    experimentexcludes=False,
    conditions="all",
    conditionincludes=False,
    conditionexcludes=False,
    figs=False,
    log=True,
    frac=0.33,
):
    """
    Correct OD or fluorescence for that of the media.

    Data from wells marked Null is used.

    Uses lowess to smooth measurements of from all Null wells and subtracts
    this smoothed time series from the raw data.

    Parameters
    ----------
    datatypes: string or list of strings
        Data types to be corrected.
    commonmedia: string
        A condition containing Null wells that should be used to correct
        media for other conditions.
    experiments: string or list of strings
        The experiments to include.
    conditions: string or list of strings
        The conditions to include.
    experimentincludes: string, optional
        Selects only experiments that include the specified string in their
        name.
    experimentexcludes: string, optional
        Ignores experiments that include the specified string in their
        name.
    conditionincludes: string, optional
        Selects only conditions that include the specified string in their
        name.
    conditionexcludes: string, optional
        Ignores conditions that include the specified string in their name.
    figs: boolean, optional
        If True, display fits to data for the Null wells.
    frac: float
        The fraction of the data used for smoothing via lowess.
        https://www.statsmodels.org/devel/generated/statsmodels.nonparametric.smoothers_lowess.lowess.html

    Examples
    --------
    >>> p.correctmedia()
    >>> p.correctmedia('OD')
    >>> p.correctmedia(commonmedia= '1% Glu')
    """
    exps = sunder.getset(
        self,
        experiments,
        experimentincludes,
        experimentexcludes,
        "experiment",
        nonull=True,
    )
    cons = sunder.getset(
        self,
        conditions,
        conditionincludes,
        conditionexcludes,
        "condition",
        nonull=True,
        nomedia=False,
    )
    for exp in exps:
        # data types
        expdatatypes = (
            self.datatypes[exp]
            if datatypes == "all"
            else gu.makelist(datatypes)
        )
        # correct for media
        for dtype in expdatatypes:
            for c in cons:
                if self.sc[
                    (self.sc.experiment == exp) & (self.sc.condition == c)
                ][dtype + "_corrected_for_media"].any():
                    print(
                        f"{exp}: {dtype} is already corrected for media"
                        f" in {c}."
                    )
                else:
                    if c in self.allconditions[exp]:
                        print(f"{exp}: Correcting {dtype} for media in {c}.")
                        cm = commonmedia if commonmedia else c
                        # update r dataframe
                        (
                            success,
                            negvalues,
                        ) = performmediacorrection(
                            self.r, dtype, exp, c, figs, cm, frac
                        )
                        if success:
                            self.sc.loc[
                                (self.sc.experiment == exp)
                                & (self.sc.condition == c),
                                dtype + "_corrected_for_media",
                            ] = True
                            if negvalues:
                                if not self.progress["negativevalues"][exp]:
                                    self.progress["negativevalues"][
                                        exp
                                    ] = negvalues
                                else:
                                    self.progress["negativevalues"][
                                        exp
                                    ] += negvalues
        if self.progress["negativevalues"][exp]:
            print(
                "\nWarning: correcting media has created negative "
                f"values in {exp} for"
            )
            print(self.progress["negativevalues"][exp])
    # update s dataframe
    admin.update_s(self)


@clogger.log
def correctauto(
    self,
    f=["GFP", "AutoFL"],
    refstrain="WT",
    figs=True,
    useGPs=True,
    flcvfn="matern",
    bd=None,
    nosamples=1000,
    experiments="all",
    experimentincludes=False,
    experimentexcludes=False,
    conditions="all",
    conditionincludes=False,
    conditionexcludes=False,
    strains="all",
    strainincludes=False,
    strainexcludes=False,
    **kwargs,
):
    """
    Correct fluorescence for autofluorescence.

    The correction is made using the fluorescence of an untagged
    reference strain.

    The reference strain is used to estimate the autofluorescence via
    either the method of Lichten et al., 2014, where measurements of
    fluorescence at two wavelengths is required, or by using the
    fluorescence of the reference strain interpolated to the OD of the
    strain of interest (Berthoumieux et al., 2013).

    Using two measurements of fluorescence is thought to be more accurate,
    particularly for low fluorescence measurements (Mihalcescu et al.,
    2015).

    Arguments
    --
    f: string or list of strings
        The fluorescence measurements, typically either ['mCherry'] or
        ['GFP', 'AutoFL'].
    refstrain: string
        The reference strain.
    figs: boolean
        If True, display plots showing the fits to the reference strain's
        fluorescence.
    useGPs: boolean
        If True, use Gaussian processes to generate extra samples from
        the replicates. Recommended, particularly if there are only a few
        replicates, but slower.
    flcvfn: str, optional
        The covariance function to use for the Gaussian process applied
        to the logarithm of the fluorescence.
    bd: dict, optional
        Specifies the bounds on the hyperparameters for the Gaussian
        process applied to the logarithm of the fluorescence,
        e.g. {2: (-2, 0)}.
    experiments: string or list of strings
        The experiments to include.
    conditions: string or list of strings
        The conditions to include.
    strains: string or list of strings
        The strains to include.
    experimentincludes: string, optional
        Selects only experiments that include the specified string in
        their name.
    experimentexcludes: string, optional
        Ignores experiments that include the specified string in their
        name.
    conditionincludes: string, optional
        Selects only conditions that include the specified string in their
        name.
    conditionexcludes: string, optional
        Ignores conditions that include the specified string in their name.
    strainincludes: string, optional
        Selects only strains that include the specified string in their
        name.
    strainexcludes: string, optional
        Ignores strains that include the specified string in their name.
    **kwargs: passed to fitderiv when fitting the fluorescence.

    Notes
    -----
    In principle

    >>> p.correctmedia()

    should be run before running correctauto when processing data with two
    fluorescence measurements.

    It is unnecessary with only one fluorescence measurement because the
    normalisation is then done directly with the reference strain's
    fluorescence and this fluorescence can include the fluorescence from
    the media.

    In practice, running correctmedia may generate negative values of the
    fluorescence at some time points. These negative values will create
    NaNs in the corrected fluorescence, which are normally harmless.

    With sufficiently many negative values of the fluorescence, however,
    correcting data with two fluorescence measurements can become
    corrupted.

    If correctmedia generates negative fluorescence values, we therefore
    recommend comparing the corrected fluorescence between

    >>> p.correctmedia()
    >>> p.correctauto(['GFP', 'AutoFL')

    and

    >>> p.correctauto('GFP')

    to determine if these negative values are deleterious.

    Examples
    --------
    To correct data with one type of fluorescence measurement, use:

    >>> p.correctauto('GFP')
    >>> p.correctauto('mCherry', refstrain= 'BY4741')

    To correct data with two types of fluorescence measurement, use:

    >>> p.correctauto(['GFP', 'AutoFL'])
    >>> p.correctauto(['GFP', 'AutoFL'], refstrain= 'wild-type')

    References
    ----------
    S Berthoumieux, H De Jong, G Baptist, C Pinel, C Ranquet, D Ropers,
    J Geiselmann (2013).
    Shared control of gene expression in bacteria by transcription factors
    and global physiology of the cell.
    Mol Syst Biol, 9, 634.

    CA Lichten, R White, IB Clark, PS Swain (2014).
    Unmixing of fluorescence spectra to resolve quantitative time-series
    measurements of gene expression in plate readers.
    BMC Biotech, 14, 1-11.

    I Mihalcescu, MVM Gateau, B Chelli, C Pinel, JL Ravanat (2015).
    Green autofluorescence, a double edged monitoring tool for bacterial
    growth and activity in micro-plates.
    Phys Biol, 12, 066016.

    """
    admin.check_kwargs(kwargs)
    f = gu.makelist(f)
    exps, cons, _ = sunder.getall(
        self,
        experiments,
        experimentincludes,
        experimentexcludes,
        conditions,
        conditionincludes,
        conditionexcludes,
        strains,
        strainincludes,
        strainexcludes,
        nonull=True,
    )
    # check for negative fluorescence values
    for e in exps:
        for c in cons:
            if self.progress["negativevalues"][e]:
                for datatype in f:
                    if (
                        datatype in self.progress["negativevalues"][e]
                        and c in self.progress["negativevalues"][e]
                    ):
                        print(
                            f"{e}: The negative values for {datatype}"
                            f" in {c} will generate NaNs."
                        )
    # going ahead
    print(f"Using {refstrain} as the reference.")
    # correct for autofluorescence
    if len(f) == 2:
        correctauto2(
            self,
            f,
            refstrain,
            figs,
            useGPs,
            flcvfn,
            bd,
            nosamples,
            experiments,
            experimentincludes,
            experimentexcludes,
            conditions,
            conditionincludes,
            conditionexcludes,
            strains,
            strainincludes,
            strainexcludes,
            **kwargs,
        )
    elif len(f) == 1:
        correctauto1(
            self,
            f,
            refstrain,
            figs,
            useGPs,
            flcvfn,
            bd,
            nosamples,
            experiments,
            experimentincludes,
            experimentexcludes,
            conditions,
            conditionincludes,
            conditionexcludes,
            strains,
            strainincludes,
            strainexcludes,
            **kwargs,
        )
    else:
        print(f"f = {f} must be a list of length 1 or 2.")


def corrections(
    self,
    experiments="all",
    conditions="all",
    strains="all",
    experimentincludes=False,
    experimentexcludes=False,
    conditionincludes=False,
    conditionexcludes=False,
    strainincludes=False,
    strainexcludes=False,
):
    """
    Display the current corrections to the data.

    Parameters
    ----------
    experiments: string or list of strings
        The experiments to include.
    conditions: string or list of strings
        The conditions to include.
    strains: string or list of strings
        The strains to include.
    experimentincludes: string, optional
        Selects only experiments that include the specified string in their
        name.
    experimentexcludes: string, optional
        Ignores experiments that include the specified string in their
        name.
    conditionincludes: string, optional
        Selects only conditions that include the specified string in their
        name.
    conditionexcludes: string, optional
        Ignores conditions that include the specified string in their name.
    strainincludes: string, optional
        Selects only strains that include the specified string in their
        name.
    strainexcludes: string, optional
        Ignores strains that include the specified string in their name.

    Returns
    -------
    df: dataframe
        Contains the status of the corrections for the specified strains,
        conditions, and experiments.

    Examples
    --------
    >>> p.corrections()
    >>> p.corrections(strainincludes= 'GAL')
    """
    exps, cons, strs = sunder.getall(
        self,
        experiments,
        experimentincludes,
        experimentexcludes,
        conditions,
        conditionincludes,
        conditionexcludes,
        strains,
        strainincludes,
        strainexcludes,
        nonull=True,
    )
    df = self.sc.query(
        "experiment == @exps and condition == @cons and strain == @strs"
    )
    # only show corrections and not stats
    df = df[
        ["experiment", "strain", "condition"]
        + [col for col in df.columns if "correct" in col]
    ]
    df = df.T
    return df


def read_dilution_data(fname):
    d = import_files.read_text("omniplate", fname)
    res = np.array(re.split(r"\n|\t", d)[:-1]).astype(float)
    od, dilfac = res[::2], res[1::2]
    if fname == "dilution_data_xiao.tsv":
        # missing replicate - use mean of existing ones
        dilfac = np.insert(dilfac, 0, dilfac[0])
        od = np.insert(od, 0, np.mean(od[:2]))
    else:
        raise SystemExit("Dilution data unrecognised.")
    return od, dilfac


def arrange_into_replicates(od, dilfac):
    """Rearrange so that data from each replicate is in a column."""
    udilfac, indices, counts = np.unique(
        dilfac, return_inverse=True, return_counts=True
    )
    ucounts = np.unique(counts)
    if len(ucounts) == 1:
        noreps = np.unique(counts)[0]
        dilfac_reps = np.tile(np.atleast_2d(udilfac).T, noreps)
        od_reps = np.array(
            [od[indices == i] for i in range(udilfac.size)]
        ).reshape((udilfac.size, noreps))
        return od_reps, dilfac_reps
    else:
        raise Exception(
            "There are inconsistent numbers of replicates"
            " in the OD correction data."
        )


def findODcorrection(wdirpath, ODfname, figs, bd, gp_results, odmatch_min):
    """
    Determine a function to correct OD.

    Use a Gaussian process to fit serial dilution data to correct for
    non-linearities in the relationship between OD and cell density.

    The data are either loaded from file ODfname or the default
    data for haploid yeast growing in glucose are used.
    """
    print("Fitting dilution data for OD correction for non-linearities.")
    if ODfname is not None:
        try:
            od_df = pd.read_csv(
                str(wdirpath / ODfname), sep=None, engine="python", header=None
            )
            print(f"Using {ODfname}")
            od_data = od_df.to_numpy()
            od, dilfac = od_data[:, 0], od_data[:, 1]
        except (FileNotFoundError, OSError):
            raise errors.FileNotFound(str(wdirpath / ODfname))
    else:
        print("Using default data.")
        fname = "dilution_data_xiao.tsv"
        od, dilfac = read_dilution_data(fname)
    od, dilfac = arrange_into_replicates(od, dilfac)
    # run nunchaku
    X = np.mean(dilfac, 1)
    nc = Nunchaku(X, od.T, estimate_err=True, prior=[-5, 5])
    num_regions, _ = nc.get_number(4)
    bds, _ = nc.get_iboundaries(num_regions)
    # find linear region, which starts from origin
    odmatch_pts = np.mean(od, 1)[bds]
    # pick end point with OD at least equal to odmatch_min
    ipick = np.where(odmatch_pts > odmatch_min)[0][0]
    odmatch = odmatch_pts[ipick]
    dilfacmatch = X[bds[ipick]]
    # process data
    dilfac = dilfac.flatten()[np.argsort(od.flatten())]
    od = np.sort(od.flatten())
    # rescale so that OD and dilfac match
    y = dilfac / dilfacmatch * odmatch
    # set up Gaussian process
    bds = {0: (-4, 4), 1: (-4, 4), 2: (-3, -1)}
    # find bounds
    if bd is not None:
        bds = gu.mergedicts(original=bds, update=bd)
    gc = gp.maternGP(bds, od, y)
    # run Gaussian process
    gc.findhyperparameters(noruns=5, exitearly=True, quiet=True)
    if gp_results:
        gc.results()
    gc.predict(od)
    if figs:
        plt.figure()
        gc.sketch(".")
        plt.plot(odmatch, odmatch, "bs")
        plt.grid(True)
        plt.xlabel("OD")
        plt.ylabel("corrected OD (relative cell numbers)")
        if ODfname:
            plt.title("Fitting " + ODfname)
        else:
            plt.title("for haploid budding yeast in glucose")
        plt.show(block=False)
    return gc, odmatch


def performmediacorrection(
    r_df, dtype, exp, condition, figs, commonmedia, frac
):
    """
    Correct data of type dtype for any signal from the media.

    Use lowess to smooth over time the media data from the Null
    wells and subtract the smoothed values from the data.
    """
    # find data for correction with condition equal to commonmedia
    df = r_df.query(
        "experiment == @exp and condition == @commonmedia"
        " and strain == 'Null'"
    )
    if df.empty:
        # no data
        print(
            f' No well annotated "Null" was found for {commonmedia} '
            f"in experiment {exp}."
        )
        print(
            f"-> Correcting for media for {dtype} in {commonmedia} abandoned!"
        )
        return False, None
    else:
        # there is data - change r dataframe
        rtest = (r_df.experiment == exp) & (r_df.condition == condition)
        t, data = df["time"].to_numpy(), df[dtype].to_numpy()
        if ~np.any(data[~np.isnan(data)]):
            # all data is NaN
            return False, None
        # find correction
        res = lowess(data, t, frac=frac)
        correctionfn = interp1d(
            res[:, 0],
            res[:, 1],
            fill_value=(res[0, 1], res[-1, 1]),
            bounds_error=False,
        )
        if figs:
            plt.figure()
            plt.plot(t, data, "ro", res[:, 0], res[:, 1], "b-")
            plt.xlabel("time (hours)")
            plt.title(
                exp + ": media correction for " + dtype + " in " + condition
            )
            plt.show(block=False)
        # perform correction
        r_df.loc[rtest, dtype] = r_df[rtest][dtype] - correctionfn(
            r_df[rtest]["time"]
        )
        # check for any negative values
        negvalues = ""
        for s in np.unique(r_df[rtest]["strain"][r_df[rtest][dtype] < 0]):
            if s != "Null":
                wstr = f"\t{dtype}: {s} in {condition} for wells "
                for well in np.unique(
                    r_df[rtest][r_df[rtest].strain == s]["well"][
                        r_df[rtest][dtype] < 0
                    ]
                ):
                    wstr += f"{well}, "
                wstr = wstr[:-2] + "\n"
                negvalues += wstr
        return True, negvalues


def correctauto1(
    self,
    f,
    refstrain,
    figs,
    useGPs,
    flcvfn,
    bd,
    nosamples,
    experiments,
    experimentincludes,
    experimentexcludes,
    conditions,
    conditionincludes,
    conditionexcludes,
    strains,
    strainincludes,
    strainexcludes,
    **kwargs,
):
    """
    Correct autofluorescence for measurements with emissions at one wavelength.

    Corrects for autofluorescence for data with emissions measured at one
    wavelength using the fluorescence of the reference strain
    interpolated to the OD of the tagged strain.

    This method in principle corrects too for the fluorescence of the medium,
    although running correctmedia is still recommended.
    """
    print("Using one fluorescence wavelength.")
    print(f"Correcting autofluorescence using {f[0]}.")
    for e in sunder.getset(
        self,
        experiments,
        experimentincludes,
        experimentexcludes,
        "experiment",
        nonull=True,
    ):
        for c in sunder.getset(
            self,
            conditions,
            conditionincludes,
            conditionexcludes,
            "condition",
            nonull=True,
            nomedia=True,
        ):
            # process reference strain
            if c in self.allconditions[e]:
                refstrfn = processref1(self, f, refstrain, figs, e, c)
            else:
                refstrfn = None
            if refstrfn is None:
                continue
            # correct strains
            for s in sunder.getset(
                self,
                strains,
                strainincludes,
                strainexcludes,
                "strain",
                nonull=True,
            ):
                if f"{s} in {c}" in self.allstrainsconditions[e]:
                    t, (od, rawfl) = sunder.extractwells(
                        self.r, self.s, e, c, s, ["OD", f[0]]
                    )
                    # no data
                    if od.size == 0 or rawfl.size == 0:
                        print(f"\n-> No data found for {e}: {s} in {c}.\n")
                        continue
                    # correct autofluorescence for each replicate
                    fl = np.transpose(
                        [
                            rawfl[:, i] - refstrfn(od[:, i])
                            for i in range(od.shape[1])
                        ]
                    )
                    fl[fl < 0] = np.nan
                    if useGPs:
                        # get time
                        flperod = samplewithGPs(
                            self,
                            f,
                            t,
                            fl,
                            od,
                            flcvfn,
                            bd,
                            nosamples,
                            e,
                            c,
                            s,
                            figs,
                            **kwargs,
                        )
                    else:
                        # use only the replicates
                        flperod = np.transpose(
                            [fl[:, i] / od[:, i] for i in range(od.shape[1])]
                        )
                        flperod[flperod < 0] = np.nan
                    # check number of NaN
                    nonans = np.count_nonzero(np.isnan(fl))
                    if np.any(nonans):
                        if nonans == fl.size:
                            print(
                                f"\n-> Corrected fluorescence is all NaN for {e}: {s} in {c}.\n"
                            )
                        else:
                            print(
                                f"Warning - {e}: {s} in {c}\n"
                                f"{nonans} corrected data points are"
                                " NaN because the corrected fluorescence"
                                " was negative.",
                            )
                    # store results
                    bname = "c" + f[0]
                    autofdict = {
                        "experiment": e,
                        "condition": c,
                        "strain": s,
                        "time": t,
                        bname: np.nanmean(fl, 1),
                        bname + "_err": nanstdzeros2nan(fl, 1),
                        bname + "perOD": np.nanmean(flperod, 1),
                        bname + "perOD_err": nanstdzeros2nan(flperod, 1),
                    }
                    addtodataframes(self, autofdict, bname)
                    addrecord(self, f[0], e, c, s)


def processref1(self, f, refstrain, figs, experiment, condition):
    """
    Process reference strain for data with one fluorescence measurement.

    Use lowess to smooth the fluorescence of the reference
    strain as a function of OD.

    Parameters
    ----------
    f: string
        The fluorescence to be corrected. For example, ['mCherry'].
    refstrain: string
        The reference strain. For example, 'WT'.
    figs: boolean
        If True, display fits of the reference strain's fluorescence.
    experiment: string
        The experiment to be corrected.
    condition: string
        The condition to be corrected.

    Returns
    -------
    refstrfn: function
        The reference strain's fluorescence as a function of OD.
    """
    e, c = experiment, condition
    print(f"\n{e}: Processing reference strain {refstrain} for {f[0]} in {c}.")
    _, (od, fl) = sunder.extractwells(
        self.r, self.s, e, c, refstrain, ["OD", f[0]]
    )
    if od.size == 0 or fl.size == 0:
        raise errors.CorrectAuto(f"{e}: {refstrain} not found in {c}.")
    else:
        odf = od.flatten("F")
        flf = fl.flatten("F")
        if ~np.any(flf[~np.isnan(flf)]):
            return None
        # smooth fluorescence as a function of OD using lowess to minimize
        # refstrain's autofluorescence

        def choosefrac(frac):
            res = lowess(flf, odf, frac=frac)
            refstrfn = interp1d(
                res[:, 0],
                res[:, 1],
                fill_value=(res[0, 1], res[-1, 1]),
                bounds_error=False,
            )
            # max gives smoother fits than mean
            return np.max(np.abs(flf - refstrfn(odf)))

        res = minimize_scalar(choosefrac, bounds=(0.1, 0.99), method="bounded")
        # choose the optimum frac
        frac = res.x if res.success else 0.33
        res = lowess(flf, odf, frac=frac)
        refstrfn = interp1d(
            res[:, 0],
            res[:, 1],
            fill_value=(res[0, 1], res[-1, 1]),
            bounds_error=False,
        )
        if figs:
            # plot fit
            plt.figure()
            plt.plot(odf, flf, ".", alpha=0.5)
            plt.plot(res[:, 0], res[:, 1])
            plt.xlabel("OD")
            plt.ylabel(f[0])
            plt.title(e + ": " + refstrain + " for " + c)
            plt.show(block=False)
        return refstrfn


def correctauto2(
    self,
    f,
    refstrain,
    figs,
    useGPs,
    flcvfn,
    bd,
    nosamples,
    experiments,
    experimentincludes,
    experimentexcludes,
    conditions,
    conditionincludes,
    conditionexcludes,
    strains,
    strainincludes,
    strainexcludes,
    **kwargs,
):
    """
    Correct autofluorescence for measurements with two emission wavelengths.

    Corrects for autofluorescence using spectral unmixing for data with
    measured emissions at two wavelengths.

    References
    ----------
    CA Lichten, R White, IB Clark, PS Swain (2014). Unmixing of fluorescence
    spectra to resolve quantitative time-series measurements of gene
    expression in plate readers.
    BMC Biotech, 14, 1-11.
    """
    # correct for autofluorescence
    print("Using two fluorescence wavelengths.")
    print(f"Correcting autofluorescence using {f[0]} and {f[1]}.")
    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,
        ):
            # process reference strain
            refqrfn = processref2(self, f, refstrain, figs, e, c)
            if refqrfn is None:
                # too many NaNs in reference strain
                continue
            # process other 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]
                ):
                    t, (fl_0, fl_1, od) = sunder.extractwells(
                        self.r, self.s, e, c, s, f.copy() + ["OD"]
                    )
                    if fl_0.size == 0 or fl_1.size == 0:
                        print(f"Warning: No data found for {e}: {s} in {c} !!")
                        continue
                    nodata, nr = fl_0.shape
                    # set negative values to NaNs
                    fl_0[fl_0 < 0] = np.nan
                    fl_1[fl_1 < 0] = np.nan
                    # use mean OD for correction
                    odmean = np.nanmean(od, axis=1)
                    # correct autofluorescence
                    ra = refqrfn(odmean)
                    fl = applyautoflcorrection(self, ra, fl_0, fl_1)
                    fl[fl < 0] = np.nan
                    # store corrected fluorescence
                    bname = "c" + f[0]
                    autofdict = {
                        "experiment": e,
                        "condition": c,
                        "strain": s,
                        "time": t,
                        bname: np.nanmean(fl, 1),
                        bname + " err": nanstdzeros2nan(fl, 1),
                    }
                    if useGPs:
                        flperod = samplewithGPs(
                            self,
                            f[0],
                            t,
                            fl,
                            od,
                            flcvfn,
                            bd,
                            nosamples,
                            e,
                            c,
                            s,
                            figs,
                            **kwargs,
                        )
                    else:
                        # use only the replicates
                        flperod = fl / od
                        flperod[flperod < 0] = np.nan
                    # update dict
                    autofdict[bname + "perOD_err"] = naniqrzeros2nan(
                        flperod, 1
                    )
                    autofdict[bname + "perOD"] = np.nanmean(flperod, 1)
                    # add to data frames
                    addtodataframes(self, autofdict, bname)
                    addrecord(self, f[0], e, c, s)


def processref2(self, f, refstrain, figs, experiment, condition):
    """
    Process reference strain for spectral unmixing.

    Requires data with two fluorescence measurements.

    Use lowess to smooth the ratio of emitted fluorescence measurements
    so that the reference strain's data is corrected to zero as best
    as possible.

    Parameters
    ----------
    f: list of strings
        The fluorescence measurements. For example, ['GFP', 'AutoFL'].
    refstrain: string
        The reference strain. For example, 'WT'.
    figs: boolean
        If True, display fits of the fluorescence ratios.
    experiment: string
        The experiment to be corrected.
    condition: string
        The condition to be corrected.

    Returns
    -------
    qrfn: function
        The ratio of the two fluorescence values for the reference strain
        as a function of OD.
    """
    e, c = experiment, condition
    print(f"\n{e}: Processing reference strain {refstrain} for {f[0]} in {c}.")
    # refstrain data
    t, (f0, f1, od) = sunder.extractwells(
        self.r, self.s, e, c, refstrain, f + ["OD"]
    )
    if f0.size == 0 or f1.size == 0 or od.size == 0:
        raise errors.CorrectAuto(f"{e}: {refstrain} not found in {c}.")
    else:
        f0[f0 < 0] = np.nan
        f1[f1 < 0] = np.nan
        odf = od.flatten("F")
        odrefmean = np.mean(od, 1)
        qrf = (f1 / f0).flatten("F")
        if np.all(np.isnan(qrf)):
            print(f"{e}: {refstrain} in {c} has too many NaNs.")
            return
        # smooth to minimise autofluorescence in refstrain

        def choosefrac(frac):
            res = lowess(qrf, odf, frac)
            qrfn = interp1d(
                res[:, 0],
                res[:, 1],
                fill_value=(res[0, 1], res[-1, 1]),
                bounds_error=False,
            )
            flref = applyautoflcorrection(self, qrfn(odrefmean), f0, f1)
            return np.max(np.abs(flref))

        res = minimize_scalar(choosefrac, bounds=(0.1, 0.99), method="bounded")
        # calculate the relationship between qr and OD
        frac = res.x if res.success else 0.95
        res = lowess(qrf, odf, frac)
        qrfn = interp1d(
            res[:, 0],
            res[:, 1],
            fill_value=(res[0, 1], res[-1, 1]),
            bounds_error=False,
        )
        if figs:
            plt.figure()
            plt.plot(odf, qrf, ".", alpha=0.5)
            plt.plot(res[:, 0], res[:, 1])
            plt.xlabel("OD")
            plt.ylabel(f[1] + "/" + f[0])
            plt.title(e + ": " + refstrain + " in " + c)
            plt.show(block=False)
        # check autofluorescence correction for reference strain
        flref = applyautoflcorrection(self, qrfn(odrefmean), f0, f1)
        flrefperod = flref / od
        # set negative values to NaNs
        flref[flref < 0] = np.nan
        flrefperod[flrefperod < 0] = np.nan
        # store results
        bname = "c" + f[0]
        autofdict = {
            "experiment": e,
            "condition": c,
            "strain": refstrain,
            "time": t,
            bname: np.nanmean(flref, 1),
            bname + "perOD": np.nanmean(flrefperod, 1),
            bname + "_err": nanstdzeros2nan(flref, 1),
            bname + "perOD_err": nanstdzeros2nan(flrefperod, 1),
        }
        addtodataframes(self, autofdict, bname)
        return qrfn


def applyautoflcorrection(self, ra, f0data, f1data):
    """Correct for autofluorescence returning an array of replicates."""
    nr = f0data.shape[1]
    raa = np.reshape(np.tile(ra, nr), (np.size(ra), nr), order="F")
    return (raa * f0data - f1data) / (
        raa - self._gamma * np.ones(np.shape(raa))
    )


def addtodataframes(self, autofdict, bname):
    """Added dict of autofluorescence corrections to data frames."""
    autofdf = pd.DataFrame(autofdict)
    if bname not in self.s.columns:
        # extend dataframe
        self.s = pd.merge(self.s, autofdf, how="outer")
    else:
        # update dataframe
        self.s = gu.absorbdf(
            self.s,
            autofdf,
            ["experiment", "condition", "strain", "time"],
        )


def addrecord(self, fname, e, c, s):
    """Record correction."""
    self.sc.loc[
        (self.sc.experiment == e)
        & (self.sc.condition == c)
        & (self.sc.strain == s),
        fname + "_corrected_for_autofluorescence",
    ] = True


def nanstdzeros2nan(a, axis=None):
    """Like nanstd but setting zeros to nan."""
    err = np.nanstd(a, axis)
    err[err == 0] = np.nan
    return err


def naniqrzeros2nan(a, axis=None):
    """Interquartile range but setting zeros to nan."""
    iqr = np.nanquantile(a, 0.75, axis) - np.nanquantile(a, 0.25, axis)
    iqr[iqr == 0] = np.nan
    return iqr


def gethypers(self, exp, con, s, dtype="gr"):
    """Find parameters for GP from sc data frame."""
    sdf = self.sc[
        (self.sc.experiment == exp)
        & (self.sc.condition == con)
        & (self.sc.strain == s)
    ]
    try:
        cvfn = sdf["gp_for_gr"].values[0]
        hypers = [
            sdf[col].values[0]
            for col in sorted(sdf.columns)
            if ("hyper" in col and "gr" in col)
        ]
        if np.any(np.isnan(hypers)):
            return None, None
        else:
            return hypers, cvfn
    except KeyError:
        return None, None


def instantiateGP(hypers, cvfn, x, y):
    """Instantiate a Gaussian process."""
    x1 = np.tile(x, y.shape[1])
    y1 = np.reshape(y, y.size, order="F")
    y1 = y1[np.argsort(x1)]
    x1 = np.sort(x1)
    # bounds are irrelevant because parameters are optimal
    go = getattr(gp, cvfn + "GP")({0: (-5, 5), 1: (-4, 4), 2: (-5, 2)}, x1, y1)
    go.lth_opt = hypers
    # make predictions so that samples can be generated
    go.predict(x, derivs=2)
    return go


def samplewithGPs(
    self, flname, t, fl, od, flcvfn, bd, nosamples, e, c, s, figs, **kwargs
):
    """Generate extra samples of flperod using Gaussian processes."""
    hypers, cvfn = gethypers(self, e, c, s)
    if hypers is None or cvfn is None:
        raise SystemExit(
            f"\nYou first must run getstats for {s} in {c} for {e}."
        )
    fitvar = f"log_{flname}"
    derivname = f"d/dt_{fitvar}"
    od[od < 0] = np.finfo(float).eps
    # OD replicates with no NaN
    od_reps = np.setdiff1d(
        np.arange(od.shape[1]), np.unique(np.argwhere(np.isnan(od))[:, 1])
    )
    # run GP for log fluorescence; omfitderiv deals with NaNs
    if np.any(fl[~np.isnan(fl)]) and np.any(od_reps):
        ff, _ = runfitderiv(
            self,
            t,
            fl,
            fitvar,
            derivname,
            e,
            c,
            s,
            bd=bd,
            cvfn=flcvfn,
            logs=True,
            figs=figs,
            **kwargs,
        )
        if not ff.success:
            print(f"\n-> Fitting fluorescence failed for {e}: {s} in {c}.\n")
            return np.nan * np.ones((t.size, nosamples))
        # initialise GP for log ODs
        go = instantiateGP(hypers, cvfn, t, np.log(od)[:, od_reps])
        # sample
        lod_samples = go.sample(nosamples)
        lfl_samples = ff.fitderivsample(nosamples)[0]
        flperod_samples = np.exp(lfl_samples - lod_samples)
        # insert into or replace matrix of NaNs
        flperod = np.nan * np.ones((t.size, nosamples))
        flperod = flperod_samples
    else:
        # all NaN
        flperod = np.nan * np.ones((t.size, nosamples))
    return flperod
