# Author: Robert Luke <mail@robertluke.net>
#
# License: BSD (3-clause)

import os.path as op

import mne
import numpy as np
import pandas as pd
from mne.io import BaseRaw
from mne.io.constants import FIFF
from mne.transforms import _get_trans, apply_trans
from mne.utils import _check_fname, _validate_type, warn


def _read_fold_xls(fname, atlas="Juelich"):
    """Read fOLD toolbox xls file.

    The values are then manipulated in to a tidy dataframe.

    Note the xls files are not included as no license is provided.

    Parameters
    ----------
    fname : str
        Path to xls file.
    atlas : str
        Requested atlas.
    """
    page_reference = {"AAL2": 2, "AICHA": 5, "Brodmann": 8, "Juelich": 11, "Loni": 14}

    tbl = pd.read_excel(fname, sheet_name=page_reference[atlas])

    # Remove the spacing between rows
    empty_rows = np.where(np.isnan(tbl["Specificity"]))[0]
    tbl = tbl.drop(empty_rows).reset_index(drop=True)

    # Empty values in the table mean its the same as above
    for row_idx in range(1, tbl.shape[0]):
        for col_idx, col in enumerate(tbl.columns):
            if not isinstance(tbl[col][row_idx], str):
                if np.isnan(tbl[col][row_idx]):
                    tbl.iloc[row_idx, col_idx] = tbl.iloc[row_idx - 1, col_idx]

    tbl["Specificity"] = tbl["Specificity"] * 100
    tbl["brainSens"] = tbl["brainSens"] * 100
    return tbl


def _generate_montage_locations():
    """Get standard MNI montage locations in dataframe.

    Data is returned in the same format as the eeg_positions library.
    """
    # standard_1020 and standard_1005 are in MNI (fsaverage) space already,
    # but we need to undo the scaling that head_scale will do
    montage = mne.channels.make_standard_montage(
        "standard_1005", head_size=0.09700884729534559
    )
    for d in montage.dig:
        d["coord_frame"] = FIFF.FIFFV_MNE_COORD_MNI_TAL
    montage.dig[:] = montage.dig[3:]
    montage.add_mni_fiducials()  # now in fsaverage space
    coords = pd.DataFrame.from_dict(montage.get_positions()["ch_pos"]).T
    coords["label"] = coords.index
    coords = coords.rename(columns={0: "x", 1: "y", 2: "z"})

    return coords.reset_index(drop=True)


def _find_closest_standard_location(position, reference, *, out="label"):
    """Return closest montage label to coordinates.

    Parameters
    ----------
    position : array, shape (3,)
        Coordinates.
    reference : dataframe
        As generated by _generate_montage_locations.
    trans_pos : str
        Apply a transformation to positions to specified frame.
        Use None for no transformation.
    """
    from scipy.spatial.distance import cdist

    p0 = np.array(position)
    p0.shape = (-1, 3)
    head_mri_t, _ = _get_trans("fsaverage", "head", "mri")
    p0 = apply_trans(head_mri_t, p0)
    dists = cdist(p0, np.asarray(reference[["x", "y", "z"]], float))

    if out == "label":
        min_idx = np.argmin(dists)
        return reference["label"][min_idx]
    else:
        assert out == "dists"
        return dists


def fold_landmark_specificity(
    raw, landmark, fold_files=None, atlas="Juelich", interpolate=False
):
    """Return the specificity of each channel to a specified brain landmark.

    Parameters
    ----------
    raw : BaseRaw
        The fNIRS data.
    landmark : str
        Landmark of interest. Must be present in fOLD toolbox data file.
    fold_files : list | path-like | None
        If None, will use the MNE_NIRS_FOLD_PATH config variable.
        If path-like, should be a path to a directory containing '10-10.xls'
        and '10-5.xls'. If list, should be paths to the fold toolbox files.
        See the Notes section of :func:`~mne_nirs.io.fold_channel_specificity`
        for details.
    atlas : str
        Brain atlas to use.
    interpolate : bool
        If the optimal source-detector pair is not found in the fOLD files
        False (default) will yield no results for that pairing, whereas
        True will use the next closest match. See Notes of
        :func:`mne_nirs.io.fold_channel_specificity` for an example.

        .. warning::
           The sensitivity profile can differ substantially for nearest
           neighbors, so use ``interpolate=True`` with caution.

    Returns
    -------
    spec : array
        Specificity values for each channel to brain landmark.

    See Also
    --------
    fold_landmark_specificity

    Notes
    -----
    Specificity values are provided by the fOLD toolbox
    :footcite:`morais2018fnirs` excel files. See the Notes section of
    :func:`~mne_nirs.io.fold_channel_specificity` for more details.

    References
    ----------
    .. footbibliography::
    """
    _validate_type(landmark, str, "landmark")
    _validate_type(raw, BaseRaw, "raw")

    reference_locations = _generate_montage_locations()

    fold_tbl = _check_load_fold(fold_files, atlas)

    specificity = np.zeros(len(raw.ch_names))
    for cidx in range(len(raw.ch_names)):
        tbl = _source_detector_fold_table(
            raw, cidx, reference_locations, fold_tbl, interpolate
        )

        if len(tbl) > 0:
            tbl["ContainsLmk"] = [landmark in la for la in tbl["Landmark"]]
            tbl = tbl.query("ContainsLmk == True")["Specificity"]

        if len(tbl) == 0:
            continue
            # print(f"No data for {src_name}-{det_name}")
        elif len(tbl) == 1:
            specificity[cidx] = tbl.values[0]
        else:
            raise RuntimeError("Multiple specificity values returned")

    return np.array(specificity)


def fold_channel_specificity(raw, fold_files=None, atlas="Juelich", interpolate=False):
    """Return the landmarks and specificity a channel is sensitive to.

    Parameters
    ----------
    raw : BaseRaw
        The fNIRS data.
    fold_files : list | path-like | None
        If None, will use the MNE_NIRS_FOLD_PATH config variable.
        If path-like, should be a path to a directory containing '10-10.xls'
        and '10-5.xls'. If list, should be paths to the fold toolbox files.
        See Notes for details.
    atlas : str
        Brain atlas to use.
    interpolate : bool
        If the optimal source-detector pair is not found in the fOLD files
        False (default) will yield no results for that pairing, whereas
        True will use the next closest match. See Notes for an example.

        .. warning:: The sensitivity profile can differ substantially for
                     nearest neighbors, so use ``interpolate=True`` with
                     caution.

    Returns
    -------
    spec : list of DataFrame
        List of dataframes, one for each channel.

    See Also
    --------
    fold_landmark_specificity

    Notes
    -----
    **fOLD Toolbox**

    Specificity values are provided by the fOLD toolbox
    :footcite:`morais2018fnirs` excel files.
    For licensing reasons, these files are not distributed with MNE-NIRS.
    You need to download them from
    `the author's website <https://github.com/nirx/fOLD-public>`__.
    To automatically utilize the ``MNE_NIRS_FOLD_PATH`` config for the
    ``fold_files`` parameter, you can download the entire ``fOLD-public``
    repository `as a zip <https://github.com/nirx/fOLD-public/archive/refs/heads/master.zip>`__
    and expand it to some suitable location like
    ``~/mne_data/fOLD/fOLD-public-master``, and then set the config value
    on your machine by using :func:`mne:mne.set_config` like::

        >>> mne.set_config('MNE_NIRS_FOLD_PATH', '~/mne_data/fOLD/fOLD-public-master/Supplementary')

    From then on, :func:`~mne_nirs.io.fold_channel_specificity` and
    :func:`~mne_nirs.io.fold_landmark_specificity` will automatically use this
    directory to find the fOLD xls files when you pass ``fold_files=None``
    (which is the default). We recommend following this procedure so that
    the files can be reused automatically.

    **Interpolation**

    For an example of interpolation, consider the pairings P5-PO7 nor P6-PO8,
    neither of which are listed in the fOLD toolbox:

    - With ``interpolate=False``, the returned ``spec`` will not have entries
      for these channels. When
    - With ``interpolate=True``, entries like P7-PO7 and P8-PO8, respectively,
      might be used instead. A warning is emitted if such substitutions are
      made.

    References
    ----------
    .. footbibliography::
    """  # noqa: E501
    _validate_type(raw, BaseRaw, "raw")

    reference_locations = _generate_montage_locations()

    fold_tbl = _check_load_fold(fold_files, atlas)

    chan_spec = list()
    for cidx in range(len(raw.ch_names)):
        tbl = _source_detector_fold_table(
            raw, cidx, reference_locations, fold_tbl, interpolate
        )
        chan_spec.append(tbl.reset_index(drop=True))

    return chan_spec


def _check_load_fold(fold_files, atlas):
    _validate_type(fold_files, (list, "path-like", None), "fold_files")
    if fold_files is None:
        fold_files = mne.get_config("MNE_NIRS_FOLD_PATH")
        if fold_files is None:
            raise ValueError(
                "MNE_NIRS_FOLD_PATH not set, either set it using "
                "mne.set_config or pass fold_files as str or list"
            )
    if not isinstance(fold_files, list):  # path-like
        fold_files = _check_fname(
            fold_files,
            overwrite="read",
            must_exist=True,
            name="fold_files",
            need_dir=True,
        )
        fold_files = [op.join(fold_files, f"10-{x}.xls") for x in (5, 10)]

    fold_tbl = pd.DataFrame()
    for fi, fname in enumerate(fold_files):
        fname = _check_fname(
            fname, overwrite="read", must_exist=True, name=f"fold_files[{fi}]"
        )
        fold_tbl = pd.concat(
            [fold_tbl, _read_fold_xls(fname, atlas=atlas)], ignore_index=True
        )
    return fold_tbl


def _source_detector_fold_table(raw, cidx, reference, fold_tbl, interpolate):
    src = raw.info["chs"][cidx]["loc"][3:6]
    det = raw.info["chs"][cidx]["loc"][6:9]

    ref_lab = list(reference["label"])
    dists = _find_closest_standard_location([src, det], reference, out="dists")
    src_min, det_min = np.argmin(dists, axis=1)
    src_name, det_name = ref_lab[src_min], ref_lab[det_min]

    tbl = fold_tbl.query("Source == @src_name and Detector == @det_name")
    dist = np.linalg.norm(dists[[0, 1], [src_min, det_min]])
    # Try reversing source and detector
    if len(tbl) == 0:
        tbl = fold_tbl.query("Source == @det_name and Detector == @src_name")
    if len(tbl) == 0 and interpolate:
        # Try something hopefully not too terrible: pick the one with the
        # smallest net distance
        good = np.isin(fold_tbl["Source"], reference["label"]) & np.isin(
            fold_tbl["Detector"], reference["label"]
        )
        assert good.any()
        tbl = fold_tbl[good]
        assert len(tbl)
        src_idx = [ref_lab.index(src) for src in tbl["Source"]]
        det_idx = [ref_lab.index(det) for det in tbl["Detector"]]
        # Original
        tot_dist = np.linalg.norm([dists[0, src_idx], dists[1, det_idx]], axis=0)
        assert tot_dist.shape == (len(tbl),)
        idx = np.argmin(tot_dist)
        dist_1 = tot_dist[idx]
        src_1, det_1 = ref_lab[src_idx[idx]], ref_lab[det_idx[idx]]
        # And the reverse
        tot_dist = np.linalg.norm([dists[0, det_idx], dists[1, src_idx]], axis=0)
        idx = np.argmin(tot_dist)
        dist_2 = tot_dist[idx]
        src_2, det_2 = ref_lab[det_idx[idx]], ref_lab[src_idx[idx]]
        if dist_1 < dist_2:
            new_dist, src_use, det_use = dist_1, src_1, det_1
        else:
            new_dist, src_use, det_use = dist_2, det_2, src_2
        warn(
            "No fOLD table entry for best matching source/detector pair "
            f"{src_name}/{det_name} (RMS distance to the channel positions "
            f"was {1000 * dist:0.1f} mm) for channel index {cidx}, "
            "using next smallest available "
            f"src/det pairing {src_use}/{det_use} (RMS distance "
            f"{1000 * new_dist:0.1f} mm). Consider setting your channel "
            "positions to standard 10-05 locations using raw.set_montage "
            "if your pair does show up in the tables.",
            module="mne_nirs",
            ignore_namespaces=("mne", "mne_nirs"),
        )
        tbl = fold_tbl.query("Source == @src_use and Detector == @det_use")
        tbl = tbl.copy()
        tbl["BestSource"] = src_name
        tbl["BestDetector"] = det_name
        tbl["BestMatchDistance"] = dist
        tbl["MatchDistance"] = new_dist
        assert len(tbl)
    else:
        tbl = tbl.copy()
        tbl["BestSource"] = src_name
        tbl["BestDetector"] = det_name
        tbl["BestMatchDistance"] = dist
        tbl["MatchDistance"] = dist

    tbl = tbl.copy()  # don't get warnings about setting values later
    return tbl
