import multiprocessing
import re
import warnings
from abc import ABC, ABCMeta, abstractmethod
from argparse import ArgumentParser
from collections import OrderedDict
from datetime import datetime
from functools import cached_property, wraps
from importlib.metadata import version
from itertools import product
from numbers import Number
from operator import truediv
from pathlib import Path
from traceback import print_exc

import numpy as np
import ome_types
import yaml
from ome_types import model, ureg
from pint import set_application_registry
from tiffwrite import IJTiffFile
from tqdm.auto import tqdm

from .jvm import JVM
from .transforms import Transform, Transforms

try:
    __version__ = version(Path(__file__).parent.name)
except Exception:  # noqa
    __version__ = 'unknown'

try:
    with open(Path(__file__).parent.parent / '.git' / 'HEAD') as g:
        head = g.read().split(':')[1].strip()
    with open(Path(__file__).parent.parent / '.git' / head) as h:
        __git_commit_hash__ = h.read().rstrip('\n')
except Exception:  # noqa
    __git_commit_hash__ = 'unknown'

ureg.default_format = '~P'
set_application_registry(ureg)
warnings.filterwarnings('ignore', 'Reference to unknown ID')


class ReaderNotFoundError(Exception):
    pass


class TransformTiff(IJTiffFile):
    """ transform frames in a parallel process to speed up saving """
    def __init__(self, image, *args, **kwargs):
        self.image = image
        super().__init__(*args, **kwargs)

    def compress_frame(self, frame):
        return super().compress_frame(np.asarray(self.image(*frame)).astype(self.dtype))


class DequeDict(OrderedDict):
    def __init__(self, maxlen=None, *args, **kwargs):
        self.maxlen = maxlen
        super().__init__(*args, **kwargs)

    def __truncate__(self):
        if self.maxlen is not None:
            while len(self) > self.maxlen:
                self.popitem(False)

    def __setitem__(self, *args, **kwargs):
        super().__setitem__(*args, **kwargs)
        self.__truncate__()

    def update(self, *args, **kwargs):
        super().update(*args, **kwargs)
        self.__truncate__()


def find(obj, **kwargs):
    for item in obj:
        if all([getattr(item, key) == value for key, value in kwargs.items()]):
            return item


def try_default(fun, default, *args, **kwargs):
    try:
        return fun(*args, **kwargs)
    except Exception:  # noqa
        return default


def get_ome(path):
    from .readers.bfread import jars
    try:
        jvm = JVM(jars)  # noqa
        ome_meta = jvm.metadata_tools.createOMEXMLMetadata()
        reader = jvm.image_reader()
        reader.setMetadataStore(ome_meta)
        reader.setId(str(path))
        ome = ome_types.from_xml(str(ome_meta.dumpXML()), parser='lxml')
    except Exception:  # noqa
        print_exc()
        ome = model.OME()
    finally:
        jvm.kill_vm()  # noqa
    return ome


class Shape(tuple):
    def __new__(cls, shape, axes='xyczt'):
        if isinstance(shape, Shape):
            axes = shape.axes
        instance = super().__new__(cls, shape)
        instance.axes = axes.lower()
        return instance

    def __getitem__(self, n):
        if isinstance(n, str):
            if len(n) == 1:
                return self[self.axes.find(n.lower())] if n.lower() in self.axes else 1
            else:
                return tuple(self[i] for i in n)
        return super().__getitem__(n)

    @cached_property
    def xyczt(self):
        return tuple(self[i] for i in 'xyczt')


class Imread(np.lib.mixins.NDArrayOperatorsMixin, ABC):
    """ class to read image files, while taking good care of important metadata,
        currently optimized for .czi files, but can open anything that bioformats can handle
            path: path to the image file
            optional:
            series: in case multiple experiments are saved in one file, like in .lif files
            dtype: datatype to be used when returning frames
            meta: define metadata, used for pickle-ing

            NOTE: run imread.kill_vm() at the end of your script/program, otherwise python might not terminate

            modify images on the fly with a decorator function:
                define a function which takes an instance of this object, one image frame,
                and the coordinates c, z, t as arguments, and one image frame as return
                >> imread.frame_decorator = fun
                then use imread as usually

            Examples:
                >> im = imread('/DATA/lenstra_lab/w.pomp/data/20190913/01_YTL639_JF646_DefiniteFocus.czi')
                >> im
                 << shows summary
                >> im.shape
                 << (256, 256, 2, 1, 600)
                >> plt.imshow(im(1, 0, 100))
                 << plots frame at position c=1, z=0, t=100 (python type indexing), note: round brackets; always 2d array
                    with 1 frame
                >> data = im[:,:,0,0,:25]
                 << retrieves 5d numpy array containing first 25 frames at c=0, z=0, note: square brackets; always 5d array
                >> plt.imshow(im.max(0, None, 0))
                 << plots max-z projection at c=0, t=0
                >> len(im)
                 << total number of frames
                >> im.pxsize
                 << 0.09708737864077668 image-plane pixel size in um
                >> im.laserwavelengths
                 << [642, 488]
                >> im.laserpowers
                 << [0.02, 0.0005] in %

                See __init__ and other functions for more ideas.

            Subclassing:
                Subclass AbstractReader to add more file types. A subclass should always have at least the following
                methods:
                    staticmethod _can_open(path): returns True when the subclass can open the image in path
                    __metadata__(self): pulls some metadata from the file and do other format specific things,
                                        it needs to define a few properties, like shape, etc.
                    __frame__(self, c, z, t): this should return a single frame at channel c, slice z and time t
                    optional close(self): close the file in a proper way
                    optional field priority: subclasses with lower priority will be tried first, default = 99
                    Any other method can be overridden as needed
            wp@tl2019-2023 """

    def __new__(cls, path=None, *args, **kwargs):
        if cls is not Imread:
            return super().__new__(cls)
        if len(AbstractReader.__subclasses__()) == 0:
            raise Exception('Restart python kernel please!')
        if isinstance(path, Imread):
            return path
        path, _ = AbstractReader.split_path_series(path)
        for subclass in sorted(AbstractReader.__subclasses__(), key=lambda subclass_: subclass_.priority):
            if subclass._can_open(path):
                do_not_pickle = (AbstractReader.do_not_pickle,) if isinstance(AbstractReader.do_not_pickle, str) \
                    else AbstractReader.do_not_pickle
                subclass_do_not_pickle = (subclass.do_not_pickle,) if isinstance(subclass.do_not_pickle, str) \
                    else subclass.do_not_pickle if hasattr(subclass, 'do_not_pickle') else ()
                subclass.do_not_pickle = set(do_not_pickle).union(set(subclass_do_not_pickle))

                return super().__new__(subclass)
        raise ReaderNotFoundError(f'No reader found for {path}.')

    def __init__(self, base=None, slice=None, shape=(0, 0, 0, 0, 0), dtype=None, frame_decorator=None):
        self.base = base or self
        self.slice = slice
        self._shape = Shape(shape)
        self.dtype = dtype
        self.frame_decorator = frame_decorator
        self.transform = Transforms()
        self.flags = dict(C_CONTIGUOUS=False, F_CONTIGUOUS=False, OWNDATA=False, WRITEABLE=False,
                          ALIGNED=False, WRITEBACKIFCOPY=False, UPDATEIFCOPY=False)

    def __call__(self, c=None, z=None, t=None, x=None, y=None):
        """ same as im[] but allowing keyword axes, but slices need to made with slice() or np.s_ """
        return self[{k: slice(v) if v is None else v for k, v in dict(c=c, z=z, t=t, x=x, y=y).items()}]

    def __copy__(self):
        return self.copy()

    def __contains__(self, item):
        def unique_yield(a, b):
            for k in a:
                yield k
            for k in b:
                if k not in a:
                    yield k

        for idx in unique_yield([key[:3] for key in self.cache.keys()],
                                product(range(self.shape['c']), range(self.shape['z']), range(self.shape['t']))):
            xyczt = (slice(None), slice(None)) + idx
            in_idx = tuple(xyczt['xyczt'.find(i)] for i in self.axes)
            if item in np.asarray(self[in_idx]):
                return True
        return False

    def __enter__(self):
        return self

    def __exit__(self, *args, **kwargs):
        if not self.isclosed:
            self.isclosed = True
            if hasattr(self, 'close'):
                self.close()

    def __getitem__(self, n):
        """ slice like a numpy array but return an Imread instance """
        if self.isclosed:
            raise OSError("file is closed")
        if isinstance(n, (slice, Number)):  # None = :
            n = (n,)
        elif isinstance(n, type(Ellipsis)):
            n = (None,) * len(self.axes)
        elif isinstance(n, dict):  # allow im[dict(z=0)] etc.
            n = [n.get(i, slice(None)) for i in self.axes]
        n = list(n)

        # deal with ...
        ell = [i for i, e in enumerate(n) if isinstance(e, type(Ellipsis))]
        if len(ell) > 1:
            raise IndexError('an index can only have a single ellipsis (...)')
        if len(ell):
            if len(n) > self.ndim:
                n.remove(Ellipsis)
            else:
                n[ell[0]] = None
                while len(n) < self.ndim:
                    n.insert(ell[0], None)
        while len(n) < self.ndim:
            n.append(None)

        axes_idx = [self.shape.axes.find(i) for i in 'xyczt']
        n = [n[j] if 0 <= j < len(n) else None for j in axes_idx]  # reorder n

        new_slice = []
        for s, e in zip(self.slice, n):
            if e is None:
                new_slice.append(s)
            else:
                new_slice.append(s[e])

        # TODO: check output dimensionality when requested shape in some dimension is 1
        if all([isinstance(s, Number) or a < 0 and s.size == 1 for s, a in zip(new_slice, axes_idx)]):
            return self.block(*new_slice).item()
        else:
            new = View(self)
            new.slice = new_slice
            new._shape = Shape([1 if isinstance(s, Number) else len(s) for s in new_slice])
            new.axes = ''.join(j for j in self.axes if j in [i for i, s in zip('xyczt', new_slice)
                                                             if not isinstance(s, Number)])
            return new

    def __getstate__(self):
        return {key: value for key, value in self.__dict__.items() if key not in self.do_not_pickle}

    def __len__(self):
        return self.shape[0]

    def __repr__(self):
        return self.summary

    def __setstate__(self, state):
        """ What happens during unpickling """
        self.__dict__.update(state)
        if isinstance(self, AbstractReader):
            self.open()
        self.cache = DequeDict(16)

    def __str__(self):
        return str(self.path)

    def __array__(self, dtype=None):
        block = self.block(*self.slice)
        axes_idx = [self.shape.axes.find(i) for i in 'xyczt']
        axes_squeeze = tuple({i for i, j in enumerate(axes_idx) if j == -1}.union(
            {i for i, j in enumerate(self.slice) if isinstance(j, Number)}))
        block = block.squeeze(axes_squeeze)
        if dtype is not None:
            block = block.astype(dtype)
        if block.ndim == 0:
            return block.item()
        axes = ''.join(j for i, j in enumerate('xyczt') if i not in axes_squeeze)
        return block.transpose([axes.find(i) for i in self.shape.axes if i in axes])

    def __array_arg_fun__(self, fun, axis=None, out=None):
        """ frame-wise application of np.argmin and np.argmax """
        if axis is None:
            value = arg = None
            for idx in product(range(self.shape['c']), range(self.shape['z']), range(self.shape['t'])):
                xyczt = (slice(None), slice(None)) + idx
                in_idx = tuple(xyczt['xyczt'.find(i)] for i in self.axes)
                new = np.asarray(self[in_idx])
                new_arg = np.unravel_index(fun(new), new.shape)
                new_value = new[new_arg]
                if value is None:
                    arg = new_arg + idx
                    value = new_value
                else:
                    i = fun((value, new_value))
                    arg = (arg, new_arg + idx)[i]
                    value = (value, new_value)[i]
            axes = ''.join(i for i in self.axes if i in 'xy') + 'czt'
            arg = np.ravel_multi_index([arg[axes.find(i)] for i in self.axes], self.shape)
            if out is None:
                return arg
            else:
                out.itemset(arg)
                return out
        else:
            if isinstance(axis, str):
                axis_str, axis_idx = axis, self.axes.index(axis)
            else:
                axis_str, axis_idx = self.axes[axis], axis
            if axis_str not in self.axes:
                raise IndexError(f'Axis {axis_str} not in {self.axes}.')
            out_shape = list(self.shape)
            out_axes = list(self.axes)
            out_shape.pop(axis_idx)
            out_axes.pop(axis_idx)
            if out is None:
                out = np.zeros(out_shape, int)
            if axis_str in 'xy':
                for idx in product(range(self.shape['c']), range(self.shape['z']), range(self.shape['t'])):
                    xyczt = (slice(None), slice(None)) + idx
                    out_idx = tuple(xyczt['xyczt'.find(i)] for i in out_axes)
                    in_idx = tuple(xyczt['xyczt'.find(i)] for i in self.axes)
                    new = self[in_idx]
                    out[out_idx] = fun(np.asarray(new), new.axes.find(axis_str))
            else:
                value = np.zeros(out.shape, self.dtype)
                for idx in product(range(self.shape['c']), range(self.shape['z']), range(self.shape['t'])):
                    xyczt = (slice(None), slice(None)) + idx
                    out_idx = tuple(xyczt['xyczt'.find(i)] for i in out_axes)
                    in_idx = tuple(xyczt['xyczt'.find(i)] for i in self.axes)
                    new_value = self[in_idx]
                    new_arg = np.full_like(new_value, idx['czt'.find(axis_str)])
                    if idx['czt'.find(axis_str)] == 0:
                        value[out_idx] = new_value
                        out[out_idx] = new_arg
                    else:
                        old_value = value[out_idx]
                        i = fun((old_value, new_value), 0)
                        value[out_idx] = np.where(i, new_value, old_value)
                        out[out_idx] = np.where(i, new_arg, out[out_idx])
            return out

    def __array_fun__(self, funs, axis=None, dtype=None, out=None, keepdims=False, initials=None, where=True,
                      ffuns=None, cfun=None):
        """ frame-wise application of np.min, np.max, np.sum, np.mean and their nan equivalents """
        p = re.compile(r'\d')
        dtype = self.dtype if dtype is None else np.dtype(dtype)
        if initials is None:
            initials = [None for _ in funs]
        if ffuns is None:
            ffuns = [None for _ in funs]

        def ffun_(frame):
            return np.asarray(frame)
        ffuns = [ffun_ if ffun is None else ffun for ffun in ffuns]
        if cfun is None:
            def cfun(*res):
                return res[0]

        # TODO: smarter transforms
        if axis is None:
            for idx in product(range(self.shape['c']), range(self.shape['z']), range(self.shape['t'])):
                xyczt = (slice(None), slice(None)) + idx
                in_idx = tuple(xyczt['xyczt'.find(i)] for i in self.axes)
                w = where if where is None or isinstance(where, bool) else where[in_idx]
                initials = [fun(np.asarray(ffun(self[in_idx])), initial=initial, where=w)
                            for fun, ffun, initial in zip(funs, ffuns, initials)]
            res = cfun(*initials)
            res = (np.round(res) if dtype.kind in 'ui' else res).astype(p.sub('', dtype.name))
            if keepdims:
                res = np.array(res, dtype, ndmin=self.ndim)
            if out is None:
                return res
            else:
                out.itemset(res)
                return out
        else:
            if isinstance(axis, str):
                axis_str, axis_idx = axis, self.axes.index(axis)
            else:
                axis_idx = axis % self.ndim
                axis_str = self.axes[axis_idx]
            if axis_str not in self.axes:
                raise IndexError(f'Axis {axis_str} not in {self.axes}.')
            out_shape = list(self.shape)
            out_axes = list(self.axes)
            if not keepdims:
                out_shape.pop(axis_idx)
                out_axes.pop(axis_idx)
            if out is None:
                out = np.zeros(out_shape, dtype)
            if axis_str in 'xy':
                xy = 'xy' if self.axes.find('x') < self.axes.find('y') else 'yx'
                frame_ax = xy.find(axis_str)
                for idx in product(range(self.shape['c']), range(self.shape['z']), range(self.shape['t'])):
                    xyczt = (slice(None), slice(None)) + idx
                    out_idx = tuple(xyczt['xyczt'.find(i)] for i in out_axes)
                    in_idx = tuple(xyczt['xyczt'.find(i)] for i in self.axes)
                    w = where if where is None or isinstance(where, bool) else where[in_idx]
                    res = cfun(*[fun(ffun(self[in_idx]), frame_ax, initial=initial, where=w)
                                 for fun, ffun, initial in zip(funs, ffuns, initials)])
                    out[out_idx] = (np.round(res) if out.dtype.kind in 'ui' else res).astype(p.sub('', dtype.name))
            else:
                tmps = [np.zeros(out_shape) for _ in ffuns]
                for idx in product(range(self.shape['c']), range(self.shape['z']), range(self.shape['t'])):
                    xyczt = (slice(None), slice(None)) + idx
                    out_idx = tuple(xyczt['xyczt'.find(i)] for i in out_axes)
                    in_idx = tuple(xyczt['xyczt'.find(i)] for i in self.axes)

                    if idx['czt'.find(axis_str)] == 0:
                        w = where if where is None or isinstance(where, bool) else (where[in_idx],)
                        for tmp, fun, ffun, initial in zip(tmps, funs, ffuns, initials):
                            tmp[out_idx] = fun((ffun(self[in_idx]),), 0, initial=initial, where=w)
                    else:
                        w = where if where is None or isinstance(where, bool) else \
                            (np.ones_like(where[in_idx]), where[in_idx])
                        for tmp, fun, ffun in zip(tmps, funs, ffuns):
                            tmp[out_idx] = fun((tmp[out_idx], ffun(self[in_idx])), 0, where=w)
                out[...] = (np.round(cfun(*tmps)) if out.dtype.kind in 'ui' else
                            cfun(*tmps)).astype(p.sub('', dtype.name))
            return out

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

    @axes.setter
    def axes(self, value):
        shape = self.shape[value]
        if isinstance(shape, Number):
            shape = (shape,)
        self._shape = Shape(shape, value)

    @property
    def dtype(self):
        return self._dtype

    @dtype.setter
    def dtype(self, value):
        self._dtype = np.dtype(value)

    @cached_property
    def extrametadata(self):
        if isinstance(self.path, Path):
            if self.path.with_suffix('.pzl2').exists():
                pname = self.path.with_suffix('.pzl2')
            elif self.path.with_suffix('.pzl').exists():
                pname = self.path.with_suffix('.pzl')
            else:
                return
            try:
                return self.get_config(pname)
            except Exception:  # noqa
                return
        return

    @property
    def ndim(self):
        return len(self.shape)

    @property
    def size(self):
        return np.prod(self.shape)

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

    @shape.setter
    def shape(self, value):
        if isinstance(value, Shape):
            self._shape = value
        else:
            self._shape = Shape((value['xyczt'.find(i.lower())] for i in self.axes), self.axes)

    @property
    def summary(self):
        """ gives a helpful summary of the recorded experiment """
        s = [f"path/filename: {self.path}",
             f"series/pos:    {self.series}",
             f"reader:        {self.base.__class__.__module__.split('.')[-1]}"]
        s.extend((f"dtype:         {self.dtype}",
                  f"shape ({self.axes}):".ljust(15) + f"{' x '.join(str(i) for i in self.shape)}"))
        if self.pxsize_um:
            s.append(f'pixel size:    {1000 * self.pxsize_um:.2f} nm')
        if self.zstack and self.deltaz_um:
            s.append(f'z-interval:    {1000 * self.deltaz_um:.2f} nm')
        if self.exposuretime_s and not all(e is None for e in self.exposuretime_s):
            s.append(f'exposuretime:  {self.exposuretime_s[0]:.2f} s')
        if self.timeseries and self.timeinterval:
            s.append(f'time interval: {self.timeinterval:.3f} s')
        if self.binning:
            s.append('binning:       {}x{}'.format(*self.binning))
        if self.laserwavelengths:
            s.append('laser colors:  ' + ' | '.join([' & '.join(len(w) * ('{:.0f}',)).format(*w)
                                                     for w in self.laserwavelengths]) + ' nm')
        if self.laserpowers:
            s.append('laser powers:  ' + ' | '.join([' & '.join(len(p) * ('{:.3g}',)).format(*[100 * i for i in p])
                                                     for p in self.laserpowers]) + ' %')
        if self.objective:
            s.append(f'objective:     {self.objective.model}')
        if self.magnification:
            s.append(f'magnification: {self.magnification}x')
        if self.tubelens:
            s.append(f'tubelens:      {self.tubelens.model}')
        if self.filter:
            s.append(f'filterset:     {self.filter}')
        if self.powermode:
            s.append(f'powermode:     {self.powermode}')
        if self.collimator:
            s.append('collimator:   ' + (' {}' * len(self.collimator)).format(*self.collimator))
        if self.tirfangle:
            s.append('TIRF angle:   ' + (' {:.2f}°' * len(self.tirfangle)).format(*self.tirfangle))
        if self.gain:
            s.append('gain:         ' + (' {:.0f}' * len(self.gain)).format(*self.gain))
        if self.pcf:
            s.append('pcf:          ' + (' {:.2f}' * len(self.pcf)).format(*self.pcf))
        return '\n'.join(s)

    @property
    def T(self):
        return self.transpose()

    @cached_property
    def timeseries(self):
        return self.shape['t'] > 1

    @cached_property
    def zstack(self):
        return self.shape['z'] > 1

    @wraps(np.argmax)
    def argmax(self, *args, **kwargs):
        return self.__array_arg_fun__(np.argmax, *args, **kwargs)

    @wraps(np.argmin)
    def argmin(self, *args, **kwargs):
        return self.__array_arg_fun__(np.argmin, *args, **kwargs)

    @wraps(np.max)
    def max(self, axis=None, out=None, keepdims=False, initial=None, where=True, **kwargs):
        return self.__array_fun__([np.max], axis, None, out, keepdims, [initial], where)

    @wraps(np.mean)
    def mean(self, axis=None, dtype=None, out=None, keepdims=False, *, where=True, **kwargs):
        dtype = dtype or float
        n = np.prod(self.shape) if axis is None else self.shape[axis]

        def sfun(frame):
            return np.asarray(frame).astype(float)

        def cfun(res):
            return res / n

        return self.__array_fun__([np.sum], axis, dtype, out, keepdims, None, where, [sfun], cfun)

    @wraps(np.min)
    def min(self, axis=None, out=None, keepdims=False, initial=None, where=True, **kwargs):
        return self.__array_fun__([np.min], axis, None, out, keepdims, [initial], where)

    @wraps(np.moveaxis)
    def moveaxis(self, source, destination):
        raise NotImplementedError('moveaxis is not implemented')

    @wraps(np.nanmax)
    def nanmax(self, axis=None, out=None, keepdims=False, initial=None, where=True, **kwargs):
        return self.__array_fun__([np.nanmax], axis, None, out, keepdims, [initial], where)

    @wraps(np.nanmean)
    def nanmean(self, axis=None, dtype=None, out=None, keepdims=False, *, where=True, **kwargs):
        dtype = dtype or float

        def sfun(frame):
            return np.asarray(frame).astype(float)

        def nfun(frame):
            return np.invert(np.isnan(frame))

        return self.__array_fun__([np.nansum, np.sum], axis, dtype, out, keepdims, None, where, (sfun, nfun), truediv)

    @wraps(np.nanmin)
    def nanmin(self, axis=None, out=None, keepdims=False, initial=None, where=True, **kwargs):
        return self.__array_fun__([np.nanmin], axis, None, out, keepdims, [initial], where)

    @wraps(np.nansum)
    def nansum(self, axis=None, dtype=None, out=None, keepdims=False, initial=None, where=True, **kwargs):
        return self.__array_fun__([np.nansum], axis, dtype, out, keepdims, [initial], where)

    @wraps(np.nanstd)
    def nanstd(self, axis=None, dtype=None, out=None, ddof=0, keepdims=None, *, where=None):
        return self.nanvar(axis, dtype, out, ddof, keepdims, where=where, std=True)

    @wraps(np.nanvar)
    def nanvar(self, axis=None, dtype=None, out=None, ddof=0, keepdims=None, *, where=True, std=False):
        dtype = dtype or float

        def sfun(frame):
            return np.asarray(frame).astype(float)

        def s2fun(frame):
            return np.asarray(frame).astype(float) ** 2

        def nfun(frame):
            return np.invert(np.isnan(frame))

        if std:
            def cfun(s, s2, n):
                return np.sqrt((s2 - s ** 2 / n) / (n - ddof))
        else:
            def cfun(s, s2, n):
                return (s2 - s ** 2 / n) / (n - ddof)
        return self.__array_fun__([np.nansum, np.nansum, np.sum], axis, dtype, out, keepdims, None, where,
                                  (sfun, s2fun, nfun), cfun)

    def flatten(self, *args, **kwargs):
        return np.asarray(self).flatten(*args, **kwargs)

    @wraps(np.reshape)
    def reshape(self, *args, **kwargs):
        return np.asarray(self).reshape(*args, **kwargs)

    @wraps(np.squeeze)
    def squeeze(self, axes=None):
        new = self.copy()
        if axes is None:
            axes = tuple(i for i, j in enumerate(new.shape) if j == 1)
        elif isinstance(axes, Number):
            axes = (axes,)
        else:
            axes = tuple(new.axes.find(ax) if isinstance(ax, str) else ax for ax in axes)
        if any([new.shape[ax] != 1 for ax in axes]):
            raise ValueError('cannot select an axis to squeeze out which has size not equal to one')
        new.axes = ''.join(j for i, j in enumerate(new.axes) if i not in axes)
        return new

    @wraps(np.std)
    def std(self, axis=None, dtype=None, out=None, ddof=0, keepdims=None, *, where=True):
        return self.var(axis, dtype, out, ddof, keepdims, where=where, std=True)

    @wraps(np.sum)
    def sum(self, axis=None, dtype=None, out=None, keepdims=False, initial=None, where=True, **kwargs):
        return self.__array_fun__([np.sum], axis, dtype, out, keepdims, [initial], where)

    @wraps(np.swapaxes)
    def swapaxes(self, axis1, axis2):
        new = self.copy()
        axes = new.axes
        if isinstance(axis1, str):
            axis1 = axes.find(axis1)
        if isinstance(axis2, str):
            axis2 = axes.find(axis2)
        axes[axis1], axes[axis2] = axes[axis2], axes[axis1]
        new.axes = axes
        return new

    @wraps(np.transpose)
    def transpose(self, *axes):
        new = self.copy()
        if not axes:
            new.axes = new.axes[::-1]
        else:
            new.axes = ''.join(ax if isinstance(ax, str) else new.axes[ax] for ax in axes)
        return new

    @wraps(np.var)
    def var(self, axis=None, dtype=None, out=None, ddof=0, keepdims=None, *, where=True, std=False):
        dtype = dtype or float
        n = np.prod(self.shape) if axis is None else self.shape[axis]

        def sfun(frame):
            return np.asarray(frame).astype(float)

        def s2fun(frame):
            return np.asarray(frame).astype(float) ** 2

        if std:
            def cfun(s, s2):
                return np.sqrt((s2 - s ** 2 / n) / (n - ddof))
        else:
            def cfun(s, s2):
                return (s2 - s ** 2 / n) / (n - ddof)
        return self.__array_fun__([np.sum, np.sum], axis, dtype, out, keepdims, None, where, (sfun, s2fun), cfun)

    def asarray(self):
        return self.__array__()

    def astype(self, dtype, *args, **kwargs):
        new = self.copy()
        new.dtype = dtype
        return new

    def block(self, x=None, y=None, c=None, z=None, t=None):
        """ returns 5D block of frames """
        x, y, c, z, t = (np.arange(self.shape[i]) if e is None else np.array(e, ndmin=1)
                         for i, e in zip('xyczt', (x, y, c, z, t)))
        d = np.empty((len(x), len(y), len(c), len(z), len(t)), self.dtype)
        for (ci, cj), (zi, zj), (ti, tj) in product(enumerate(c), enumerate(z), enumerate(t)):
            d[:, :, ci, zi, ti] = self.frame(cj, zj, tj)[x][:, y]
        return d

    def copy(self):
        return View(self)

    def data(self, c=0, z=0, t=0):
        """ returns 3D stack of frames """
        c, z, t = (np.arange(self.shape[i]) if e is None else np.array(e, ndmin=1) for i, e in zip('czt', (c, z, t)))
        return np.dstack([self.frame(ci, zi, ti) for ci, zi, ti in product(c, z, t)])

    def frame(self, c=0, z=0, t=0):
        """ returns single 2D frame """
        c = self.get_channel(c)
        c %= self.base.shape['c']
        z %= self.base.shape['z']
        t %= self.base.shape['t']

        # cache last n (default 16) frames in memory for speed (~250x faster)
        key = (c, z, t, self.transform, self.frame_decorator)
        if key in self.cache:
            self.cache.move_to_end(key)
            f = self.cache[key]
        else:
            f = self.transform[self.channel_names[c], t].frame(self.__frame__(c, z, t))
            if self.frame_decorator is not None:
                f = self.frame_decorator(self, f, c, z, t)
            self.cache[key] = f
        if self.dtype is not None:
            return f.copy().astype(self.dtype)
        else:
            return f.copy()

    def get_channel(self, channel_name):
        if not isinstance(channel_name, str):
            return channel_name
        else:
            c = [i for i, c in enumerate(self.channel_names) if c.lower().startswith(channel_name.lower())]
            assert len(c) > 0, f'Channel {c} not found in {self.channel_names}'
            assert len(c) < 2, f'Channel {c} not unique in {self.channel_names}'
            return c[0]

    @staticmethod
    def get_config(file):
        """ Open a yml config file """
        loader = yaml.SafeLoader
        loader.add_implicit_resolver(
            r'tag:yaml.org,2002:float',
            re.compile(r'''^(?:
                     [-+]?(?:[0-9][0-9_]*)\.[0-9_]*(?:[eE][-+]?[0-9]+)?
                    |[-+]?(?:[0-9][0-9_]*)(?:[eE][-+]?[0-9]+)
                    |\.[0-9_]+(?:[eE][-+][0-9]+)?
                    |[-+]?[0-9][0-9_]*(?::[0-5]?[0-9])+\.[0-9_]*
                    |[-+]?\\.(?:inf|Inf|INF)
                    |\.(?:nan|NaN|NAN))$''', re.X),
            list(r'-+0123456789.'))
        with open(file) as f:
            return yaml.load(f, loader)

    def get_czt(self, c, z, t):
        czt = []
        for i, n in zip('czt', (c, z, t)):
            if n is None:
                czt.append(list(range(self.shape[i])))
            elif isinstance(n, range):
                if n.stop < 0:
                    stop = n.stop % self.shape[i]
                elif n.stop > self.shape[i]:
                    stop = self.shape[i]
                else:
                    stop = n.stop
                czt.append(list(range(n.start % self.shape[i], stop, n.step)))
            elif isinstance(n, Number):
                czt.append([n % self.shape[i]])  # noqa
            else:
                czt.append([k % self.shape[i] for k in n])
        return [self.get_channel(c) for c in czt[0]], *czt[1:]

    @staticmethod
    def get_ome(path):
        """ Use java BioFormats to make an ome metadata structure. """
        with multiprocessing.get_context('spawn').Pool(1) as pool:
            ome = pool.map(get_ome, (path,))[0]
            return ome

    @cached_property
    def ome(self):
        return self.get_ome(self.path)

    def is_noise(self, volume=None):
        """ True if volume only has noise """
        if volume is None:
            volume = self
        fft = np.fft.fftn(volume)
        corr = np.fft.fftshift(np.fft.ifftn(fft * fft.conj()).real / np.sum(volume ** 2))
        return 1 - corr[tuple([0] * corr.ndim)] < 0.0067

    @staticmethod
    def kill_vm():
        JVM().kill_vm()

    def new(self, *args, **kwargs):
        warnings.warn('Imread.new has been deprecated, use Imread.view instead.', DeprecationWarning, 2)
        return self.view(*args, **kwargs)

    def save_as_tiff(self, fname=None, c=None, z=None, t=None, split=False, bar=True, pixel_type='uint16', **kwargs):
        """ saves the image as a tif file
            split: split channels into different files """
        if fname is None:
            fname = self.path.with_suffix('.tif')
            if fname == self.path:
                raise FileExistsError(f'File {fname} exists already.')
        if not isinstance(fname, Path):
            fname = Path(fname)
        if split:
            for i in range(self.shape['c']):
                if self.timeseries:
                    self.save_as_tiff(fname.with_name(f'{fname.stem}_C{i:01d}').with_suffix('.tif'), i, 0, None, False,
                                      bar, pixel_type)
                else:
                    self.save_as_tiff(fname.with_name(f'{fname.stem}_C{i:01d}').with_suffix('.tif'), i, None, 0, False,
                                      bar, pixel_type)
        else:
            n = [c, z, t]
            for i, ax in enumerate('czt'):
                if n[i] is None:
                    n[i] = range(self.shape[ax])
                elif not isinstance(n[i], (tuple, list)):
                    n[i] = (n[i],)

            shape = [len(i) for i in n]
            with TransformTiff(self, fname.with_suffix('.tif'), shape, pixel_type,
                            pxsize=self.pxsize_um, deltaz=self.deltaz_um, **kwargs) as tif:
                for i, m in tqdm(zip(product(*[range(s) for s in shape]), product(*n)),  # noqa
                                 total=np.prod(shape), desc='Saving tiff', disable=not bar):
                    tif.save(m, *i)

    def with_transform(self, channels=True, drift=False, file=None, bead_files=()):
        """ returns a view where channels and/or frames are registered with an affine transformation
            channels: True/False register channels using bead_files
            drift: True/False register frames to correct drift
            file: load registration from file with name file, default: transform.yml in self.path.parent
            bead_files: files used to register channels, default: files in self.path.parent,
                with names starting with 'beads'
            """
        view = self.view()
        if file is None:
            file = Path(view.path.parent) / 'transform.yml'
        if not bead_files:
            try:
                bead_files = Transforms.get_bead_files(view.path.parent)
            except Exception:
                if not file.exists():
                    raise Exception('No transform file and no bead file found.')
                bead_files = ()

        if channels:
            try:
                view.transform = Transforms.from_file(file, T=drift)
            except Exception:  # noqa
                view.transform = Transforms().with_beads(view.cyllens, bead_files)
                if drift:
                    view.transform = view.transform.with_drift(view)
                view.transform.save(file.with_suffix('.yml'))
                view.transform.save_channel_transform_tiff(bead_files, file.with_suffix('.tif'))
        elif drift:
            try:
                view.transform = Transforms.from_file(file, C=False)
            except Exception:  # noqa
                view.transform = Transforms().with_drift(self)
        view.transform.adapt(view.frameoffset, view.shape.xyczt, view.channel_names)
        return view

    @staticmethod
    def split_path_series(path):
        if isinstance(path, str):
            path = Path(path)
        if isinstance(path, Path) and path.name.startswith('Pos'):
            return path.parent, int(path.name.lstrip('Pos'))
        return path, 0

    def view(self, *args, **kwargs):
        return View(self, *args, **kwargs)


class View(Imread, ABC):
    def __init__(self, base, dtype=None):
        super().__init__(base.base, base.slice, base.shape, dtype or base.dtype, base.frame_decorator)
        self.transform = base.transform

    def __getattr__(self, item):
        if not hasattr(self.base, item):
            raise AttributeError(f'{self.__class__} object has no attribute {item}')
        return self.base.__getattribute__(item)


class AbstractReader(Imread, metaclass=ABCMeta):
    priority = 99
    do_not_pickle = 'cache'
    ureg = ureg

    @staticmethod
    @abstractmethod
    def _can_open(path):  # Override this method, and return true when the subclass can open the file
        return False

    @abstractmethod
    def __frame__(self, c, z, t):  # Override this, return the frame at c, z, t
        return np.random.randint(0, 255, self.shape['xy'])

    @cached_property
    def ome(self):
        return self.get_ome(self.path)

    def open(self):  # Optionally override this, open file handles etc.
        """ filehandles cannot be pickled and should be marked such by setting do_not_pickle = 'file_handle_name' """
        return

    def close(self):  # Optionally override this, close file handles etc.
        return

    def __init__(self, path, dtype=None, axes=None):
        if isinstance(path, Imread):
            return
        super().__init__()
        self.isclosed = False
        if isinstance(path, str):
            path = Path(path)
        self.path, self.series = self.split_path_series(path)
        if isinstance(path, Path) and path.exists():
            self.title = self.path.name
            self.acquisitiondate = datetime.fromtimestamp(self.path.stat().st_mtime).strftime('%y-%m-%dT%H:%M:%S')
        else:  # ndarray
            self.title = self.__class__.__name__
            self.acquisitiondate = 'now'

        self.reader = None
        self.pcf = None
        self.powermode = None
        self.collimator = None
        self.tirfangle = None
        self.cyllens = ['None', 'None']
        self.duolink = 'None'
        self.detector = [0, 1]
        self.track = [0]
        self.cache = DequeDict(16)
        self.frameoffset = 0, 0  # how far apart the centers of frame and sensor are

        self.open()
        # extract some metadata from ome
        instrument = self.ome.instruments[0] if self.ome.instruments else None
        image = self.ome.images[0]
        pixels = image.pixels
        self.shape = pixels.size_x, pixels.size_y, pixels.size_c, pixels.size_z, pixels.size_t
        self.dtype = pixels.type.value if dtype is None else dtype
        self.pxsize = pixels.physical_size_x_quantity
        try:
            self.exposuretime = tuple(find(image.pixels.planes, the_c=c).exposure_time_quantity
                                      for c in range(self.shape['c']))
        except AttributeError:
            self.exposuretime = ()

        if self.zstack:
            self.deltaz = image.pixels.physical_size_z_quantity
            self.deltaz_um = None if self.deltaz is None else self.deltaz.to(self.ureg.um).m
        else:
            self.deltaz = self.deltaz_um = None
        if self.ome.images[0].objective_settings:
            self.objective = find(instrument.objectives, id=self.ome.images[0].objective_settings.id)
        else:
            self.objective = None
        try:
            t0 = find(image.pixels.planes, the_c=0, the_t=0, the_z=0).delta_t
            t1 = find(image.pixels.planes, the_c=0, the_t=self.shape['t'] - 1, the_z=0).delta_t
            self.timeinterval = (t1 - t0) / (self.shape['t'] - 1) if self.shape['t'] > 1 else None
        except AttributeError:
            self.timeinterval = None
        try:
            self.binning = [int(i) for i in image.pixels.channels[0].detector_settings.binning.value.split('x')]
            self.pxsize *= self.binning[0]
        except (AttributeError, IndexError, ValueError):
            self.binning = None
        self.channel_names = [channel.name for channel in image.pixels.channels]
        self.channel_names += [chr(97 + i) for i in range(len(self.channel_names), self.shape['c'])]
        self.cnamelist = self.channel_names
        try:
            optovars = [objective for objective in instrument.objectives if 'tubelens' in objective.id.lower()]
        except AttributeError:
            optovars = []
        if len(optovars) == 0:
            self.tubelens = None
        else:
            self.tubelens = optovars[0]
        if self.objective:
            if self.tubelens:
                self.magnification = self.objective.nominal_magnification * self.tubelens.nominal_magnification
            else:
                self.magnification = self.objective.nominal_magnification
            self.NA = self.objective.lens_na
        else:
            self.magnification = None
            self.NA = None

        self.gain = [find(instrument.detectors, id=channel.detector_settings.id).amplification_gain
                     for channel in image.pixels.channels
                     if channel.detector_settings
                     and find(instrument.detectors, id=channel.detector_settings.id).amplification_gain]
        self.laserwavelengths = [(channel.excitation_wavelength_quantity.to(self.ureg.nm).m,)
                                 for channel in pixels.channels if channel.excitation_wavelength_quantity]
        self.laserpowers = try_default(lambda: [(1 - channel.light_source_settings.attenuation,)
                                                for channel in pixels.channels], [])
        self.filter = try_default(lambda: [find(instrument.filter_sets, id=channel.filter_set_ref.id).model
                                           for channel in image.pixels.channels], None)
        self.pxsize_um = None if self.pxsize is None else self.pxsize.to(self.ureg.um).m
        self.exposuretime_s = [None if i is None else i.to(self.ureg.s).m for i in self.exposuretime]

        if axes is None:
            self.axes = ''.join(i for i in 'cztxy' if self.shape[i] > 1)
        elif axes.lower() == 'full':
            self.axes = 'cztxy'
        else:
            self.axes = axes
        self.slice = [np.arange(s, dtype=int) for s in self.shape.xyczt]

        m = self.extrametadata
        if m is not None:
            try:
                self.cyllens = m['CylLens']
                self.duolink = m['DLFilterSet'].split(' & ')[m['DLFilterChannel']]
                if 'FeedbackChannels' in m:
                    self.feedback = m['FeedbackChannels']
                else:
                    self.feedback = m['FeedbackChannel']
            except Exception:  # noqa
                self.cyllens = ['None', 'None']
                self.duolink = 'None'
                self.feedback = []
        try:
            self.cyllenschannels = np.where([self.cyllens[self.detector[c]].lower() != 'none'
                                             for c in range(self.shape['c'])])[0].tolist()
        except Exception:  # noqa
            pass
        try:
            s = int(re.findall(r'_(\d{3})_', self.duolink)[0]) * ureg.nm
        except Exception:  # noqa
            s = 561 * ureg.nm
        try:
            sigma = []
            for c, d in enumerate(self.detector):
                emission = (np.hstack(self.laserwavelengths[c]) + 22) * ureg.nm
                sigma.append([emission[emission > s].max(initial=0), emission[emission < s].max(initial=0)][d])
            sigma = np.hstack(sigma)
            sigma[sigma == 0] = 600 * ureg.nm
            sigma /= 2 * self.NA * self.pxsize
            self.sigma = sigma.magnitude.tolist()
        except Exception:  # noqa
            self.sigma = [2] * self.shape['c']
        if not self.NA:
            self.immersionN = 1
        elif 1.5 < self.NA:
            self.immersionN = 1.661
        elif 1.3 < self.NA < 1.5:
            self.immersionN = 1.518
        elif 1 < self.NA < 1.3:
            self.immersionN = 1.33
        else:
            self.immersionN = 1

        p = re.compile(r'(\d+):(\d+)$')
        try:
            self.track, self.detector = zip(*[[int(i) for i in p.findall(find(
                self.ome.images[0].pixels.channels, id=f'Channel:{c}').detector_settings.id)[0]]
                                              for c in range(self.shape['c'])])
        except Exception:
            pass


def main():
    parser = ArgumentParser(description='Display info and save as tif')
    parser.add_argument('file', help='image_file')
    parser.add_argument('out', help='path to tif out', type=str, default=None, nargs='?')
    parser.add_argument('-r', '--register', help='register channels', action='store_true')
    parser.add_argument('-c', '--channel', help='channel', type=int, default=None)
    parser.add_argument('-z', '--zslice', help='z-slice', type=int, default=None)
    parser.add_argument('-t', '--time', help='time', type=int, default=None)
    parser.add_argument('-s', '--split', help='split channels', action='store_true')
    parser.add_argument('-f', '--force', help='force overwrite', action='store_true')
    args = parser.parse_args()

    with Imread(args.file) as im:
        if args.register:
            im = im.with_transform()
        print(im.summary)
        if args.out:
            out = Path(args.out).absolute()
            out.parent.mkdir(parents=True, exist_ok=True)
            if out.exists() and not args.force:
                print(f'File {args.out} exists already, add the -f flag if you want to overwrite it.')
            else:
                im.save_as_tiff(out, args.channel, args.zslice, args.time, args.split)


from .readers import *
