from typing import Union

import nibabel as nib
import numpy as np

from medio.metadata.affine import Affine
from medio.metadata.metadata import MetaData


class NibIO:
    coord_sys = 'nib'
    RGB_DTYPE = np.dtype([('R', np.uint8),
                          ('G', np.uint8),
                          ('B', np.uint8)])
    RGBA_DTYPE = np.dtype([('R', np.uint8),
                           ('G', np.uint8),
                           ('B', np.uint8),
                           ('A', np.uint8)])

    @staticmethod
    def read_img(input_path, desired_axcodes=None, header=False, channels_axis=None):
        """
        Reads a NIFTI file and returns the image array and metadata
        :param input_path: path-like (str or pathlib.Path) of the image file to read
        :param desired_axcodes: str, tuple of str or None - the desired orientation of the image to be returned
        :param channels_axis: if not None and the array dtype is structured, stacks the channels along channels_axis
        :param header: whether to include a header attribute with additional metadata in the returned metadata
        :return: image array and corresponding metadata
        """
        img_struct = nib.load(input_path)
        orig_ornt_str = ''.join(nib.aff2axcodes(img_struct.affine))
        if desired_axcodes is not None:
            img_struct = NibIO.reorient(img_struct, desired_axcodes)
        img = np.asanyarray(img_struct.dataobj)
        if channels_axis is not None:
            img = NibIO.unravel_array(img, channels_axis)
        affine = Affine(img_struct.affine)
        metadata = MetaData(affine=affine, orig_ornt=orig_ornt_str, coord_sys=NibIO.coord_sys)
        if header:
            metadata.header = {key: img_struct.header[key] for key in img_struct.header.keys()}
        return img, metadata

    @staticmethod
    def save_img(filename, img, metadata, use_original_ornt=True, channels_axis=None):
        """
        Saves the given image as a NIFTI file.
        :param filename: path-like output filename, including a '.nii.gz' or '.nii' suffix
        :param img: image data array
        :param metadata: the matching metadata
        :param use_original_ornt: whether to use the original orientation of the image of not
        :param channels_axis: if not None gives the channels axis of img (for channeled images RGB/RGBA)
        """
        if channels_axis is not None:
            img = NibIO.pack_channeled_img(img, channels_axis)
        orig_coord_sys = metadata.coord_sys
        metadata.convert(NibIO.coord_sys)
        img_struct = nib.Nifti1Image(img, metadata.affine)
        desired_axcodes = metadata.orig_ornt if use_original_ornt else None
        metadata.convert(orig_coord_sys)
        img_struct = NibIO.reorient(img_struct, desired_axcodes)
        nib.save(img_struct, filename)

    @staticmethod
    def reorient(img_struct, desired_axcodes: Union[tuple, str, None]):
        """Reorient a nibabel image to a desired orientation described by desired_axcodes strings tuple, for example
        ('L', 'P', 'I'). If desired_axcodes is None it returns the given img_struct"""
        if desired_axcodes is not None:
            if isinstance(desired_axcodes, str):
                desired_axcodes = tuple(desired_axcodes)
            start_ornt = nib.orientations.io_orientation(img_struct.affine)
            end_ornt = nib.orientations.axcodes2ornt(desired_axcodes)
            ornt_tform = nib.orientations.ornt_transform(start_ornt, end_ornt)
            img_struct = img_struct.as_reoriented(ornt_tform)
        return img_struct

    @staticmethod
    def unravel_array(array, channels_axis=-1):
        """Simplify array dtype if it is a structured data type. For example, if the array if of RGB dtype:
        np.dtype([('R', 'uint8'), ('G', 'uint8'), ('B', 'uint8')])
        Convert it into an array with dtype 'uint8' and 3 channels for RGB in an additional last dimension"""
        dtype = array.dtype
        if not (hasattr(dtype, '__len__') and len(dtype) > 1):
            return array
        return np.stack([array[field] for field in dtype.fields], axis=channels_axis)

    @staticmethod
    def pack_channeled_img(img, channels_axis):
        dtype = img.dtype
        if not np.issubdtype(dtype, np.uint8):
            raise ValueError(f'RGB or RGBA images must have dtype "np.uint8", got: "{dtype}"')
        n_channels = img.shape[channels_axis]
        img = np.moveaxis(img, channels_axis, -1)
        r_channel = img[..., 0]
        if n_channels == 3:
            img_rgb = np.empty_like(r_channel, dtype=NibIO.RGB_DTYPE)
            img_rgb['R'] = r_channel
            img_rgb['G'] = img[..., 1]
            img_rgb['B'] = img[..., 2]
            return img_rgb
        elif n_channels == 4:
            img_rgba = np.empty_like(r_channel, dtype=NibIO.RGBA_DTYPE)
            img_rgba['R'] = r_channel
            img_rgba['G'] = img[..., 1]
            img_rgba['B'] = img[..., 2]
            img_rgba['A'] = img[..., 3]
            return img_rgba
        else:
            raise ValueError(f'Invalid number of channels: {n_channels}, should be 3 (RGB) or 4 (RGBA)')
