from abc import ABC, abstractmethod
import os
from pathlib import Path
from typing import Dict

import numpy as np

import czifile
import tifffile

class GenericImageDataAccessor(ABC):

    @abstractmethod
    def __init__(self):
        """
        Abstract base class that exposes an interfaces for image data, irrespective of whether it is instantiated
        from file I/O or other means.  Enforces X, Y, C, Z dimensions in that order.
        """
        pass

    @property
    def chroma(self):
        return self.shape_dict['C']

    @staticmethod
    def conform_data(data):
        if len(data.shape) > 4 or (0 in data.shape):
            raise DataShapeError(f'Cannot handle image with dimensions other than X, Y, C, and Z: {data.shape}')
        ones = [1 for i in range(0, 4 - len(data.shape))]
        return data.reshape(*data.shape, *ones)

    def is_3d(self):
        return True if self.shape_dict['Z'] > 1 else False

    def get_one_channel_data (self, channel: int):
        c = int(channel)
        return InMemoryDataAccessor(self.data[:, :, c:(c+1), :])

    @property
    def data(self):
        """
        Return data as 4d with axes in order of Y, X, C, Z
        :return: np.ndarray
        """
        return self._data

    @property
    def shape(self):
        return self._data.shape

    @property
    def shape_dict(self):
        return dict(zip(('Y', 'X', 'C', 'Z'), self.data.shape))

class InMemoryDataAccessor(GenericImageDataAccessor):
    def __init__(self, data):
        self._data = self.conform_data(data)

class GenericImageFileAccessor(GenericImageDataAccessor): # image data is loaded from a file
    def __init__(self, fpath: Path):
        """
        Interface for image data that originates in an image file
        :param fpath: absolute path to image file
        :param kwargs: variable-length keyword arguments
        """
        if not os.path.exists(fpath):
            raise FileAccessorError(f'Could not find file at {fpath}')
        self.fpath = fpath

class CziImageFileAccessor(GenericImageFileAccessor):
    """
    Image that is stored in a Zeiss .CZI file; may be multi-channel, and/or a z-stack,
    but not a time series or multiposition acquisition.
    """
    def __init__(self, fpath: Path):
        super().__init__(fpath)

        try:
            cf = czifile.CziFile(fpath)
            self.czifile = cf
        except Exception:
            raise FileAccessorError(f'Unable to access CZI data in {fpath}')

        sd = {ch: cf.shape[cf.axes.index(ch)] for ch in cf.axes}
        if (sd.get('S') and (sd['S'] > 1)) or (sd.get('T') and (sd['T'] > 1)):
            raise DataShapeError(f'Cannot handle image with multiple positions or time points: {sd}')

        idx = {k: sd[k] for k in ['Y', 'X', 'C', 'Z']}
        yxcz = np.moveaxis(
            cf.asarray(),
            [cf.axes.index(ch) for ch in idx],
            [0, 1, 2, 3]
        )

        self._data = self.conform_data(yxcz.reshape(yxcz.shape[0:4]))

    def __del__(self):
        self.czifile.close()


def write_accessor_data_to_file(fpath: Path, accessor: GenericImageDataAccessor) -> bool:
    try:
        zcyx= np.moveaxis(
            accessor.data, # yxcz
            [3, 2, 0, 1],
            [0, 1, 2, 3]
        )
        tifffile.imwrite(fpath, zcyx, imagej=True)
    except:
        raise FileWriteError(f'Unable to write data to file')
    return True


def generate_file_accessor(fpath):
    if str(fpath).upper().endswith('.CZI'):
        return CziImageFileAccessor(fpath)
    else:
        raise FileAccessorError(f'Could not match a file accessor with {fpath}')

class Error(Exception):
    pass

class FileAccessorError(Error):
    pass

class FileNotFoundError(Error):
    pass

class DataShapeError(Error):
    pass

class FileWriteError(Error):
    pass

class InvalidAxisKey(Error):
    pass

class InvalidDataShape(Error):
    pass