"""Functions to load the data and parse the annotation file."""

import re
from dataclasses import dataclass
from typing import Dict, List

import numpy as np
import pandas as pd
from tqdm import tqdm

import omniplate.admin as admin
import omniplate.clogger as clogger
import omniplate.omerrors as errors
import omniplate.omgenutils as gu
from omniplate.parseplate import parseplate


@clogger.log
def load(
    self,
    dnames,
    anames=None,
    experimenttype=None,
    platereadertype="Tecan",
    dsheets=False,
    asheets=False,
    info=True,
    combine=False,
    combinename=None,
    min_dt=0.05,
):
    """
    Load raw data files.

    Two files are required: one generated by the plate reader and the other
    the corresponding annotation files - both assumed to be xlsx.

    Parameters
    ----------
    dnames: string or list of strings, optional
        The name of the file containing the data from the plate reader
        or a list of file names.
    anames: string or list of strings, optional
        The name of file containing the corresponding annotation or a list
        of file names.
    experimenttype: string or list of strings, optional
        If specified, creates a new experiment_type column in each
        dataframe.
    platereadertype: string or list of str
        The type of plate reader, currently either 'Tecan' or 'Sunrise'
        or 'tidy'. If you have data from different plate readers, say three,
        use a list, such as ["Tecan", "Tecan", "Sunrise"].
    dsheets: list of integers or strings, optional
        The relevant sheets of the Excel files storing the data.
    asheets: list of integers or strings, optional
        The relevant sheets of the corresponding Excel files for the
        annotation.
    info: boolean
        If True (default), display summary information on the data once
        loaded.
    combine: boolean
        If True (default is False), combine all the data loaded into one
        single experiment.
    combinename: str (optional)
        If specified, the name given to the combined experiment.
    min_dt: float, optional
        The minimum time interval taken by the plate reader between
        measurements. Used only when combine=True.

    Examples
    -------
    >>> p.load('Data.xlsx', 'DataContents.xlsx')
    >>> p.load('Data.xlsx', 'DataContents.xlsx', info= False)
    >>> p.load(['HxtA.xls', 'HxtB.xlsx'], ['HxtAcontents.xlsx',
            'HxtBcontents.xlsx'], combine=True)
    >>> p.load(['HxtG.xls', 'HxtP.xlsx'],
            ['HxtGcontents.xlsx', 'HxtPcontents.xlsx'],
            ["glucose_pregrowth", "pyruvate_pregrowth"])
    """
    (
        dnames,
        anames,
        platereadertype,
        experimenttype_dict,
        dsheets,
        asheets,
    ) = run_load_checks(
        dnames=dnames,
        anames=anames,
        experimenttype=experimenttype,
        platereadertype=platereadertype,
        dsheets=dsheets,
        asheets=asheets,
    )
    alldata = {}
    for i, (dname, platereader) in enumerate(zip(dnames, platereadertype)):
        # get dataframe for raw data
        res = loaddatafiles(
            platereadertype=platereader,
            datadirpath=self.datadirpath,
            dname=dname,
            dsheet=dsheets[i],
            aname=anames[i],
            asheet=asheets[i],
        )
        self.allexperiments.append(res.experiment)
        self.allconditions[res.experiment] = res.allconditions
        self.allstrains[res.experiment] = res.allstrains
        self.datatypes[res.experiment] = res.datatypes
        self.allstrainsconditions[res.experiment] = list(
            (res.rdf.strain + " in " + res.rdf.condition).dropna().unique()
        )
        alldata.update(res.alldata)
        # make r dataframe of raw data
        self.r = (
            pd.merge(self.r, res.rdf, how="outer")
            if hasattr(self, "r")
            else res.rdf
        )
        # update progress dictionary
        admin.initialiseprogress(self, res.experiment)
    sort_attributes(self)
    # make sc dataframe for summary stats and corrections
    alldfs = []
    for exp in alldata:
        strs, cons = [], []
        for cs in alldata[exp]:
            strs.append(cs.split(" in ")[0])
            cons.append(cs.split(" in ")[1])
        corrdict = {"experiment": exp, "strain": strs, "condition": cons}
        corrdict.update(
            {dtype + "_measured": True for dtype in self.datatypes[exp]}
        )
        alldfs.append(pd.DataFrame(corrdict))
    self.sc = pd.concat(alldfs)
    self.sc = self.sc.reset_index(drop=True)
    # replace NaN with False for datatypes_measured
    for column in self.sc.columns[self.sc.columns.str.contains("_measured")]:
        self.sc[column] = (
            self.sc[column].fillna(False).infer_objects(copy=False)
        )
    # make backup dataframe of original data
    self.origr = self.r.copy()
    # make dataframe for well contents
    self.wellsdf = admin.makewellsdf(self.r)
    # make s dataframe for summary data
    self.s = admin.make_s(self)
    # add experiment_type
    if experimenttype is not None:
        for df in [self.r, self.s, self.sc]:
            df["experiment_type"] = df.experiment.map(experimenttype_dict)
    # display info on experiment, conditions and strains
    if info:
        self.info
    print('\nWarning: wells with no strains have been changed to "Null".')
    # combine into one experiment
    if combine:
        if combinename is not None and isinstance(combinename, str):
            self.combined = combinename
        combine_experiments(self, min_dt=min_dt)


def run_load_checks(
    dnames,
    anames,
    experimenttype,
    platereadertype,
    dsheets,
    asheets,
):
    """Fix inputs to loaddata."""
    dnames = gu.makelist(dnames)
    if anames is None:
        anames = [dname.split(".")[0] + "Contents.xlsx" for dname in dnames]
    else:
        anames = gu.makelist(anames)
    if experimenttype is not None:
        experimenttype = gu.makelist(experimenttype)
        experimenttype_dict = {
            dname.split(".")[0]: e_type
            for dname, e_type in zip(dnames, experimenttype)
        }
    else:
        experimenttype_dict = None
    if isinstance(platereadertype, str):
        platereadertype = [platereadertype for dname in dnames]
    if not dsheets:
        dsheets = [0 for dname in dnames]
    if not asheets:
        asheets = [0 for dname in dnames]
    return (
        dnames,
        anames,
        platereadertype,
        experimenttype_dict,
        dsheets,
        asheets,
    )


def loaddatafiles(platereadertype, datadirpath, dname, dsheet, aname, asheet):
    """Call functions to parse data and metadate from the input files."""

    @dataclass
    class res:
        rdf: pd.DataFrame
        allconditions: List[str]
        allstrains: List[str]
        alldata: Dict
        experiment: str
        datatypes: List[str]

    experiment = dname.split(".")[0]
    # import and parse plate contents file
    allconditions, allstrains, alldata, wellcontents = analyseContentsofWells(
        datadirpath, experiment, aname, asheet
    )
    # import and parse data file created by plate reader
    try:
        print("Loading", dname)
        rdf = parseplate(dname, platereadertype, datadirpath, dsheet)
    except FileNotFoundError as exc:
        raise errors.FileNotFound(str(datadirpath / dname)) from exc
    # get datatypes
    cols = rdf.columns.to_list()
    datatypes = [col for col in cols if col != "time" and col != "well"]
    # add condition and strain to dataframe
    rdf["experiment"] = experiment
    rdf["condition"] = rdf["well"].map(
        {well: wellcontents[well][0] for well in wellcontents}
    )
    rdf["strain"] = rdf["well"].map(
        {well: wellcontents[well][1] for well in wellcontents}
    )
    # drop wells with condition=strain=None
    rdf = rdf.drop(
        rdf[rdf.condition.isnull() & rdf.strain.isnull()].index
    ).reset_index(drop=True)
    return res(
        rdf=rdf,
        allconditions=allconditions,
        allstrains=allstrains,
        alldata=alldata,
        experiment=experiment,
        datatypes=datatypes,
    )


def analyseContentsofWells(datadirpath, experiment, aname, asheetnumber):
    """
    Load and parse ContentsofWells file.

    Return wellcontents, a dictionary with the contents of each well indexed
    by well, and alldata, a dictionary describing the contents of each well
    indexed by experiment.
    """
    try:
        alldata = {}
        # import contents of the wells
        anno = pd.read_excel(
            str(datadirpath / aname), index_col=0, sheet_name=asheetnumber
        )
        alldata[experiment] = []
        wellcontents = {}
        # run through and parse content of each well
        for x in np.arange(1, 13):
            for y in "ABCDEFGH":
                well = y + str(x)
                if (
                    isinstance(anno[x][y], str)
                    and anno[x][y] != "contaminated"
                ):
                    s, c = anno[x][y].split(" in ")
                    # standardise naming of wells with no strains
                    s = re.sub("(null|NULL)", "Null", s)
                    wellcontents[well] = [c.strip(), s.strip()]
                    alldata[experiment].append(
                        wellcontents[well][1] + " in " + wellcontents[well][0]
                    )
                else:
                    wellcontents[well] = [None, None]
        # create summary descriptions of the well contents
        alldata[experiment] = list(np.unique(alldata[experiment]))
        allconditions = list(
            np.unique(
                [
                    wellcontents[well][0]
                    for well in wellcontents
                    if wellcontents[well][0] is not None
                ]
            )
        )
        allstrains = list(
            np.unique(
                [
                    wellcontents[well][1]
                    for well in wellcontents
                    if wellcontents[well][0] is not None
                ]
            )
        )
        return allconditions, allstrains, alldata, wellcontents
    except FileNotFoundError as exc:
        raise errors.FileNotFound(str(datadirpath / aname)) from exc


def combine_experiments(self, min_dt, dt=None):
    """
    Combine raw data from all loaded experiments into one experiment.

    Only to be run immediately after first loading the data.

    Wells are relabelled with a prefix denoting their original
    experiment, and p.experiment_map gives the mapping back to the
    original experiment names.
    """
    print("Combining experiments...")
    # to make new r dataframe
    rdict = {col: [] for col in list(self.r.columns)}
    # find data types present in all experiments
    datatypes = self.r.select_dtypes(include="number").columns.tolist()
    dfindices = self.r.select_dtypes(include=["object"]).columns.tolist()
    # define a new universal time
    datatypes.remove("time")
    unique_t = np.sort(self.r.time.unique())
    if dt is None:
        ut_values, ut_counts = np.unique(
            np.round(np.diff(unique_t), 2), return_counts=True
        )
        dt = ut_values[ut_values > min_dt][
            np.argmax(ut_counts[ut_values > min_dt])
        ]
    print(f" Using dt = {dt:.2f}.")
    nt = np.arange(0, unique_t.max(), dt)
    # find index for all wells
    wells_index = list(self.r.groupby(dfindices).mean().index)
    # interpolate data for each well to new universal time
    print(" Interpolating data to a universal time.")
    for wi in tqdm(wells_index):
        df = self.r.query(
            " and ".join(
                [
                    f"{column} == '{value}'"
                    for column, value in zip(dfindices, wi)
                ]
            )
        )
        # local universal time depending on length of experiments
        l_nt = nt[
            np.argmin((df.time.min() - nt) ** 2) : np.argmin(
                (df.time.max() - nt) ** 2
            )
            + 1
        ]
        x = df.time.values
        xs = np.sort(x)
        # interpolate data for this well to universal time
        new_y = {}
        for datatype in datatypes:
            try:
                y = df[datatype].values
                ys = y[np.argsort(x)]
                new_y[datatype] = np.interp(l_nt, xs, ys)
            except KeyError:
                new_y[datatype] = np.nan * np.ones(l_nt.shape)
        # add to dict
        for i in range(l_nt.size):
            for column, value in zip(dfindices, wi):
                rdict[column].append(value)
            rdict["time"].append(l_nt[i])
            for datatype in datatypes:
                rdict[datatype].append(new_y[datatype][i])
    # re-make r dataframe, now with universal time
    self.r = pd.DataFrame(data=rdict)
    # make dict mapping experiment number to original name
    experiment_map = {
        expt: str(i) for i, expt in enumerate(self.r.experiment.unique())
    }
    # rename wells
    self.r.well = (
        self.r.experiment.astype("string").replace(experiment_map)
        + "_"
        + self.r.well
    )
    self.r["well"] = self.r["well"].astype("object")
    # save names of original experiments
    self.r["original_experiment"] = self.r.experiment.copy()
    self.r["experiment_id"] = (
        self.r.experiment.astype("string").replace(experiment_map).copy()
    )
    self.r["experiment_id"] = self.r["experiment_id"].astype("object")
    # write over r.experiment and sort
    self.r.experiment = self.combined
    self.r = self.r.sort_values(by=["time", "well"]).reset_index(drop=True)
    # remake summary attributes
    self.allexperiments = [self.combined]
    for mattr in [
        "allconditions",
        "allstrains",
        "allstrainsconditions",
        "datatypes",
    ]:
        uniques = set(
            [
                lattr
                for e in getattr(self, mattr)
                for lattr in getattr(self, mattr)[e]
            ]
        )
        setattr(self, mattr, {self.combined: list(uniques)})
    sort_attributes(self)
    # reset self.progress
    for key in self.progress:
        self.progress[key] = {
            self.combined: self.progress[key][next(iter(self.progress[key]))]
        }
    # remake dataframe for well content
    self.wellsdf = admin.makewellsdf(self.r)
    # remake s dataframe
    self.s = admin.make_s(self)
    # add original experiments to sc dataframe
    self.sc["original_experiment"] = self.sc.experiment.copy()
    self.sc["experiment_id"] = (
        self.sc.experiment.astype("string").replace(experiment_map).copy()
    )
    self.sc.experiment = self.combined


def sort_attributes(self):
    """Alphabetically sort attributes describing experiments."""
    self.allexperiments = sorted(self.allexperiments)
    for e in self.allconditions:
        self.allconditions[e] = sorted(
            self.allconditions[e], key=lambda x: (x.split()[-1], x.split()[0])
        )
    for e in self.allstrains:
        self.allstrains[e] = sorted(self.allstrains[e])
    for e in self.allstrainsconditions:
        self.allstrainsconditions[e] = sorted(
            self.allstrainsconditions[e],
            key=lambda x: (
                x.split(" in ")[0],
                x.split(" in ")[1].split()[-1],
                x.split(" in ")[1].split()[0],
            ),
        )
