# functions to load the data generated by the plate reader
# and to parse the annotation file
import pandas as pd
import numpy as np
import om_code.omerrors as errors
from om_code.parseplate import parseplate


def loaddatafiles(
    platereadertype, wdirpath, dname, dsheetnumber, aname, asheetnumber
):
    """Call functions to parse data and metadate from the input files."""
    experiment = dname.split(".")[0]

    # import and parse plate contents file
    allconditions, allstrains, alldata, wellcontents = analyseContentsofWells(
        wdirpath, experiment, aname, asheetnumber
    )
    # import and parse data file created by plate reader
    try:
        print("loading", dname)
        rdf = parseplate(dname, platereadertype, wdirpath, dsheetnumber)
    except FileNotFoundError:
        raise errors.FileNotFound(str(wdirpath / dname))

    # 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}
    )

    return rdf, allconditions, allstrains, alldata, experiment, datatypes


###


def analyseContentsofWells(wdirpath, 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.
    """
    import re

    try:
        alldata = {}
        # import contents of the wells
        anno = pd.read_excel(
            str(wdirpath / 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:
        raise errors.FileNotFound(str(wdirpath / aname))
