"""Utility functions for handling atlas files and labels.

This module provides functions to parse label files, load atlas volumes,
and manage surface-based atlases using FreeSurfer annotations.
"""

import os
import numpy as np


def fetch_labels(labels):
    """Parse a labels input.

    A list is returned as-is. If ``labels`` is a string, it is treated as the
    path to an XML file and parsed for entries within a
    ``<data><label><name>...</name></label></data>`` structure.

    Parameters
    ----------
    labels : list of str or str
        A list of label names or a path to an XML file containing labels.

    Returns
    -------
    list of str
        Parsed label names.

    Raises
    ------
    ValueError
        If the XML file is invalid, cannot be parsed, contains no labels, or
        if ``labels`` is neither a list nor a string.

    Examples
    --------
    >>> fetch_labels(['A', 'B'])
    ['A', 'B']
    >>> fetch_labels('atlas.xml')  # doctest: +SKIP
    ['Region1', 'Region2']
    """
    if isinstance(labels, str):
        import xml.etree.ElementTree as ET

        try:
            tree = ET.parse(labels)
            root = tree.getroot()
            data = root.find("data")
            if data is None:
                raise ValueError("Invalid XML file: missing 'data' element.")
            label_list = []
            for label in data.findall("label"):
                name_elem = label.find("name")
                if name_elem is not None:
                    label_list.append(name_elem.text)
            if not label_list:
                raise ValueError("No labels found in the XML file.")
            return label_list
        except Exception as e:
            raise ValueError(f"Error processing XML file {labels}: {e}")
    elif isinstance(labels, list):
        return labels
    else:
        raise ValueError(f"Invalid labels type: {type(labels)}")


def pack_vol_output(file):
    """Load an atlas file and return volume data and header.

    Parameters
    ----------
    file : str or Nifti1Image
        Path to a NIfTI/NPZ file or a loaded
        :class:`~nibabel.nifti1.Nifti1Image`.

    Returns
    -------
    dict
        Dictionary with ``'vol'`` and ``'hdr'`` entries.

    Raises
    ------
    ValueError
        If the file format or object type is not supported.

    Examples
    --------
    >>> pack_vol_output('atlas.nii.gz')  # doctest: +SKIP
    {'vol': array(...), 'hdr': array(...)}
    """
    if isinstance(file, str):
        path = os.path.abspath(file)
        _, ext = os.path.splitext(file)
        ext = ext.lower()

        if ext in [".nii", ".gz", ".nii.gz"]:
            import nibabel as nib

            img = nib.load(file)
            vol_data = img.get_fdata(dtype=np.float32)
            hdr_matrix = img.affine
            return {
                "vol": vol_data,
                "hdr": hdr_matrix,
            }

        elif ext == ".npz":
            arch = np.load(path, allow_pickle=True)
            vol_data = arch["vol"]
            hdr_matrix = arch["hdr"]
            return {
                "vol": vol_data,
                "hdr": hdr_matrix,
            }
        else:
            raise ValueError(f"Unrecognized file format '{ext}' for path: {path}")
    else:
        from nibabel.nifti1 import Nifti1Image

        if isinstance(file, Nifti1Image):
            vol_data = file.get_fdata(dtype=np.float32)
            hdr_matrix = file.affine
            return {
                "vol": vol_data,
                "hdr": hdr_matrix,
            }
        else:
            raise ValueError("Unsupported type for pack_vol_output")


def pack_surf_output(
    atlas_name, fetcher, subject: str = "fsaverage", subjects_dir: str = None, **kwargs
):
    """Load a surface-based atlas using FreeSurfer annotations.

    Parameters
    ----------
    atlas_name : str
        Name of the atlas (e.g., ``'aparc'``).
    fetcher : callable or None
        Function used to download the atlas if necessary.
    subject : str, optional
        Subject identifier, by default ``'fsaverage'``.
    subjects_dir : str or None, optional
        FreeSurfer subjects directory, by default ``None``.
    **kwargs
        Additional keyword arguments passed to the fetcher.

    Returns
    -------
    dict
        Dictionary with ``'vol'``, ``'hdr'``, ``'labels'``, ``'indexes'`` and
        ``'regions'`` keys where ``'regions'`` maps region names to vertex
        indices.

    Raises
    ------
    ValueError
        If the atlas cannot be located or fetched.

    Examples
    --------
    >>> pack_surf_output('aparc', None)  # doctest: +SKIP
    {'vol': [...], 'hdr': None, 'labels': array([...]), 'indexes': array([...]),
     'regions': {'Region': array([...])}}
    """
    # Determine subjects_dir: use provided or from MNE config
    import mne

    if subjects_dir is None:
        subjects_dir = mne.get_config("SUBJECTS_DIR", None)
        if subjects_dir is None:
            import os

            subjects_dir = os.path.join(mne.datasets.sample.data_path(), "subjects")
    from pathlib import Path

    subjects_dir = Path(subjects_dir)
    if fetcher is None:
        try:
            labels = mne.read_labels_from_annot(
                subject,
                atlas_name,
                subjects_dir=subjects_dir,
                **kwargs,
            )
        except Exception:
            mne.datasets.fetch_fsaverage(subjects_dir=subjects_dir, verbose=True)
            labels = mne.read_labels_from_annot(
                subject,
                atlas_name,
                subjects_dir=subjects_dir,
                **kwargs,
            )
    else:
        try:
            labels = fetcher(subject=subject, subjects_dir=subjects_dir, **kwargs)
        except Exception:
            fetcher(subjects_dir=subjects_dir, **kwargs)
            labels = mne.read_labels_from_annot(
                subject,
                atlas_name,
                subjects_dir=subjects_dir,
                **kwargs,
            )

    src = mne.setup_source_space(
        subject,
        spacing="oct6",
        subjects_dir=subjects_dir,
        add_dist=False,
    )
    lh_vert = src[0]["vertno"]  # Left hemisphere vertices
    rh_vert = src[1]["vertno"]  # Right hemisphere vertices

    # Map label names to indices in the vertex arrays.
    from collections import defaultdict

    region_vertices_lh = defaultdict(list)
    region_vertices_rh = defaultdict(list)
    labmap_lh = {}
    labmap_rh = {}

    for label in labels:
        if label.hemi == "lh":
            match = np.nonzero(np.isin(lh_vert, label.vertices))[0]
            verts = lh_vert[match]
            region_vertices_lh[label.name].extend(verts.tolist())
            for idx in match:
                labmap_lh[idx] = label.name
        elif label.hemi == "rh":
            match = np.nonzero(np.isin(rh_vert, label.vertices))[0]
            verts = rh_vert[match]
            region_vertices_rh[label.name].extend(verts.tolist())
            for idx in match:
                labmap_rh[idx] = label.name

    region_vertices_lh = {
        k: np.array(v, dtype=int) for k, v in region_vertices_lh.items()
    }
    region_vertices_rh = {
        k: np.array(v, dtype=int) for k, v in region_vertices_rh.items()
    }
    region_vertices = {**region_vertices_lh, **region_vertices_rh}

    indexes_lh = np.sort(np.array(list(labmap_lh.keys())))
    labels_lh = np.array([labmap_lh[i] for i in indexes_lh])
    vmap_lh = lh_vert[indexes_lh]

    indexes_rh = np.sort(np.array(list(labmap_rh.keys())))
    labels_rh = np.array([labmap_rh[i] for i in indexes_rh])
    vmap_rh = rh_vert[indexes_rh]

    labels_combined = np.concatenate([labels_lh, labels_rh])
    indexes_combined = np.concatenate([vmap_lh, vmap_rh])

    return {
        "vol": [lh_vert, rh_vert],
        "hdr": None,
        "labels": labels_combined,
        "indexes": indexes_combined,
        "regions": region_vertices,
    }
