"""
fs.py
This is a file that provides methods to interface with image files.

Stored Numpy arrays in computer vision can be represented in .npy, .tif, .png, .jpg etc.
And when reading and writing files, the axis order and pixel format (grey vs. RGB vs RGBA)
often complicate the code.

This file aims to provide a unified way to manipulate, read and write np arrays based on a
file path pattern e.g. "./dataset1/dog_*.png" and a few settings specifying the reading
axis order and underlying data arrangement formats
"""


import os
import shutil
import glob
import numpy as np
import nibabel as nib
from PIL import Image
import scipy
from dataclasses import dataclass, field
from cvpl_tools.array_key_dict import ArrayKeyDict
from nibabel import load as nib_load


def ensure_dir_exists(dir_path, remove_if_already_exists):
    """
    If a directory does not exist, make a new directory with the name.
    This assumes the parent directory must exists; otherwise a path not
    found error will be thrown.
    Args:
        dir_path: The path of folder
        remove_if_already_exists: if True and the folder already exists, then remove it and make a new one.
    """
    if os.path.exists(dir_path):
        if remove_if_already_exists:
            shutil.rmtree(dir_path)
            os.mkdir(dir_path)
    else:
        os.mkdir(dir_path)


class ImFileType:
    """
    Represents different file types that can be read and written by fs.py
    """
    FTYPE_UNKNOWN = -1
    FTYPE_JPX = 0
    FTYPE_JPG = 1
    FTYPE_PNG = 2
    FTYPE_NIB = 3
    FTYPE_TIF = 4
    FTYPE_NPY = 5
    FTYPE_NPZ = 6  # probabilities output by nnunet
    ext_to_ftype = {
        '.jpx': FTYPE_JPX,
        '.jpg': FTYPE_JPG,
        '.png': FTYPE_PNG,
        '.nii.gz': FTYPE_NIB,
        '.tif': FTYPE_TIF,
        '.npy': FTYPE_NPY,
        '.npz': FTYPE_NPZ,
    }
    ftype_to_ext = list(ext_to_ftype.keys())
    all_ftypes = list(ext_to_ftype.values())  # except unknown

    FORMAT_UINT8 = 0
    FORMAT_UINT16 = 1
    FORMAT_INT32 = 2
    FORMAT_FLOAT32 = 3
    format_to_np_dtype = [np.uint8, np.uint16, np.int32, np.float32]

    @staticmethod
    def ftype_from_im_path(im_path: str) -> int:
        """
        Infer the ftype from image path;
        The path can contain wildcards like "*"
        When the ftype can't be inferred, FTYPE_UNKNOWN will be returned
        """
        prefix, ext = os.path.splitext(im_path)
        ext = ext.lower()
        if ext == '.gz':
            prefix_prefix, prefix_ext = os.path.splitext(prefix)
            if prefix_ext.lower() == '.nii':
                return ImFileType.FTYPE_NIB
            else:
                return ImFileType.FTYPE_UNKNOWN
        elif ext in ImFileType.ext_to_ftype:
            return ImFileType.ext_to_ftype[ext]
        else:
            return ImFileType.FTYPE_UNKNOWN

    @staticmethod
    def get_imid_and_zstack(im_path: str, is_stack=False, keep_dirpath=False) -> tuple[str, str]:
        """
        Get the path stem as the unique image id from the path of the file;
        Second return value is the zstack number as a string
        file extensions will be stripped;
        If the file is part of a stacked image, then the z-stack numbers will be stripped from stem as will
        """
        ftype = ImFileType.ftype_from_im_path(im_path)
        if ftype == ImFileType.FTYPE_UNKNOWN:
            raise ValueError(f'Attempt to get the imid of an image file from path "{im_path}" with unknown file type!')
        elif ftype == ImFileType.FTYPE_NIB:
            imid = im_path[:-len('.nii.gz')]
        else:
            imid = os.path.splitext(im_path)[0]

        if not keep_dirpath:
            imid = os.path.split(imid)[1]

        if is_stack:
            i = -1
            while len(imid) >= -i and imid[i].isdigit():
                i -= 1
            return imid[:i + 1], imid[i + 1:]
        else:
            return imid, ''

    @staticmethod
    def get_imid(im_path: str, is_stack=False, keep_dirpath=False) -> str:
        return ImFileType.get_imid_and_zstack(im_path, is_stack, keep_dirpath)[0]

    @staticmethod
    def load_im_as_np(im_format, path, true_im_ndim: int, stack_axis=None,
                      concat_instead=False, allow_pickle=False) -> tuple:
        """
        Load individual file as a np array, or a stack of images as a np array if a list of path is given
        Returns a tuple (im, im_ftype, im_meta) where im_meta may be needed in order to save the im back to a file
        for certain file types
        """
        if isinstance(path, str):
            ftype = ImFileType.ftype_from_im_path(path)
            assert stack_axis is None, 'ERROR: Passing stack_axis in when loading single file'
            assert ftype != ImFileType.FTYPE_UNKNOWN, 'ERROR: Attempt to load an unknown image type!'
            if ftype == ImFileType.FTYPE_NPZ:
                assert im_format == ImFileType.FORMAT_FLOAT32
                im = np.load(path, allow_pickle=allow_pickle)['probabilities']
                im = np.transpose(im, (0, 3, 2, 1))
                im_meta = [ftype, im.shape, None, path, None]
            elif ftype == ImFileType.FTYPE_NPY:
                im = np.load(path, allow_pickle=allow_pickle)
                if allow_pickle:
                    im = im.item()['masks']
                im_meta = [ftype, im.shape, None, path, None]
            elif ftype == ImFileType.FTYPE_NIB:
                nii_im = nib_load(path)
                im = np.array(nii_im.get_fdata(), dtype=ImFileType.format_to_np_dtype[im_format])
                im_meta = [ftype, im.shape, None, path, nii_im.affine]
            else:
                im = np.array(Image.open(path), dtype=ImFileType.format_to_np_dtype[im_format])
                im_meta = [ftype, im.shape, None, path, None]

            assert im.dtype == ImFileType.format_to_np_dtype[im_format], \
                f'ERROR: Expect file type {ImFileType.format_to_np_dtype[im_format]} but got an image of dtype {im.dtype} at {path}'
            if len(im.shape) > true_im_ndim:
                axes = range(true_im_ndim, len(im.shape))
                for i in axes:
                    assert im.shape[i] == 1, f'Extra axes to true_im_ndim must be of width 1. true_im_ndim={true_im_ndim}, im.shape={im.shape}'
                im = im.squeeze(axis=tuple(axes))
            assert len(im.shape) == true_im_ndim, (f'ERROR: Expected true dimension {true_im_ndim}, but found '
                                                  f'image of shape {im.shape}!')
            assert type(im_meta) is list, 'ERROR: im_meta is not properly initialized!'
            return im, im_meta
        else:
            assert isinstance(path, list)
            im_arr = []
            im_meta = None
            for filepath in path:
                if concat_instead:
                    true_slice_dim = true_im_ndim
                else:
                    true_slice_dim = true_im_ndim - 1
                tup = ImFileType.load_im_as_np(im_format, filepath, true_slice_dim)
                im_arr.append(tup[0])
                cur_im_meta = tup[1]
                if im_meta is None:
                    im_meta = list(cur_im_meta)  # shallow copy
                    ftype_meta = im_meta[-1]
                    im_meta[-1] = [ftype_meta]
                else:
                    assert im_meta[0] == cur_im_meta[0]  # same file type
                    assert im_meta[1] == cur_im_meta[1]  # each image has the same shape
                    im_meta[-1].append(cur_im_meta[-1])
            if stack_axis is None:
                raise ValueError(f"Please specify stack_axis!")
                # stack_axis = true_im_ndim - 1  # default stack on the last axis
            if concat_instead:  # concat on existing axis
                assert 0 <= stack_axis < true_im_ndim - 1, f'Attempt to concat on axis {stack_axis} with true dim={true_im_ndim}'
                im = np.concatenate(im_arr, axis=stack_axis)
            else:  # stack on new axis
                assert 0 <= stack_axis < true_im_ndim, f'Attempt to concat on axis {stack_axis} with true dim={true_im_ndim}'
                im = np.stack(im_arr, axis=stack_axis)

            im_meta[2] = stack_axis
            im_meta[3] = path

            return im, im_meta

    @staticmethod
    def save_np_as_im(target_path: str, im: np.ndarray, ftype=None, stack_axis=None,
                      concat_instead=False, ftype_meta=None, allow_pickle=False, im_format=None):
        """
        If ftype is None (default), it will be inferred from path
        """
        true_im_ndim = len(im.shape)
        if type(target_path) is str:
            assert not os.path.exists(target_path), f'{target_path} exists!'
            if ftype is None:
                ftype = ImFileType.ftype_from_im_path(target_path)
            if im_format is not None:
                assert im.dtype == ImFileType.format_to_np_dtype[im_format]
            assert ftype != ImFileType.FTYPE_UNKNOWN, 'ERROR: Attempt to save an unknown image type!'
            if ftype == ImFileType.FTYPE_NPZ:
                assert im.dtype == np.float32
                im = np.transpose(im, (0, 3, 2, 1))
                np.savez(target_path, probabilities=im, allow_pickle=allow_pickle)
            elif ftype == ImFileType.FTYPE_NPY:
                if allow_pickle:
                    im = {'masks': im}
                np.save(target_path, im, allow_pickle=allow_pickle)
            elif ftype == ImFileType.FTYPE_NIB:
                assert ftype_meta.shape == (4, 4), f'Affine should have shape 4, 4, but get {ftype_meta[-1].shape}'
                nii_im = nib.Nifti1Image(im, ftype_meta)
                nib.save(nii_im, target_path)
            else:
                Image.fromarray(im).save(target_path)
        else:
            if stack_axis is None:
                stack_axis = true_im_ndim - 1
            for i in range(im.shape[stack_axis]):
                if concat_instead:
                    im_slice = im.take(indices=range(i, i + 1), axis=stack_axis)
                else:
                    im_slice = im.take(indices=i, axis=stack_axis)
                target = target_path[i]
                if ftype_meta is None:
                    cur_ftype_meta = None
                else:
                    cur_ftype_meta = ftype_meta[i]
                # recursive call!
                ImFileType.save_np_as_im(
                    target_path=target,
                    im=im_slice,
                    ftype=ftype,
                    stack_axis=stack_axis,
                    concat_instead=False,
                    ftype_meta=cur_ftype_meta,
                    allow_pickle=allow_pickle)


@dataclass
class ImReadSetting:
    """
    im_format (int) - specifies the data storage layout e.g. ImFileType.FORMAT_UINT8
    stack_axis (int) - if not None, combine several files with different zstacks to the same imid; specify the axis to stack on
    concat_instead (bool) - if True, the stacking are done by np.concatenate() which preserves im dimensions; over the same axis
    true_im_ndim (int) - # of dimensions of a single image; includes number of channels (for 2d BGR image this would be 3)
    allow_pickle (bool) - set allow_pickle for np.load and np.save
    imid_keep_dirpath (bool) - keeping the directory path when extracting imid from paths; default to False
    """
    im_format: int = ImFileType.FORMAT_UINT8
    stack_axis: int | None = None
    concat_instead: bool = False
    true_im_ndim: int = 3
    allow_pickle: bool = False
    imid_keep_dirpath: bool = False


@dataclass
class ImWriteSetting:
    """
    im_format (int) - specifies the data storage layout e.g. ImFileType.FORMAT_UINT8
    stack_axis (int) - if not None, combine several files with different zstacks to the same imid; specify the axis to stack on
    concat_instead (bool) - if True, the stacking are done by np.concatenate() which preserves im dimensions; over the same axis
    allow_pickle (bool) - set allow_pickle for np.load and np.save
    ftype (int) - file type to write e.g. ImFileType.FTYPE_NPZ
    """
    im_format: int = ImFileType.FORMAT_UINT8
    stack_axis: int | None = None
    concat_instead: bool = False
    allow_pickle: bool = False
    ftype: int | None = None


class ImIO:
    @staticmethod
    def read_single_image(read_setting: ImReadSetting, path):
        im, im_meta = ImFileType.load_im_as_np(read_setting.im_format, path,
                                               true_im_ndim=read_setting.true_im_ndim,
                                               stack_axis=read_setting.stack_axis,
                                               concat_instead=read_setting.concat_instead,
                                               allow_pickle=read_setting.allow_pickle)
        return im, im_meta

    @staticmethod
    def write_single_image(write_setting: ImWriteSetting, path, im, im_meta):
        ImFileType.save_np_as_im(path, im=im, ftype=write_setting.ftype,
                                 stack_axis=write_setting.stack_axis,
                                 concat_instead=write_setting.concat_instead,
                                 ftype_meta=im_meta[-1],
                                 allow_pickle=write_setting.allow_pickle)

    @staticmethod
    def read_filenames(read_setting: ImReadSetting, pattern: str | None = None, path: list | None = None):
        """
        The keys (imids) of the ArrayKeyDict returned by this function will be sorted alphabetically
        Parameters
            path (list) - optional and passed in as a list of all paths if pattern is None
        Return
            imid_to_path (Dict) - maps from imids to either a filename or a list of filenames (if image is stacked);
                The keys of this dictionary is a sorted list of image imids, act as unique ids of images, sorted alphabetically
        """
        imid_to_path = ImIO._get_imid_to_path_mapping(read_setting, path, pattern)
        imid_to_path.reorder_keys(sorted(imid_to_path.ordered_keys()))
        return imid_to_path

    @staticmethod
    def read_image(read_setting: ImReadSetting, paths: list):
        """
        Parameter
            paths (list) - can be nested in case of stacked images
        Return
            an iterator whose elements are (im, im_meta), im is the np array of the image and im_meta contains much
            additional metadata returned by ImFileType.load_im_as_np()
        """
        assert type(paths) is not ArrayKeyDict, f'Attempt to pass an ArrayKeyDict to read_image()!'
        for path in paths:
            im, im_meta = ImIO.read_single_image(read_setting, path)
            yield im, im_meta

    @staticmethod
    def write_image(write_setting: ImWriteSetting, paths: list, im_iter: iter):
        """
        Parameter
            paths (list) - can be nested in case of stacked images
            im_iter (iter) - An iterator returning tuples (im, im_meta), the order of these images should be the same as
                the order of their paths in the paths variable
        """
        for path in paths:
            im, im_meta = next(im_iter)
            ImIO.write_single_image(write_setting, path, im, im_meta)

    @staticmethod
    def _get_imid_to_path_mapping(read_setting: ImReadSetting,
                                  paths: list[str] | None = None, pattern: str | None = None):
        if paths is None:
            paths = glob.glob(pattern)
        elif pattern is not None:
            print('WARNING: both paths and pattern are not None! Is this intended?')

        imid_to_path = ArrayKeyDict()
        for path in paths:
            ftype = ImFileType.ftype_from_im_path(path)
            if ftype == ImFileType.FTYPE_UNKNOWN:
                continue
            imid, zstack_str = ImFileType.get_imid_and_zstack(path, is_stack=read_setting.stack_axis is not None,
                                                              keep_dirpath=read_setting.imid_keep_dirpath)
            if zstack_str == '':
                imid_to_path[imid] = path
            else:
                tup = (zstack_str, path)
                if imid not in imid_to_path:
                    imid_to_path[imid] = [tup]
                else:
                    imid_to_path[imid].append(tup)
        for k in list(imid_to_path.ordered_keys()):
            v = imid_to_path[k]
            if type(v) is list:
                imid_to_path[k] = [tup[1] for tup in sorted(v)]
        return imid_to_path


def im_resize(target_size: tuple, im: np.ndarray, order: int = 0) -> np.ndarray:
    assert len(im.shape) == len(target_size), f'Image has shape {im.shape} but intend to resize ' \
                                              f'to incompatible shape {target_size}'
    zoom_factors = (target_size[i] / im.shape[i] for i in range(len(target_size)))
    resized_im = scipy.ndimage.zoom(im, zoom_factors, order=order)  # order=0 -> use nearest neighbor
    return resized_im


class TmpDirectory:
    def __init__(self, dirpath: str):
        self.cur_idx = 0
        self.dirpath = dirpath

    def __enter__(self):
        if os.path.exists(self.dirpath):
            for _ in os.listdir():
                raise FileExistsError('Temporary directory path must not contain existing files, please check if any '
                                      'file exists.')
        ensure_dir_exists(self.dirpath, remove_if_already_exists=True)
        return self

    def assign_tmpdir(self):
        """Return a path prefix that is guaranteed to be empty within the temporary directory

        Files and directories with this prefix can be created then
        """
        tmppath = f'{self.dirpath}/{self.cur_idx}_'
        self.cur_idx += 1
        return tmppath

    def __exit__(self, exc_type, exc_val, exc_tb):
        shutil.rmtree(self.dirpath)


def test():
    assert ImFileType.ftype_from_im_path('./some_root/hello.nii.gz') == ImFileType.FTYPE_NIB
    assert ImFileType.ftype_from_im_path('./some_root/hello_0001.JPG') == ImFileType.FTYPE_JPG
    assert ImFileType.ftype_from_im_path('./a.py') == ImFileType.FTYPE_UNKNOWN
    assert ImFileType.get_imid('./some_root/hello.png') == 'hello'
    assert (ImFileType.get_imid('./some_root/nibfile0001.nii.gz', is_stack=True, keep_dirpath=True) ==
            './some_root/nibfile')

    print('All tests successfully completed')
