"""Read Waters masslynx files."""

from datetime import datetime
from pathlib import Path
import re
import rainbow as rb

from ..chromatogram import Chromatogram1D, Chromatogram2D, Chromatograms


def can_read(f):
    """Determine whether the given file/folder is a Waters MassLynx RAW file.

    Parameters
    ----------
    f : str
        The file or folder to evaluate

    Returns
    -------
    bool
        True if the given path is a MassLynx RAW file, False otherwise.
    """
    f = Path(f)

    if f.suffix.lower() != ".raw":
        return False
    if not f.is_dir():
        return False
    if not (f / "_HEADER.TXT").exists() or not (f / "_INLET.INF").exists():
        return False

    return True


def read(path, detector=None, metadata=None, precision=3):
    """Read the given MassLynx RAW file.

    Parameters
    ----------
    path : str
        The path of the RAW file to read
    detector : str
        The detector to read from. If None, returns data from all detectors.
    metadata : dict
        The metadata to use to select chromatogram(s) in the file. The given fields will have to be matched by the chromatogram's metadata to be selected. For example, `{'polarity`: '+'}` to select ESI+ MS data. If None, returns all chromatograms.
    precision : int
        The precision for reading the different channels/mz of the chromatogram. The higher the number, the more digits.

    Returns
    --------
    chromatograpy.Chromatogram or list of
        The chromatogram(s) contained in the file.
    """
    raw_file = MassLynxRaw(path, precision)

    if metadata is not None:
        if detector is not None:
            metadata["Detector"] = detector
        chromatograms = raw_file.by_metadata(metadata)
    elif detector is not None:
        chromatograms = raw_file.by_detector(detector)
    else:
        chromatograms = raw_file.chromatograms

    if len(chromatograms) == 0:
        raise ValueError(
            f"MassLynx RAW file `{path}` does not contain a chromatogram matching the selection criteria provided"
        )
    if len(chromatograms) == 1:
        return chromatograms[0]
    return chromatograms


def translate_position(pos):
    """Translate a MassLynx bottle number into a chromatograpy position."""
    return pos.replace(",", "")


metadata_name_map = {
    "Acquired Name": "Name",
    "Inlet Method": "Method",
    "Bottle Number": "Position",
    "Job Code": "Sequence",
}


def translate_metadata(metadata):
    """Translate metadata from MassLynx format to Chromatograpy format.

    Parameters
    ----------
    metadata : dict
        The metadata from the _HEADER.TXT file from a MassLynx .RAW file, as returned by `read_metadata`

    To do:
    ------
    - If the locale of the MassLynx computer is not the same as the analysis computer, reading the datetime may fail.

    Returns
    -------
    dict
        The metadata dictionary with chromatograpy-based variable names
    """
    translated_metadata = {
        k if k not in metadata_name_map.keys() else metadata_name_map[k]: v
        for k, v in metadata.items()
        if v != "" and k not in ["Acquired Date", "Acquired Time", "Bottle Number"]
    }

    str_timestamp = metadata["Acquired Date"] + " " + metadata["Acquired Time"]
    translated_metadata["Timestamp"] = datetime.strptime(
        str_timestamp, "%d-%b-%Y %H:%M:%S"
    )

    translated_metadata["Position"] = translate_position(metadata["Bottle Number"])

    return translated_metadata


class MassLynxRaw:
    """A  .raw file generated by the Waters MassLynx software"""

    def __init__(self, f, prec):
        """Open the given MassLynx .RAW file"""
        chromatograms = read_data(f, prec=prec)
        self.metadata = read_metadata(f)

        self.chromatograms = Chromatograms()
        for chrom, f, detector, chrom_metadata in chromatograms:
            chrom.metadata = self.metadata.copy()
            chrom.metadata["Detector"] = detector
            chrom_metadata["Raw file"] = f
            chrom.metadata.update(chrom_metadata)
            self.chromatograms.append(chrom)

    @property
    def detectors(self):
        return set([chrom.metadata["Detector"] for chrom in self.chromatograms])

    def by_detector(self, detector):
        """Return the chromatograms for the given detector name.

        Parameters
        -----------
        detector: str
            The name of the detector to retrieve data from. E.g. "MS" for QDa, "UV" for PDA.

        Returns
        -------
        list
            The list of chromatograms for the given detector
        """
        return [
            chrom
            for chrom in self.chromatograms
            if chrom.metadata["Detector"] == detector
        ]

    def by_metadata(self, *args, **kwargs):
        """Return the chromatograms matching the queried metadata

        Parameters
        ----------
        A dict with the metadata to match
        OR
        The key and value for the metadata to match as keyword arguments

        Returns
        -------
        list
            The list of chromatograms matching the query
        """
        if len(args) > 0:
            to_match = args[0]
        else:
            to_match = {}
        to_match.update(kwargs)

        return [
            chrom
            for chrom in self.chromatograms
            if all(
                [
                    k in chrom.metadata and chrom.metadata[k] == v
                    for k, v in to_match.items()
                ]
            )
        ]


def read_data(f, prec=3):
    """Read the data in a given MassLynx .RAW file."""
    # Note: Conversion to string required because rainbow checks that the path is a string.
    raw_data = rb.read(str(f), prec=prec)

    data = []
    for df in raw_data.datafiles:
        file_data = df.data.squeeze()
        if len(file_data.shape) == 1:
            chrom = Chromatogram1D(file_data, index=df.xlabels)
        elif len(file_data.shape) == 2:
            chrom = Chromatogram2D(file_data, index=df.xlabels, columns=df.ylabels)
        data.append((chrom, df.name, df.detector, df.metadata))
    return data


def read_metadata(f):
    """Read the metadata of a given MassLynx .RAW file."""
    metadata = {}
    with open(Path(f) / "_HEADER.TXT", "r") as metadata_file:
        for line in metadata_file.read().splitlines():
            k, v = line.split(":", 1)
            metadata[k[3:]] = v.strip()

    with open(Path(f) / "_INLET.INF", "r", encoding="iso-8859-1") as metadata_file:
        injection_volume = re.findall(
            r".*Injection Volume \(ul\)   -   (\d*.\d*)", metadata_file.read()
        )
        metadata["Injection volume"] = float(injection_volume[0])

    return translate_metadata(metadata)
