"""Preprocessing module.

Convert raw LA-ICP-MS ``.csv`` files from Agilent or ThermoFisher
quadrupole instruments into a single :class:`pandas.DataFrame` ready
for processing with :class:`~lasertram.tram.LaserTRAM`.
"""

from __future__ import annotations

import re
from pathlib import Path

import numpy as np
import pandas as pd
from dateutil.parser import parse
from tabulate import tabulate
from tqdm import tqdm

data_examples_dir = Path(__file__).parent.parent / "data"


def extract_agilent_data(file: str | Path) -> dict:
    """Read raw output from an Agilent quadrupole ``.csv`` file.

    Parameters
    ----------
    file : str or pathlib.Path
        Path to the ``.csv`` file to be extracted.

    Returns
    -------
    dict
        Keys ``"timestamp"``, ``"file"``, ``"sample"``, ``"data"``.
    """
    # import data
    # extract sample name
    # extract timestamp
    # extract data and make headers ready for lasertram

    df = pd.read_csv(file, sep="\t", header=None)

    sample = df.iloc[0, 0].split("\\")[-1].split(".")[0].replace("_", "-")

    timestamp = parse(df.iloc[2, 0].split(" ")[7] + " " + df.iloc[2, 0].split(" ")[8])

    data = pd.DataFrame([sub.split(",") for sub in df.iloc[3:-1, 0]])

    header = data.iloc[0, :]
    data = data[1:]
    data.columns = header
    newcols = []
    for s in data.columns.tolist():
        l = re.findall(r"(\d+|[A-Za-z]+)", s)
        if "Time" in l:
            newcols.append(l[0])
        else:

            newcols.append(l[1] + l[0])
    data.columns = newcols

    return {"timestamp": timestamp, "file": file, "sample": sample, "data": data}


def extract_thermo_data(file: str | Path) -> dict | None:
    """Read raw output from a ThermoFisher quadrupole ``.csv`` file.

    Parameters
    ----------
    file : str or pathlib.Path
        Path to the ``.csv`` file to be extracted.

    Returns
    -------
    dict or None
        Keys ``"timestamp"``, ``"file"``, ``"sample"``, ``"data"``.
        Returns ``None`` if the file is empty.
    """
    assert isinstance(file, (str, Path)), "file must be a str or pathlib.Path object"

    if isinstance(file, Path):
        pass
    else:
        file = Path(file)

    # gets the top row in your csv and turns it into a pandas series
    try:
        top = pd.read_csv(file, nrows=0)

        # since it is only 1 long it is also the column name
        # extract that as a list
        sample = list(top.columns)

        # turn that list value to a string
        sample = str(sample[0])

        # because its a string it can be split
        # split at : removes the time stamp
        sample = sample.split(":")[0]

        # .strip() removes leading and trailing spaces
        sample = sample.strip()

        # replace middle spaces with _ because spaces are bad
        nospace = sample.replace(" ", "_")

        # get the timestamp by splitting the string by the previously
        # designated sample. Also drops the colon in front of the date
        timestamp = top.columns.tolist()[0].split(sample)[1:][0][1:]

        timestamp = parse(timestamp)

        # import data
        # remove the top rows. Double check that your header is the specified
        # amount of rows to be skipped in 'skiprows' argument
        data = pd.read_csv(file, skiprows=13)
        # drop empty column at the end
        data.drop(data.columns[len(data.columns) - 1], axis=1, inplace=True)

        # remove dwell time row beneath header row
        data = data.dropna()

        return {"timestamp": timestamp, "file": file, "sample": nospace, "data": data}

    except pd.errors.EmptyDataError:
        return None


def make_lt_ready_folder(folder: str | Path, quad_type: str) -> pd.DataFrame:
    """Combine a folder of raw ``.csv`` files into a LaserTRAM-ready DataFrame.

    Parameters
    ----------
    folder : str or pathlib.Path
        Path to the folder containing ``.csv`` files. The folder should
        contain **only** data files.
    quad_type : str
        ``"agilent"`` or ``"thermo"``.

    Returns
    -------
    pandas.DataFrame
        DataFrame ready for :meth:`~lasertram.tram.LaserTRAM.get_data`.

    Examples
    --------
    >>> from lasertram.helpers import preprocessing
    >>> df = preprocessing.make_lt_ready_folder("path/to/csvs", "thermo")
    """

    if isinstance(folder, Path):
        pass
    else:
        folder = Path(folder)
    assert (
        folder.is_dir()
    ), f"{folder} is not a directory, please choose a directory to your data .csv files"

    print(
        "Processing the all .csv files in the following folder for LaserTRAM input format:\n"
    )
    print(tabulate([[str(folder)]], headers=["Folder Path"], tablefmt="pipe"))
    print("\n")
    my_dict = {}
    files = [f for f in folder.glob("*.csv")]
    # GET METADATA
    # establish progress bar
    pbar = tqdm(
        files, desc="Extracting metadata from files", unit="file", total=len(files)
    )
    temp = None
    empty_files = []
    for i in pbar:
        # rename description to match file
        pbar.set_description(f"Extracting metadata from {i.name}")

        if quad_type == "thermo":
            temp = extract_thermo_data(i)

        elif quad_type == "agilent":
            temp = extract_agilent_data(i)

        # if the file is empty extract_thermo_data will be none
        if temp is None:
            empty_files.append(i)
        else:
            my_dict[temp["timestamp"]] = temp

    my_dict = dict(sorted(my_dict.items()))
    # if any empty files display them in a table
    if len(empty_files) > 0:
        print("the following files were skipped because they contained no data:\n")
        table = [[e.name] for e in empty_files]
        print(tabulate(table, headers=["Empty files"], tablefmt="pipe"))
        print("\n")
    # GET DATA FROM DICTIONARY AND CONCAT ALL THE DATA INTO ONE DF
    outdf = pd.DataFrame()
    pbar2 = tqdm(my_dict, desc="Combining individual files", unit="file")
    for timestamp in pbar2:

        pbar2.set_description(
            f"Adding data from {Path(my_dict[timestamp]['file']).name}"
        )

        samplelabel = pd.DataFrame(
            np.repeat(
                my_dict[timestamp]["sample"], my_dict[timestamp]["data"].shape[0]
            ),
            columns=["SampleLabel"],
            index=my_dict[timestamp]["data"].index,
        )
        ts = pd.DataFrame(
            np.repeat(
                my_dict[timestamp]["timestamp"], my_dict[timestamp]["data"].shape[0]
            ),
            columns=["timestamp"],
            index=my_dict[timestamp]["data"].index,
        )
        df = pd.concat([ts, samplelabel, my_dict[timestamp]["data"]], axis="columns")

        outdf = pd.concat([outdf, df])
        outdf.index = np.arange(outdf.shape[0], dtype=int)
    print("Success.\n")
    return outdf


def make_lt_ready_file(file: str | Path, quad_type: str) -> pd.DataFrame:
    """Convert a single raw ``.csv`` file into a LaserTRAM-ready DataFrame.

    Parameters
    ----------
    file : str or pathlib.Path
        Path to the ``.csv`` file.
    quad_type : str
        ``"agilent"`` or ``"thermo"``.

    Returns
    -------
    pandas.DataFrame
        DataFrame ready for :meth:`~lasertram.tram.LaserTRAM.get_data`.

    Raises
    ------
    ValueError
        If ``quad_type`` is not ``"thermo"`` or ``"agilent"``.
    """

    if isinstance(file, Path):
        pass
    else:
        file = Path(file)

    assert file.name.endswith(".csv"), f"File '{file}' does not have a CSV extension."

    if quad_type == "thermo":
        temp = extract_thermo_data(file)

    elif quad_type == "agilent":
        temp = extract_agilent_data(file)
    else:
        temp = None

    if temp:
        outdf = temp["data"]
        outdf.insert(0, "SampleLabel", temp["sample"])
        outdf.insert(0, "timestamp", temp["timestamp"])

    else:
        raise ValueError("please choose either 'thermo' or 'agilent' for quad_type")

    return outdf


def load_test_rawdata() -> pd.DataFrame:
    """Load example raw LA-ICP-MS data.

    The data accompany the following manuscript:

        Lubbers, J., Kent, A., & Russo, C. (2025). lasertram: a Python
        library for time resolved analysis of laser ablation inductively
        coupled plasma mass spectrometry data.

    Returns
    -------
    pandas.DataFrame
        Raw counts-per-second data indexed by ``SampleLabel``.

    Examples
    --------
    >>> from lasertram.helpers import preprocessing
    >>> raw_data = preprocessing.load_test_rawdata()
    >>> raw_data.head()
    """

    # current_path = Path(__file__).parent
    # lt_ready = pd.read_excel(
    #     current_path.parents[1]
    #     / "test_data"
    #     / "computers_and_geosciences_examples"
    #     / "2022-05-10_LT_ready.xlsx"
    # ).set_index("SampleLabel")

    lt_ready = pd.read_excel(
        data_examples_dir
        / "computers_and_geosciences_examples"
        / "2022-05-10_LT_ready.xlsx"
    ).set_index("SampleLabel")

    return lt_ready


def load_test_intervals() -> pd.DataFrame:
    """Load example interval selections.

    The data accompany the following manuscript:

        Lubbers, J., Kent, A., & Russo, C. (2025). lasertram: a Python
        library for time resolved analysis of laser ablation inductively
        coupled plasma mass spectrometry data.

    Returns
    -------
    pandas.DataFrame
        Interval definitions indexed by ``Spot``.

    Examples
    --------
    >>> from lasertram.helpers import preprocessing
    >>> intervals = preprocessing.load_test_intervals()
    """

    # current_path = Path(__file__).parent

    # intervals = pd.read_excel(
    #     current_path.parents[1]
    #     / "test_data"
    #     / "computers_and_geosciences_examples"
    #     / "example_intervals.xlsx"
    # ).set_index("Spot")

    intervals = pd.read_excel(
        data_examples_dir
        / "computers_and_geosciences_examples"
        / "example_intervals.xlsx"
    ).set_index("Spot")

    return intervals


def load_test_int_std_comps() -> pd.DataFrame:
    """Load example internal standard compositions.

    The data accompany the following manuscript:

        Lubbers, J., Kent, A., & Russo, C. (2025). lasertram: a Python
        library for time resolved analysis of laser ablation inductively
        coupled plasma mass spectrometry data.

    Returns
    -------
    pandas.DataFrame
        Internal standard concentrations and uncertainties.

    Examples
    --------
    >>> from lasertram.helpers import preprocessing
    >>> int_std_comps = preprocessing.load_test_int_std_comps()
    """

    # current_path = Path(__file__).parent

    # concentrations = pd.read_excel(
    #     current_path.parents[1]
    #     / "test_data"
    #     / "computers_and_geosciences_examples"
    #     / "example_internal_std.xlsx"
    # )

    concentrations = pd.read_excel(
        data_examples_dir
        / "computers_and_geosciences_examples"
        / "example_internal_std.xlsx"
    )

    return concentrations
