"""Read test or example data."""

import gzip
import json
from os.path import dirname, exists, join as pjoin
import pickle

import numpy as np
from scipy.sparse import load_npz

from dipy.core.gradients import GradientTable, gradient_table
from dipy.core.sphere import HemiSphere, Sphere
from dipy.data.fetcher import (
    fetch_30_bundle_atlas_hcp842,
    fetch_bundle_atlas_hcp842,
    fetch_bundle_fa_hcp,
    fetch_bundle_warp_dataset,
    fetch_bundles_2_subjects,
    fetch_cenir_multib,
    fetch_cfin_multib,
    fetch_deepn4_test,
    fetch_deepn4_tf_weights,
    fetch_deepn4_torch_weights,
    fetch_disco1_dataset,
    fetch_disco2_dataset,
    fetch_disco3_dataset,
    fetch_disco_dataset,
    fetch_evac_test,
    fetch_evac_tf_weights,
    fetch_evac_torch_weights,
    fetch_gold_standard_io,
    fetch_hbn,
    fetch_hcp,
    fetch_isbi2013_2shell,
    fetch_ivim,
    fetch_mni_template,
    fetch_ptt_minimal_dataset,
    fetch_resdnn_tf_weights,
    fetch_resdnn_torch_weights,
    fetch_scil_b0,
    fetch_sherbrooke_3shell,
    fetch_stanford_hardi,
    fetch_stanford_labels,
    fetch_stanford_pve_maps,
    fetch_stanford_t1,
    fetch_stanford_tracks,
    fetch_syn_data,
    fetch_synb0_test,
    fetch_synb0_weights,
    fetch_taiwan_ntu_dsi,
    fetch_target_tractogram_hcp,
    fetch_tissue_data,
    get_bundle_atlas_hcp842,
    get_fnames,
    get_target_tractogram_hcp,
    get_two_hcp842_bundles,
    read_DiB_70_lte_pte_ste,
    read_DiB_217_lte_pte_ste,
    read_bundles_2_subjects,
    read_cenir_multib,
    read_cfin_dwi,
    read_cfin_t1,
    read_five_af_bundles,
    read_isbi2013_2shell,
    read_ivim,
    read_mni_template,
    read_qte_lte_pte,
    read_scil_b0,
    read_sherbrooke_3shell,
    read_stanford_hardi,
    read_stanford_labels,
    read_stanford_pve_maps,
    read_stanford_t1,
    read_syn_data,
    read_taiwan_ntu_dsi,
    read_tissue_data,
)
from dipy.io.image import load_nifti
from dipy.testing.decorators import warning_for_keywords
from dipy.tracking.streamline import relist_streamlines
from dipy.utils.arrfuncs import as_native_array

__all__ = [
    "fetch_30_bundle_atlas_hcp842",
    "fetch_bundle_atlas_hcp842",
    "fetch_bundle_fa_hcp",
    "fetch_bundle_warp_dataset",
    "fetch_bundles_2_subjects",
    "fetch_cenir_multib",
    "fetch_cfin_multib",
    "fetch_deepn4_test",
    "fetch_deepn4_tf_weights",
    "fetch_deepn4_torch_weights",
    "fetch_disco1_dataset",
    "fetch_disco2_dataset",
    "fetch_disco3_dataset",
    "fetch_disco_dataset",
    "fetch_evac_test",
    "fetch_evac_tf_weights",
    "fetch_evac_torch_weights",
    "fetch_gold_standard_io",
    "fetch_hbn",
    "fetch_hcp",
    "fetch_isbi2013_2shell",
    "fetch_ivim",
    "fetch_mni_template",
    "fetch_ptt_minimal_dataset",
    "fetch_resdnn_tf_weights",
    "fetch_resdnn_torch_weights",
    "fetch_scil_b0",
    "fetch_sherbrooke_3shell",
    "fetch_stanford_hardi",
    "fetch_stanford_labels",
    "fetch_stanford_pve_maps",
    "fetch_stanford_t1",
    "fetch_stanford_tracks",
    "fetch_syn_data",
    "fetch_synb0_test",
    "fetch_synb0_weights",
    "fetch_taiwan_ntu_dsi",
    "fetch_target_tractogram_hcp",
    "fetch_tissue_data",
    "get_bundle_atlas_hcp842",
    "get_fnames",
    "get_target_tractogram_hcp",
    "get_two_hcp842_bundles",
    "read_DiB_70_lte_pte_ste",
    "read_DiB_217_lte_pte_ste",
    "read_bundles_2_subjects",
    "read_cenir_multib",
    "read_cfin_dwi",
    "read_cfin_t1",
    "read_five_af_bundles",
    "read_isbi2013_2shell",
    "read_ivim",
    "read_mni_template",
    "read_qte_lte_pte",
    "read_scil_b0",
    "read_sherbrooke_3shell",
    "read_stanford_hardi",
    "read_stanford_labels",
    "read_stanford_pve_maps",
    "read_stanford_t1",
    "read_syn_data",
    "read_taiwan_ntu_dsi",
    "read_tissue_data",
]


def loads_compat(byte_data):
    return pickle.loads(byte_data, encoding="latin1")


DATA_DIR = pjoin(dirname(__file__), "files")
SPHERE_FILES = {
    "symmetric362": pjoin(DATA_DIR, "evenly_distributed_sphere_362.npz"),
    "symmetric642": pjoin(DATA_DIR, "evenly_distributed_sphere_642.npz"),
    "symmetric724": pjoin(DATA_DIR, "evenly_distributed_sphere_724.npz"),
    "repulsion724": pjoin(DATA_DIR, "repulsion724.npz"),
    "repulsion100": pjoin(DATA_DIR, "repulsion100.npz"),
    "repulsion200": pjoin(DATA_DIR, "repulsion200.npz"),
}


class DataError(Exception):
    pass


@warning_for_keywords()
def get_sim_voxels(*, name="fib1"):
    """provide some simulated voxel data

    Parameters
    ----------
    name : str, which file?
        'fib0', 'fib1' or 'fib2'

    Returns
    -------
    dix : dictionary, where dix['data'] returns a 2d array
        where every row is a simulated voxel with different orientation

    Examples
    --------
    >>> from dipy.data import get_sim_voxels
    >>> sv=get_sim_voxels(name='fib1')
    >>> sv['data'].shape == (100, 102)
    True
    >>> sv['fibres']
    '1'
    >>> sv['gradients'].shape == (102, 3)
    True
    >>> sv['bvals'].shape == (102,)
    True
    >>> sv['snr']
    '60'
    >>> sv2=get_sim_voxels(name='fib2')
    >>> sv2['fibres']
    '2'
    >>> sv2['snr']
    '80'

    Notes
    -----
    These sim voxels were provided by M.M. Correia using Rician noise.

    """
    if name == "fib0":
        fname = pjoin(DATA_DIR, "fib0.pkl.gz")
    if name == "fib1":
        fname = pjoin(DATA_DIR, "fib1.pkl.gz")
    if name == "fib2":
        fname = pjoin(DATA_DIR, "fib2.pkl.gz")
    return loads_compat(gzip.open(fname, "rb").read())


@warning_for_keywords()
def get_skeleton(*, name="C1"):
    """Provide skeletons generated from Local Skeleton Clustering (LSC).

    Parameters
    ----------
    name : str, 'C1' or 'C3'

    Returns
    -------
    dix : dictionary

    Examples
    --------
    >>> from dipy.data import get_skeleton
    >>> C=get_skeleton(name='C1')
    >>> len(C.keys())
    117
    >>> for c in C: break
    >>> sorted(C[c].keys())
    ['N', 'hidden', 'indices', 'most']

    """
    if name == "C1":
        fname = pjoin(DATA_DIR, "C1.pkl.gz")
    if name == "C3":
        fname = pjoin(DATA_DIR, "C3.pkl.gz")
    return loads_compat(gzip.open(fname, "rb").read())


@warning_for_keywords()
def get_sphere(*, name="symmetric362"):
    """provide triangulated spheres

    Parameters
    ----------
    name : str
        which sphere - one of:
        * 'symmetric362'
        * 'symmetric642'
        * 'symmetric724'
        * 'repulsion724'
        * 'repulsion100'
        * 'repulsion200'

    Returns
    -------
    sphere : a dipy.core.sphere.Sphere class instance

    Examples
    --------
    >>> import numpy as np
    >>> from dipy.data import get_sphere
    >>> sphere = get_sphere(name="symmetric362")
    >>> verts, faces = sphere.vertices, sphere.faces
    >>> verts.shape == (362, 3)
    True
    >>> faces.shape == (720, 3)
    True
    >>> verts, faces = get_sphere(name="not a sphere name") #doctest: +IGNORE_EXCEPTION_DETAIL
    Traceback (most recent call last):
        ...
    DataError: No sphere called "not a sphere name"

    """  # noqa: E501
    fname = SPHERE_FILES.get(name)
    if fname is None:
        raise DataError(f'No sphere called "{name}"')
    res = np.load(fname)
    # Set to native byte order to avoid errors in compiled routines for
    # big-endian platforms, when using these spheres.
    return Sphere(
        xyz=as_native_array(res["vertices"]), faces=as_native_array(res["faces"])
    )


default_sphere = HemiSphere.from_sphere(get_sphere(name="repulsion724"))
small_sphere = HemiSphere.from_sphere(get_sphere(name="symmetric362"))


def _gradient_from_file(filename):
    """Reads a gradient file saved as a text file compatible with np.loadtxt
    and saved in the dipy data directory"""

    def gtab_getter():
        gradfile = pjoin(DATA_DIR, filename)
        grad = np.loadtxt(gradfile, delimiter=",")
        gtab = GradientTable(grad)
        return gtab

    return gtab_getter


get_3shell_gtab = _gradient_from_file("gtab_3shell.txt")
get_isbi2013_2shell_gtab = _gradient_from_file("gtab_isbi2013_2shell.txt")
get_gtab_taiwan_dsi = _gradient_from_file("gtab_taiwan_dsi.txt")


def dsi_voxels():
    fimg, fbvals, fbvecs = get_fnames(name="small_101D")
    bvals = np.loadtxt(fbvals)
    bvecs = np.loadtxt(fbvecs).T
    data, _ = load_nifti(fimg)
    gtab = gradient_table(bvals, bvecs=bvecs)
    return data, gtab


def dsi_deconv_voxels():
    from dipy.sims.voxel import sticks_and_ball

    gtab = gradient_table(np.loadtxt(get_fnames(name="dsi515btable")))
    data = np.zeros((2, 2, 2, 515))
    for ix in range(2):
        for iy in range(2):
            for iz in range(2):
                data[ix, iy, iz], _ = sticks_and_ball(
                    gtab,
                    d=0.0015,
                    S0=1.0,
                    angles=[(0, 0), (90, 0)],
                    fractions=[50, 50],
                    snr=None,
                )
    return data, gtab


def mrtrix_spherical_functions():
    """Spherical functions represented by spherical harmonic coefficients and
    evaluated on a discrete sphere.

    Returns
    -------
    func_coef : array (2, 3, 4, 45)
        Functions represented by the coefficients associated with the
        mrtrix spherical harmonic basis of maximal order (l) 8.
    func_discrete : array (2, 3, 4, 81)
        Functions evaluated on `sphere`.
    sphere : Sphere
        The discrete sphere, points on the surface of a unit sphere, used to
        evaluate the functions.

    Notes
    -----
    These coefficients were obtained by using the dwi2SH command of mrtrix.

    """
    func_discrete, _ = load_nifti(pjoin(DATA_DIR, "func_discrete.nii.gz"))
    func_coef, _ = load_nifti(pjoin(DATA_DIR, "func_coef.nii.gz"))
    gradients = np.loadtxt(pjoin(DATA_DIR, "sphere_grad.txt"))
    # gradients[0] and the first volume of func_discrete,
    # func_discrete[..., 0], are associated with the b=0 signal.
    # gradients[:, 3] are the b-values for each gradient/volume.
    sphere = Sphere(xyz=gradients[1:, :3])
    return func_coef, func_discrete[..., 1:], sphere


dipy_cmaps = None


def get_cmap(name):
    """Make a callable, similar to maptlotlib.pyplot.get_cmap."""
    global dipy_cmaps
    if dipy_cmaps is None:
        filename = pjoin(DATA_DIR, "dipy_colormaps.json")
        with open(filename) as f:
            dipy_cmaps = json.load(f)

    desc = dipy_cmaps.get(name)
    if desc is None:
        return None

    def simple_cmap(v):
        """Emulates matplotlib colormap callable"""
        rgba = np.ones((len(v), 4))
        for i, color in enumerate(("red", "green", "blue")):
            x, y0, y1 = zip(*desc[color])
            # Matplotlib allows more complex colormaps, but for users who do
            # not have Matplotlib dipy makes a few simple colormaps available.
            # These colormaps are simple because y0 == y1, and therefore we
            # ignore y1 here.
            rgba[:, i] = np.interp(v, x, y0)
        return rgba

    return simple_cmap


def two_cingulum_bundles():
    fname = get_fnames(name="cb_2")
    res = np.load(fname)
    cb1 = relist_streamlines(res["points"], res["offsets"])
    cb2 = relist_streamlines(res["points2"], res["offsets2"])
    return cb1, cb2


def matlab_life_results():
    matlab_rmse = np.load(pjoin(DATA_DIR, "life_matlab_rmse.npy"))
    matlab_weights = np.load(pjoin(DATA_DIR, "life_matlab_weights.npy"))
    return matlab_rmse, matlab_weights


@warning_for_keywords()
def load_sdp_constraints(model_name, *, order=None):
    """Import semidefinite programming constraint matrices for different models.

    Generated as described for example in :footcite:p:`DelaHaije2020`.

    Parameters
    ----------
    model_name : string
        A string identifying the model that is to be constrained.
    order : unsigned int, optional
        A non-negative integer that represent the order or instance of the
        model.
        Default: None.

    Returns
    -------
    sdp_constraints : array
        Constraint matrices

    Notes
    -----
    The constraints will be loaded from a file named 'id_constraint_order.npz'.

    References
    ----------
    .. footbibliography::

    """

    file = model_name + "_constraint"
    if order is not None:
        file += f"_{str(order)}"
    if order is None:
        file += "_SC"
    file += ".npz"

    path = pjoin(DATA_DIR, file)

    if not exists(path):
        raise ValueError(f"Constraints file '{file}' not found.")

    try:
        array = load_npz(path)
        n, x = array.shape
        sdp_constraints = [array[i * x : (i + 1) * x] for i in range(n // x)]
        return sdp_constraints
    except Exception as e:
        raise ValueError(f"Failed to read constraints file '{file}'.") from e
