"""Utility functions for file formats"""

import logging
import numbers
import os

import nibabel as nib
from nibabel import Nifti1Image
from nibabel.streamlines import detect_format
import numpy as np
from trx import trx_file_memmap

import dipy
from dipy.testing.decorators import warning_for_keywords
from dipy.utils.optpkg import optional_package

pd, have_pd, _ = optional_package("pandas")

if have_pd:
    import pandas as pd


def nifti1_symmat(image_data, *args, **kwargs):
    """Returns a Nifti1Image with a symmetric matrix intent

    Parameters
    ----------
    image_data : array-like
        should have lower triangular elements of a symmetric matrix along the
        last dimension
    *args
        Passed to Nifti1Image
    **kwargs
        Passed to Nifti1Image

    Returns
    -------
    image : Nifti1Image
        5d, extra dimensions added before the last. Has symmetric matrix intent
        code

    """
    image_data = make5d(image_data)
    last_dim = image_data.shape[-1]
    n = (np.sqrt(1 + 8 * last_dim) - 1) / 2
    if (n % 1) != 0:
        raise ValueError("input_data does not seem to have matrix elements")

    image = Nifti1Image(image_data, *args, **kwargs)
    hdr = image.header
    hdr.set_intent("symmetric matrix", (n,))
    return image


def make5d(data):
    """reshapes the input to have 5 dimensions, adds extra dimensions just
    before the last dimension
    """
    data = np.asarray(data)
    if data.ndim > 5:
        raise ValueError("input is already more than 5d")
    shape = data.shape
    shape = shape[:-1] + (1,) * (5 - len(shape)) + shape[-1:]
    return data.reshape(shape)


@warning_for_keywords()
def decfa(img_orig, *, scale=False):
    """
    Create a nifti-compliant directional-encoded color FA image.

    Parameters
    ----------
    img_orig : Nifti1Image class instance.
        Contains encoding of the DEC FA image with a 4D volume of data, where
        the elements on the last dimension represent R, G and B components.

    scale: bool.
        Whether to scale the incoming data from the 0-1 to the 0-255 range
        expected in the output.

    Returns
    -------
    img : Nifti1Image class instance with dtype set to store tuples of
        uint8 in (R, G, B) order.


    Notes
    -----
    For a description of this format, see:

    https://nifti.nimh.nih.gov/nifti-1/documentation/nifti1fields/nifti1fields_pages/datatype.html
    """

    dest_dtype = np.dtype([("R", "uint8"), ("G", "uint8"), ("B", "uint8")])
    out_data = np.zeros(img_orig.shape[:3], dtype=dest_dtype)

    data_orig = np.asanyarray(img_orig.dataobj)

    if scale:
        data_orig = (data_orig * 255).astype("uint8")

    for ii in np.ndindex(img_orig.shape[:3]):
        val = data_orig[ii]
        out_data[ii] = (val[0], val[1], val[2])

    new_hdr = img_orig.header
    new_hdr["dim"][4] = 1
    new_hdr.set_intent(1001, name="Color FA")
    new_hdr.set_data_dtype(dest_dtype)

    return Nifti1Image(out_data, affine=img_orig.affine, header=new_hdr)


def decfa_to_float(img_orig):
    """
    Convert a nifti-compliant directional-encoded color FA image into a
    nifti image with RGB encoded in floating point resolution.

    Parameters
    ----------
    img_orig : Nifti1Image class instance.
        Contains encoding of the DEC FA image with a 3D volume of data, where
        each element is a (R, G, B) tuple in uint8.

    Returns
    -------
    img : Nifti1Image class instance with float dtype.

    Notes
    -----
    For a description of this format, see:

    https://nifti.nimh.nih.gov/nifti-1/documentation/nifti1fields/nifti1fields_pages/datatype.html
    """

    data_orig = np.asanyarray(img_orig.dataobj)
    out_data = np.zeros(data_orig.shape + (3,), dtype=np.uint8)

    for ii in np.ndindex(img_orig.shape[:3]):
        val = data_orig[ii]
        out_data[ii] = np.array([val[0], val[1], val[2]])

    new_hdr = img_orig.header
    new_hdr["dim"][4] = 3

    # Remove the original intent
    new_hdr.set_intent(0)
    new_hdr.set_data_dtype(float)

    return Nifti1Image(out_data, affine=img_orig.affine, header=new_hdr)


def is_reference_info_valid(affine, dimensions, voxel_sizes, voxel_order):
    """Validate basic data type and value of spatial attribute.

    Does not ensure that voxel_sizes and voxel_order are self-coherent with
    the affine.

    Only verifies the following:
        - affine is of the right type (float) and dimension (4,4)
        - affine contain values in the rotation part
        - dimensions is of right type (int) and length (3)
        - voxel_sizes is of right type (float) and length (3)
        - voxel_order is of right type (str) and length (3)

    The listed parameters are what is expected, provide something else and this
    function should fail (cover common mistakes).

    Parameters
    ----------
    affine: ndarray (4,4)
        Transformation of VOX to RASMM
    dimensions: ndarray (3,), int16
        Volume shape for each axis
    voxel_sizes:  ndarray (3,), float32
        Size of voxel for each axis
    voxel_order: string
        Typically 'RAS' or 'LPS'

    Returns
    -------
    output : bool
        Does the input represent a valid 'state' of spatial attribute

    """
    all_valid = True
    only_3d_warning = False

    if not affine.shape == (4, 4):
        all_valid = False
        logging.warning("Transformation matrix must be 4x4")

    if not affine[0:3, 0:3].any():
        all_valid = False
        logging.warning("Rotation matrix cannot be all zeros")

    if not len(dimensions) >= 3:
        all_valid = False
        only_3d_warning = True

    for i in dimensions:
        if not isinstance(i, numbers.Integral):
            all_valid = False
            logging.warning("Dimensions must be int.")
        if i <= 0:
            all_valid = False
            logging.warning("Dimensions must be above 0.")

    if not len(voxel_sizes) >= 3:
        all_valid = False
        only_3d_warning = True
    for i in voxel_sizes:
        if not isinstance(i, numbers.Number):
            all_valid = False
            logging.warning("Voxel size must be int/float.")
        if i <= 0:
            all_valid = False
            logging.warning("Voxel size must be above 0.")

    if not len(voxel_order) >= 3:
        all_valid = False
        only_3d_warning = True
    for i in voxel_order:
        if not isinstance(i, str):
            all_valid = False
            logging.warning("Voxel order must be string/char.")
        if i not in ["R", "A", "S", "L", "P", "I"]:
            all_valid = False
            logging.warning("Voxel order does not follow convention.")

    if only_3d_warning:
        logging.warning("Only 3D (and above) reference are considered valid.")

    return all_valid


def split_name_with_gz(filename):
    """
    Returns the clean basename and extension of a file.
    Means that this correctly manages the ".nii.gz" extensions.

    Parameters
    ----------
    filename: str
        The filename to clean

    Returns
    -------
        base, ext : tuple(str, str)
        Clean basename and the full extension
    """
    base, ext = os.path.splitext(filename)

    if ext.lower() == ".gz":
        # Test if we have a .nii additional extension
        temp_base, add_ext = os.path.splitext(base)

        if add_ext.lower() == ".nii" or add_ext.lower() == ".trk":
            ext = add_ext + ext
            base = temp_base

    return base, ext


def get_reference_info(reference):
    """Will compare the spatial attribute of 2 references.

    Parameters
    ----------
    reference : Nifti or Trk filename, Nifti1Image or TrkFile, Nifti1Header or
        trk.header (dict), TrxFile or trx.header (dict)
        Reference that provides the spatial attribute.

    Returns
    -------
    output : tuple
        - affine ndarray (4,4), np.float32, transformation of VOX to RASMM
        - dimensions ndarray (3,), int16, volume shape for each axis
        - voxel_sizes  ndarray (3,), float32, size of voxel for each axis
        - voxel_order, string, Typically 'RAS' or 'LPS'
    """

    is_nifti = False
    is_trk = False
    is_sft = False
    is_trx = False
    if isinstance(reference, str):
        _, ext = split_name_with_gz(reference)
        ext = ext.lower()
        if ext in [".nii", ".nii.gz"]:
            header = nib.load(reference).header
            is_nifti = True
        elif ext == ".trk":
            header = nib.streamlines.load(reference, lazy_load=True).header
            is_trk = True
        elif ext == ".trx":
            header = trx_file_memmap.load(reference).header
            is_trx = True
    elif isinstance(reference, trx_file_memmap.TrxFile):
        header = reference.header
        is_trx = True
    elif isinstance(reference, nib.nifti1.Nifti1Image):
        header = reference.header
        is_nifti = True
    elif isinstance(reference, nib.streamlines.trk.TrkFile):
        header = reference.header
        is_trk = True
    elif isinstance(reference, nib.nifti1.Nifti1Header):
        header = reference
        is_nifti = True
    elif isinstance(reference, dict) and "magic_number" in reference:
        header = reference
        is_trk = True
    elif isinstance(reference, dict) and "NB_VERTICES" in reference:
        header = reference
        is_trx = True
    elif isinstance(reference, dipy.io.stateful_tractogram.StatefulTractogram):
        is_sft = True

    if is_nifti:
        affine = header.get_best_affine()
        dimensions = header["dim"][1:4]
        voxel_sizes = header["pixdim"][1:4]

        if not affine[0:3, 0:3].any():
            raise ValueError(
                "Invalid affine, contains only zeros."
                "Cannot determine voxel order from transformation"
            )
        voxel_order = "".join(nib.aff2axcodes(affine))
    elif is_trk:
        affine = header["voxel_to_rasmm"]
        dimensions = header["dimensions"]
        voxel_sizes = header["voxel_sizes"]
        voxel_order = header["voxel_order"]
    elif is_sft:
        affine, dimensions, voxel_sizes, voxel_order = reference.space_attributes
    elif is_trx:
        affine = header["VOXEL_TO_RASMM"]
        dimensions = header["DIMENSIONS"]
        voxel_sizes = nib.affines.voxel_sizes(affine)
        voxel_order = "".join(nib.aff2axcodes(affine))
    else:
        raise TypeError("Input reference is not one of the supported format")

    if isinstance(voxel_order, np.bytes_):
        voxel_order = voxel_order.decode("utf-8")

    is_reference_info_valid(affine, dimensions, voxel_sizes, voxel_order)

    return affine.astype(np.float32), dimensions, voxel_sizes, voxel_order


def is_header_compatible(reference_1, reference_2):
    """Will compare the spatial attribute of 2 references

    Parameters
    ----------
    reference_1 : Nifti or Trk filename, Nifti1Image or TrkFile,
        Nifti1Header or trk.header (dict)
        Reference that provides the spatial attribute.
    reference_2 : Nifti or Trk filename, Nifti1Image or TrkFile,
        Nifti1Header or trk.header (dict)
        Reference that provides the spatial attribute.

    Returns
    -------
    output : bool
        Does all the spatial attribute match
    """

    affine_1, dimensions_1, voxel_sizes_1, voxel_order_1 = get_reference_info(
        reference_1
    )
    affine_2, dimensions_2, voxel_sizes_2, voxel_order_2 = get_reference_info(
        reference_2
    )

    identical_header = True
    if not np.allclose(affine_1, affine_2, rtol=1e-03, atol=1e-03):
        logging.error("Affine not equal")
        identical_header = False

    if not np.array_equal(dimensions_1, dimensions_2):
        logging.error("Dimensions not equal")
        identical_header = False

    if not np.allclose(voxel_sizes_1, voxel_sizes_2, rtol=1e-03, atol=1e-03):
        logging.error("Voxel_size not equal")
        identical_header = False

    if voxel_order_1 != voxel_order_2:
        logging.error("Voxel_order not equal")
        identical_header = False

    return identical_header


def create_tractogram_header(
    tractogram_type, affine, dimensions, voxel_sizes, voxel_order
):
    """Write a standard trk/tck header from spatial attribute"""
    if isinstance(tractogram_type, str):
        tractogram_type = detect_format(tractogram_type)

    new_header = tractogram_type.create_empty_header()
    new_header[nib.streamlines.Field.VOXEL_SIZES] = tuple(voxel_sizes)
    new_header[nib.streamlines.Field.DIMENSIONS] = tuple(dimensions)
    new_header[nib.streamlines.Field.VOXEL_TO_RASMM] = affine
    new_header[nib.streamlines.Field.VOXEL_ORDER] = voxel_order

    return new_header


def create_nifti_header(affine, dimensions, voxel_sizes):
    """Write a standard nifti header from spatial attribute"""
    new_header = nib.Nifti1Header()
    new_header.set_sform(affine)
    new_header["dim"][1:4] = dimensions
    new_header["pixdim"][1:4] = voxel_sizes

    new_header.affine = new_header.get_best_affine()

    return new_header


@warning_for_keywords()
def save_buan_profiles_hdf5(fname, dt, *, key=None):
    """Saves the given input dataframe to .h5 file

    Parameters
    ----------
    fname : string
        file name for saving the hdf5 file
    dt : Pandas DataFrame
        DataFrame to be saved as .h5 file
    key : str, optional
        Key to retrieve the contents in the HDF5 file. The file rootname will
        be used if not provided.

    """

    df = pd.DataFrame(dt)
    filename_hdf5 = fname + ".h5"

    if key is None:
        base_name_parts, _ = os.path.splitext(os.path.basename(fname))
        key = base_name_parts.split(".")[0]

    store = pd.HDFStore(filename_hdf5, complevel=9)
    store.append(key, df, data_columns=True, complevel=9)
    store.close()


@warning_for_keywords()
def read_img_arr_or_path(data, *, affine=None):
    """
    Helper function that handles inputs that can be paths, nifti img or arrays

    Parameters
    ----------
    data : array or nib.Nifti1Image or str.
        Either as a 3D/4D array or as a nifti image object, or as
        a string containing the full path to a nifti file.

    affine : 4x4 array, optional.
        Must be provided for `data` provided as an array. If provided together
        with Nifti1Image or str `data`, this input will over-ride the affine
        that is stored in the `data` input. Default: use the affine stored
        in `data`.

    Returns
    -------
    data, affine : ndarray and 4x4 array
    """
    if isinstance(data, np.ndarray) and affine is None:
        raise ValueError(
            "If data is provided as an array, an affine has ", "to be provided as well"
        )
    if isinstance(data, str):
        data = nib.load(data)
    if isinstance(data, nib.Nifti1Image):
        if affine is None:
            affine = data.affine
        data = data.get_fdata()
    return data, affine
