from __future__ import annotations

from datetime import datetime
from pathlib import Path
from typing import TYPE_CHECKING, Any, ClassVar, Literal

import itk
import itk.support.types as itkt
import numpy as np

from medio.metadata.affine import Affine
from medio.metadata.dcm_uid import generate_uid
from medio.metadata.itk_orientation import itk_orientation_code
from medio.metadata.metadata import MetaData, check_dcm_ornt
from medio.utils.files import is_dicom, make_dir, parse_series_uids

if TYPE_CHECKING:
    import os

    from numpy.typing import NDArray

    from medio.metadata.metadata import CoordSys


class ItkIO:
    coord_sys: ClassVar[Literal["itk"]] = "itk"
    DEFAULT_COMPONENTS_AXIS: int = 0  # in the transposed image
    # default image type:
    dimension: int = 3
    pixel_type = itk.ctype("short")  # signed short - int16
    image_type = itk.Image[pixel_type, dimension]

    @staticmethod
    def read_img(
        input_path: str | os.PathLike[str],
        desired_axcodes: str | tuple[str, ...] | None = None,
        header: bool = False,
        components_axis: int | None = None,
        pixel_type: itkt.PixelTypes | None = pixel_type,
        fallback_only: bool = True,
        series: str | int | None = None,
        private_tags: bool = False,
    ) -> tuple[NDArray[np.generic], MetaData[object]]:
        """
        The main reader function, reads images and performs reorientation and unpacking
        :param input_path: path of image file or directory containing dicom series
        :param desired_axcodes: string or tuple - e.g. 'LPI', ('R', 'A', 'S')
        :param header: whether to include a header attribute with additional metadata in the returned metadata
        :param components_axis: if not None and the image is channeled (e.g. RGB) move the channels to channels_axis
        :param pixel_type: preferred itk pixel type for the image
        :param fallback_only: if True, finds the pixel_type automatically and uses pixel_type only if failed
        :param series: str or int of the series to read (in the case of multiple series in a directory)
        :return: numpy image and metadata object which includes pixdim, affine, original orientation string and
        coordinates system
        """
        if header:
            # Currently, only DICOM will use imageio (GDCM). For NIfTI, we'll use ITK default reader.
            imageio = itk.GDCMImageIO.New()
            if private_tags:
                imageio.LoadPrivateTagsOn()
        else:
            imageio = None
        input_path = Path(input_path)
        if input_path.is_dir():
            # We assume that the directory contains dicom series, do imageio will work, if used.
            if imageio is not None:
                # fallback_only=True would skip the imageio, so disable it
                fallback_only = False
            img = ItkIO.read_dir(str(input_path), pixel_type, fallback_only, series, imageio)
        elif input_path.is_file():
            # If the input file is not a dicom (e.g. NIfTI), fallback to not use imageio.
            use_dicom_imageio = imageio is not None and imageio.CanReadFile(str(input_path))
            if use_dicom_imageio:
                # fallback_only=True would skip the imageio, so disable it
                fallback_only = False
            else:
                imageio = None
            img = ItkIO.read_img_file(str(input_path), pixel_type, fallback_only, imageio)
        else:
            raise FileNotFoundError(f'No such file or directory: "{input_path}"')

        affine = ItkIO.get_img_aff(img)
        metadata = MetaData(affine=affine, coord_sys=ItkIO.coord_sys)
        if (desired_axcodes is None) or (desired_axcodes == metadata.ornt):
            image_np = ItkIO.itk_img_to_array(img)
        else:
            orig_ornt = metadata.ornt
            img, _ = ItkIO.reorient(img, desired_axcodes)
            image_np, affine = ItkIO.unpack_img(img)
            metadata = MetaData(affine=affine, orig_ornt=orig_ornt, coord_sys=ItkIO.coord_sys)
        if header:
            metadict = imageio.GetMetaDataDictionary() if imageio else img.GetMetaDataDictionary()
            metadata.header = {key: metadict[key] for key in metadict.GetKeys() if not key.startswith("ITK_")}

        # TODO: consider unifying with PdcmIO.move_channels_axis
        n_components = img.GetNumberOfComponentsPerPixel()
        if (n_components > 1) and (components_axis is not None):
            # assert image_np.shape[ItkIO.DEFAULT_COMPONENTS_AXIS] == n_components
            image_np = np.moveaxis(image_np, ItkIO.DEFAULT_COMPONENTS_AXIS, components_axis)

        return image_np, metadata

    @staticmethod
    def read_meta(
        input_path: str | os.PathLike[str],
        desired_axcodes: str | tuple[str, ...] | None = None,
        header: bool = False,
        pixel_type: itkt.PixelTypes | None = pixel_type,
        fallback_only: bool = True,
        series: str | int | None = None,
        private_tags: bool = False,
    ) -> MetaData[object]:
        """
        Read only the metadata (affine, orientation, spatial shape) without loading pixel data.
        Uses ITK's UpdateOutputInformation() which processes only the file header.
        :param input_path: path of image file or directory containing dicom series
        :param desired_axcodes: optional orientation string, e.g. 'LPI'
        :param header: whether to include a header attribute with additional metadata
        :param pixel_type: preferred itk pixel type
        :param fallback_only: if True, auto-detect pixel type
        :param series: series to read when a directory has multiple series
        :param private_tags: if True, also load private DICOM tags (requires header=True)
        :return: MetaData with spatial_shape set
        """

        from medio.backends.nib_io import _reorient_affine
        from medio.metadata.convert_nib_itk import convert_affine, inv_axcodes

        input_path = Path(input_path)
        image_type = itk.Image[pixel_type, ItkIO.dimension]

        if input_path.is_dir():
            filenames = ItkIO.extract_series(str(input_path), series)
            reader = itk.ImageSeriesReader[image_type].New()
            if isinstance(filenames, str):
                reader.SetFileNames([filenames])
            else:
                reader.SetFileNames(filenames)
            if header:
                imageio = itk.GDCMImageIO.New()
                if private_tags:
                    imageio.LoadPrivateTagsOn()
                reader.SetImageIO(imageio)
            else:
                imageio = None
            reader.UpdateOutputInformation()
            img = reader.GetOutput()
        elif input_path.is_file():
            reader = itk.ImageFileReader[image_type].New()
            reader.SetFileName(str(input_path))
            if header and is_dicom(str(input_path), check_exist=False):
                imageio = itk.GDCMImageIO.New()
                if private_tags:
                    imageio.LoadPrivateTagsOn()
                reader.SetImageIO(imageio)
            else:
                imageio = None
            reader.UpdateOutputInformation()
            img = reader.GetOutput()
        else:
            raise FileNotFoundError(f'No such file or directory: "{input_path}"')

        affine = ItkIO.get_img_aff(img)

        # Spatial shape in ITK is z, y, x from GetSize(); after the .T transpose in itk_img_to_array it becomes x, y, z
        itk_size = img.GetLargestPossibleRegion().GetSize()
        spatial_shape: tuple[int, ...] = tuple(itk_size[i] for i in range(ItkIO.dimension))

        metadata: MetaData[object] = MetaData(affine=affine, coord_sys=ItkIO.coord_sys, spatial_shape=spatial_shape)

        if desired_axcodes is not None and desired_axcodes != metadata.ornt:
            orig_ornt = metadata.ornt
            # Reorient using pure affine math (no pixel data)
            # Convert to nib convention for _reorient_affine
            nib_affine = convert_affine(affine)
            nib_desired = (
                inv_axcodes(desired_axcodes)
                if isinstance(desired_axcodes, str)
                else tuple(inv_axcodes(c) for c in desired_axcodes)
            )
            new_nib_affine, new_shape = _reorient_affine(nib_affine, spatial_shape, nib_desired)
            new_affine = convert_affine(new_nib_affine)
            metadata = MetaData(
                affine=new_affine, orig_ornt=orig_ornt, coord_sys=ItkIO.coord_sys, spatial_shape=new_shape
            )

        if header:
            try:
                metadict = imageio.GetMetaDataDictionary() if imageio else img.GetMetaDataDictionary()
                metadata.header = {key: metadict[key] for key in metadict.GetKeys() if not key.startswith("ITK_")}
            except Exception:
                metadata.header = {}

        return metadata

    @staticmethod
    def save_img(
        filename: str | os.PathLike[str],
        image_np: NDArray[np.generic],
        metadata: MetaData[object],
        use_original_ornt: bool = True,
        components_axis: int | None = None,
        allow_dcm_reorient: bool = False,
        compression: bool = False,
    ) -> None:
        """
        Save an image file with itk
        :param filename: the filename to save, str or os.PathLike
        :param image_np: the image's numpy array
        :param metadata: the corresponding metadata
        :param use_original_ornt: whether to save in the original orientation or not
        :param components_axis: if not None - the image has more than 1 component (e.g. RGB) and the components are in
        components_axis
        :param allow_dcm_reorient: whether to allow automatic reorientation to a right handed orientation or not
        :param compression: use compression or not
        """
        is_dcm = is_dicom(filename, check_exist=False)
        if is_dcm:
            image_np = ItkIO.prepare_dcm_array(image_np, is_vector=components_axis is not None)
        image = ItkIO.prepare_image(
            image_np,
            metadata,
            use_original_ornt,
            components_axis=components_axis,
            is_dcm=is_dcm,
            allow_dcm_reorient=allow_dcm_reorient,
        )
        ItkIO.save_img_file(image, str(filename), compression=compression)

    @staticmethod
    def prepare_image(
        image_np: NDArray[np.generic],
        metadata: MetaData[object],
        use_original_ornt: bool,
        components_axis: int | None = None,
        is_dcm: bool = False,
        allow_dcm_reorient: bool = False,
    ) -> object:
        """Prepare image for saving"""
        orig_coord_sys = metadata.coord_sys
        metadata.convert(ItkIO.coord_sys)
        desired_ornt = metadata.orig_ornt if use_original_ornt else None
        if is_dcm:
            # checking right-handed orientation before saving a dicom file/series
            desired_ornt = check_dcm_ornt(desired_ornt, metadata, allow_dcm_reorient=allow_dcm_reorient)
        image = ItkIO.pack2img(image_np, metadata.affine, components_axis=components_axis)
        if (desired_ornt is not None) and (desired_ornt != metadata.ornt):
            image, _ = ItkIO.reorient(image, desired_ornt)
        metadata.convert(orig_coord_sys)
        return image

    @staticmethod
    def prepare_dcm_array(image_np: NDArray[np.generic], is_vector: bool = False) -> NDArray[np.generic]:
        """Change image_np to correct data type for saving a single dicom file"""
        if is_vector:
            dcm_dtypes = [np.uint8]
        else:
            # for 3d image the supported data types are:
            dcm_dtypes = [np.uint8, np.uint16]
            # if the image is 2d it can be signed
            if np.squeeze(image_np).ndim == 2:
                dcm_dtypes = [np.int16, *dcm_dtypes]

        if image_np.dtype in dcm_dtypes:
            return image_np

        for dtype in dcm_dtypes:
            arr = image_np.astype(dtype, copy=False)
            if np.array_equal(arr, image_np):
                return arr

        raise NotImplementedError(
            "Saving a single dicom file with ItkIO is currently supported only for \n"
            "1. 2d images - int16, uint16, uint8\n"
            "2. 3d images with integer nonnegative values - uint8, uint16\n"
            "3. 2d/3d RGB[A] images - uint8 (with channels_axis)\n"
            "For negative values, try to save a dicom directory or use PdcmIO.save_arr2dcm_file"
        )

    @staticmethod
    def read_img_file(
        filename: str,
        pixel_type: itkt.PixelTypes | None = None,
        fallback_only: bool = False,
        imageio: object | None = None,
    ) -> object:
        """Common pixel types: itk.SS (int16), itk.US (uint16), itk.UC (uint8)"""
        return itk.imread(filename, pixel_type, fallback_only, imageio)

    @staticmethod
    def read_img_file_long(filename: str, image_type: type = image_type) -> tuple[object, str]:
        """Longer version of itk.imread that returns the itk image and io engine string"""
        reader = itk.ImageFileReader[image_type].New()
        reader.LoadPrivateTagsOn()
        reader.SetFileName(filename)
        reader.Update()
        image_io = str(reader.GetImageIO()).split(" ")[0]
        image = reader.GetOutput()
        return image, image_io

    @staticmethod
    def save_img_file(image: object, filename: str, compression: bool = False) -> None:
        itk.imwrite(image, filename, compression)

    @staticmethod
    def save_img_file_long(image: object, filename: str, compression: bool = False) -> None:
        image_type = type(image)
        writer = itk.ImageFileWriter[image_type].New()
        if compression:
            writer.UseCompressionOn()
        writer.UseInputMetaDataDictionaryOn()
        writer.SetFileName(filename)
        writer.SetInput(image)
        writer.Update()

    @staticmethod
    def itk_img_to_array(img_itk: object) -> NDArray[np.generic]:
        """
        Swap the axes to the usual x, y, z convention in RAI orientation
        (originally z, y, x)
        """
        # the transpose here is equivalent to keep_axes=True
        img_array = itk.array_from_image(img_itk).T
        return img_array

    @staticmethod
    def array_to_itk_img(img_array: NDArray[np.generic], components_axis: int | None = None) -> object:
        """Set components_axis to not None for vector images, e.g. RGB"""
        is_vector = False
        if components_axis is not None:
            img_array = np.moveaxis(img_array, components_axis, ItkIO.DEFAULT_COMPONENTS_AXIS)
            is_vector = True
        # copy is crucial for the ordering
        img_itk = itk.image_from_array(img_array.T.copy(), is_vector=is_vector)
        return img_itk

    @staticmethod
    def unpack_img(img: object) -> tuple[NDArray[np.generic], Affine]:
        image_np = ItkIO.itk_img_to_array(img)
        affine = ItkIO.get_img_aff(img)
        return image_np, affine

    @staticmethod
    def from_itk_img(
        itk_img: object,
        desired_ornt: str | None = None,
        coord_sys: CoordSys | None = "itk",
        components_axis: int | None = None,
    ) -> tuple[NDArray[np.generic], MetaData[object]]:
        """
        Convert an ITK image object to a medio (numpy array, MetaData) pair.
        :param itk_img: an itk.Image object
        :param desired_ornt: optional orientation for reorientation in the convention of `coord_sys`.
            When coord_sys is None, must be in ITK convention.
        :param coord_sys: convention of desired_ornt and returned MetaData ('itk', 'nib', or None).
            None means no conversion is applied (ITK convention is returned).
        :param components_axis: for multi-component images (e.g. RGB), move components to this axis
        :return: (np_image, metadata) — same format as read_img
        """
        from medio.metadata.convert_nib_itk import inv_axcodes

        itk_desired_ornt: str | None = desired_ornt
        if (coord_sys is not None) and (coord_sys != ItkIO.coord_sys):
            itk_desired_ornt = inv_axcodes(desired_ornt)

        affine = ItkIO.get_img_aff(itk_img)
        metadata: MetaData[object] = MetaData(affine=affine, coord_sys=ItkIO.coord_sys)
        if (itk_desired_ornt is None) or (itk_desired_ornt == metadata.ornt):
            img = itk_img
            image_np = ItkIO.itk_img_to_array(img)
        else:
            orig_ornt = metadata.ornt
            img, _ = ItkIO.reorient(itk_img, itk_desired_ornt)
            image_np, affine = ItkIO.unpack_img(img)
            metadata = MetaData(affine=affine, orig_ornt=orig_ornt, coord_sys=ItkIO.coord_sys)

        n_components = img.GetNumberOfComponentsPerPixel()
        if (n_components > 1) and (components_axis is not None):
            image_np = np.moveaxis(image_np, ItkIO.DEFAULT_COMPONENTS_AXIS, components_axis)

        if coord_sys is not None:
            metadata.convert(coord_sys)
        return image_np, metadata

    @staticmethod
    def to_itk_img(
        np_image: NDArray[np.generic],
        metadata: MetaData[object],
        components_axis: int | None = None,
    ) -> object:
        """
        Convert a medio (numpy array, MetaData) pair to an ITK image object.
        Handles any coord_sys automatically; the caller's MetaData is never mutated.
        :param np_image: the image numpy array
        :param metadata: the MetaData object (any coord_sys)
        :param components_axis: axis in np_image holding image components (e.g. RGB channels)
        :return: itk.Image with spatial metadata applied
        """
        metadata_itk = metadata.clone()
        metadata_itk.convert(ItkIO.coord_sys)
        return ItkIO.pack2img(np_image, metadata_itk.affine, components_axis=components_axis)

    @staticmethod
    def get_img_aff(img: object) -> Affine:
        direction = itk.array_from_vnl_matrix(img.GetDirection().GetVnlMatrix().as_matrix())
        spacing = itk.array_from_vnl_vector(img.GetSpacing().GetVnlVector())
        origin = itk.array_from_vnl_vector(img.GetOrigin().GetVnlVector())
        return Affine(direction=direction, spacing=spacing, origin=origin)

    @staticmethod
    def pack2img(
        image_np: NDArray[np.generic],
        affine: Affine | NDArray[np.floating],
        components_axis: int | None = None,
    ) -> object:
        image = ItkIO.array_to_itk_img(image_np, components_axis)
        ItkIO.set_img_aff(image, affine)
        return image

    @staticmethod
    def set_img_aff(image: object, affine: Affine | NDArray[np.floating]) -> None:
        if not isinstance(affine, Affine):
            affine = Affine(affine)

        dimension = image.GetImageDimension()
        direction_arr, spacing, origin = affine.direction, affine.spacing, affine.origin

        # setting metadata
        spacing_vec = itk.Vector[itk.D, dimension]()
        spacing_vec.SetVnlVector(itk.vnl_vector_from_array(spacing.astype("float")))
        image.SetSpacing(spacing_vec)

        image.SetOrigin(origin.astype("float"))

        direction_mat = itk.vnl_matrix_from_array(direction_arr.astype("float"))
        direction = itk.Matrix[itk.D, dimension, dimension](direction_mat)
        image.SetDirection(direction)

    @staticmethod
    def reorient(img: object, desired_orientation: int | tuple[str, ...] | str | None) -> tuple[object, int | None]:
        if desired_orientation is None:
            return img, None

        image_type = type(img)
        orient = itk.OrientImageFilter[image_type, image_type].New()
        orient.UseImageDirectionOn()
        orient.SetInput(img)
        if isinstance(desired_orientation, (str, tuple)):
            desired_orientation = itk_orientation_code(desired_orientation)
        orient.SetDesiredCoordinateOrientation(desired_orientation)
        orient.Update()
        reoriented_itk_img = orient.GetOutput()
        original_orientation_code = orient.GetGivenCoordinateOrientation()
        return reoriented_itk_img, original_orientation_code

    @staticmethod
    def read_dir(
        dirname: str,
        pixel_type: itk.itkCType | None = None,
        fallback_only: bool = False,
        series: str | int | None = None,
        imageio: object | None = None,
    ) -> object:
        """
        Read a dicom directory. If there is more than one series in the directory an error is raised
        (unless the series argument is used properly).
        Shorter option for a single series (provided the slices order is known):
        >>> itk.imread([filename0, filename1, ...])
        """
        filenames = ItkIO.extract_series(dirname, series)
        return itk.imread(filenames, pixel_type, fallback_only, imageio)

    @staticmethod
    def extract_series(dirname: str, series: str | int | None = None) -> list[str] | str:
        """Extract series filenames from the directory dirname"""
        names_generator = itk.GDCMSeriesFileNames.New()
        names_generator.SetDirectory(dirname)

        series_uids = names_generator.GetSeriesUIDs()
        series_uid = parse_series_uids(dirname, series_uids, series)

        filenames = names_generator.GetFileNames(series_uid)
        if len(filenames) == 1:
            filenames = filenames[0]  # there is a single image in the series

        return filenames

    @staticmethod
    def save_dcm_dir(
        dirname: str | os.PathLike[str],
        image_np: NDArray[np.generic],
        metadata: MetaData[object],
        use_original_ornt: bool = True,
        components_axis: int | None = None,
        parents: bool = False,
        exist_ok: bool = False,
        allow_dcm_reorient: bool = False,
        **kwargs: Any,
    ) -> None:
        """
        Save a 3d numpy array image_np as a dicom series of 2d dicom slices in the directory dirname
        :param dirname: the directory to save in the files, str or pathlib.Path. If it exists - must be empty
        :param image_np: the image's numpy array
        :param metadata: the corresponding metadata
        :param use_original_ornt: whether to save in the original orientation or not
        :param components_axis: if not None - the image has more than 1 component (e.g. RGB) and the components are in
        components_axis
        :param parents: if True, creates also the parents of dirname
        :param exist_ok: if True, non-empty existing directory will not raise an error
        :param allow_dcm_reorient: whether to allow automatic reorientation to a right-handed orientation or not
        :param kwargs: optional kwargs passed to ItkIO.dcm_metadata: pattern, metadata_dict
        """
        image = ItkIO.prepare_image(
            image_np,
            metadata,
            use_original_ornt,
            components_axis=components_axis,
            is_dcm=True,
            allow_dcm_reorient=allow_dcm_reorient,
        )
        image_type = type(image)
        _, (pixel_type, _) = itk.template(image)
        image2d_type = itk.Image[pixel_type, 2]
        writer = itk.ImageSeriesWriter[image_type, image2d_type].New()
        make_dir(dirname, parents, exist_ok)
        # Generate necessary metadata and filenames per slice:
        mdict_list, filenames = ItkIO.dcm_series_metadata(image, dirname, **kwargs)
        metadict_vec = itk.vector[itk.MetaDataDictionary](mdict_list)
        writer.SetMetaDataDictionaryArray(metadict_vec)
        writer.SetFileNames(filenames)
        dicom_io = itk.GDCMImageIO.New()
        dicom_io.KeepOriginalUIDOn()
        writer.SetImageIO(dicom_io)
        writer.SetInput(image)
        writer.Update()

    @staticmethod
    def dcm_series_metadata(
        image: object,
        dirname: str | os.PathLike[str],
        pattern: str = "IM{}.dcm",
        metadata_dict: dict[str, str] | None = None,
    ) -> tuple[list[object], list[str]]:
        """
        Return dicom series metadata per slice and filenames
        :param image: the full itk image to be saved as dicom series
        :param dirname: the directory name
        :param pattern: str pattern for the filenames to save, including a placeholder ('{}') for the slice number
        :param metadata_dict: dictionary of metadata for adding tags or overriding the default values. For example,
        metadata_dict={'0008|0060': 'US'} will override the default 'CT' modality and set it to 'US' (ultrasound)
        :return: metadata dictionaries per slice, slice filenames
        """
        # The number of slices
        n = image.GetLargestPossibleRegion().GetSize().GetElement(2)

        # Shared properties for all the n slices:
        mdict = itk.MetaDataDictionary()

        # Series Instance UID
        mdict["0020|000e"] = generate_uid()
        # Study Instance UID
        mdict["0020|000d"] = generate_uid()

        date, time = datetime.now().strftime("%Y%m%d %H%M%S.%f").split()
        # Study Date
        mdict["0008|0020"] = date
        # Series Date
        mdict["0008|0021"] = date
        # Content Date
        mdict["0008|0023"] = date
        # Study Time
        mdict["0008|0030"] = time
        # Series Time
        mdict["0008|0031"] = time

        # Pixel Spacing (not necessary - automatically saved)
        spacing = image.GetSpacing()
        mdict["0028|0030"] = f"{spacing[0]}\\{spacing[1]}"
        # Spacing Between Slices
        mdict["0018|0088"] = str(spacing[2])
        # Image Orientation (Patient)
        orientation_str = "\\".join(
            [str(image.GetDirection().GetVnlMatrix().get(i, j)) for j in range(2) for i in range(3)]
        )
        mdict["0020|0037"] = orientation_str
        # Patient Position
        mdict["0018|5100"] = ""

        # Number of Frames
        mdict["0028|0008"] = "1"
        # Number of Slices
        mdict["0054|0081"] = str(n)
        # Modality
        mdict["0008|0060"] = "CT"

        if metadata_dict is not None:
            for key, val in metadata_dict.items():
                mdict[key] = val

        # Per slice properties:
        mdict_list = []
        filenames = []

        for i in range(n):
            # copy the shared properties dict:
            mdict_i = itk.MetaDataDictionary(mdict)
            # Instance Number
            mdict_i["0020|0013"] = str(i + 1)
            # Image Position (Patient)
            position = image.TransformIndexToPhysicalPoint([0, 0, i])
            position_str = "\\".join([str(position[i]) for i in range(3)])
            mdict_i["0020|0032"] = position_str

            mdict_list += [mdict_i]
            filenames += [str(Path(dirname) / pattern.format(i + 1))]

        return mdict_list, filenames
