"""
File Purpose: MainDimensionsChunker

Manages "chunks" options for main dimensions (see also: MainDimensionsHaver).

See docstring for MainDimensionsChunker for details.
See also: https://gitlab.com/Sevans7/plasmacalcs/-/issues/37

--- CAUTION ---
    derivative results different with/without chunks, unless deriv_before_slice=True.
    All other results should be the same with/without chunking.
"""

import xarray as xr

from .main_dimensions import MainDimensionsHaver
from ..errors import ChunkError, ChunkDimensionalityError
from ..tools import (
    UNSET,
    simple_property,
    TaskList,
)


class Chunker():
    '''manages choosing slices based on chunks.

    --- CAUTION ---
        derivative results different with/without chunks, unless deriv_before_slice=True.
        All other results should be the same with/without chunking.

    chunks: dict
        keys should be maindim possibly with '_size' or '_slices' appended, e.g. 'x' or 'y_size'.
        x --> break x into this many chunks total.
        x_size --> break x into chunks of this size.
        x_slices --> list of slices to use for x.
    maindims_sizes: dict
        {dim: full size} for dims in maindims.
    init_slices: None or dict
        dict of slices which already exist.
        un-chunked dims' slices will be maintained when creating chunk slices.
        invalid chunking if any chunk dims appear in `slices`.
        None --> internally store as empty dict.
    lenient: bool
        whether to allow any ChunkSlicer inputs, even if invalid.
        default False, i.e. self.check_valid() (which will raise ChunkError if not valid).
    '''
    def __init__(self, chunks, maindims_sizes, init_slices=None, *, lenient=False):
        self.chunks = chunks
        self.maindims_sizes = maindims_sizes
        self.init_slices = {} if init_slices is None else init_slices
        if not lenient:
            self.check_valid()

    def __repr__(self):
        return f'{type(self).__name__}({self.chunks}, maindims_sizes={self.maindims_sizes})'

    @property
    def chunk_dims(self):
        '''the tuple of dims which are being chunked.
        E.g., if chunks={'x':10, 'y_size':50}, then chunking_dims='x', 'y'.
        Computed as set of self.chunks keys, with '_size' or '_slices' removed if present.
        '''
        result = tuple({dim[:-len('_slices')] if dim.endswith('_slices') else
                        dim[:-len('_size')] if dim.endswith('_size') else dim
                        for dim in self.chunks})
        return result

    @property
    def maindims(self):
        '''the maindims. Just the tuple of self.maindims_sizes keys.'''
        return tuple(self.maindims_sizes)

    def check_valid(self):
        '''checks that self is a valid ChunkSlicer. crash with ChunkError if not.
        Valid must satisfy:
            - all chunk dims are maindims
            - no overlap between chunks and slices (e.g. 'x' in chunks and slices)
        '''
        for x in self.chunk_dims:
            if x not in self.maindims:
                raise ChunkError(f'Invalid chunk dim {x!r}; not found in maindims={self.maindims}')
            if x in self.init_slices:
                raise ChunkError(f'Cannot chunk dim {x!r}; it is already being sliced: slices={self.init_slices[x]!r}')

    def dim_chunk_info(self, x):
        '''return chunking info for dim x: (mode, value).
        mode will be 'n', 'size', or 'slices',
        correspond to x, x_size, or x_slices appearing in self.chunks.
        value will be int, int, or iterable of slices, respectively.
        '''
        if x not in self.chunk_dims:
            raise ChunkError(f'Invalid chunk dim {x!r}; not found in self.chunks={self.chunks}')
        if x in self.chunks:
            return 'n', self.chunks[x]
        elif f'{x}_size' in self.chunks:
            return 'size', self.chunks[f'{x}_size']
        elif f'{x}_slices' in self.chunks:
            return 'slices', self.chunks[f'{x}_slices']

    def _iter_slicers_x(self, x):
        '''iterates through all chunk slices for dim x.'''
        mode, chunk = self.dim_chunk_info(x)
        full_size = self.maindims_sizes[x]
        if mode == 'n' or mode == 'size':
            if mode == 'size':
                x_size = chunk
            else:
                n = chunk
                x_size = (full_size // n) + (0 if full_size % n == 0 else 1)
            for x_start in range(0, full_size, x_size):
                x_stop = min(x_start + x_size, full_size)
                yield slice(x_start, x_stop)
        elif mode == 'slices':
            for s in chunk:
                yield s
        else:
            raise ChunkError(f'Invalid chunking mode {mode!r} for dim {x!r}')

    def iter_slicers(self):
        '''iterates through all chunk slices, yielding slices dict including any non-chunked slices.

        E.g. chunks=dict(x_size=50), slices=dict(z=7) --> yield:
            dict(x=slice(0,50), z=7),
            dict(x=slice(50,100), z=7),
            ...
        '''
        chunk_dims = self.chunk_dims
        if len(chunk_dims) == 0:
            # no chunking; just yield initial slices once
            yield self.init_slices
        elif len(chunk_dims) == 1:
            # single chunk dim; simple loop
            x = chunk_dims[0]
            for xslice in self._iter_slicers_x(x):
                slices = self.init_slices.copy()
                slices[x] = xslice
                yield slices
        elif len(chunk_dims) == 2:
            # two chunk dims, also sort of simple to loop:
            x, y = chunk_dims
            for xslice in self._iter_slicers_x(x):
                for yslice in self._iter_slicers_x(y):
                    slices = self.init_slices.copy()
                    slices[x] = xslice
                    slices[y] = yslice
                    yield slices
        else:
            # could easily extend pattern to 3D+, but should generalize if doing that.
            # but, probably not necessary!
            raise NotImplementedError('[TODO] chunking along 3 or more dims')

    def slicers(self):
        '''list of all slices to use for chunking. Equivalent: list(self.iter_slicers()).'''
        return list(self.iter_slicers())

    def __len__(self):
        '''the number of chunk slices which will be generated by self.iter_slicers().
        [EFF] perhaps should make this more efficient instead of generating all slices...
        '''
        return len(self.slicers())


class MainDimensionsChunker(MainDimensionsHaver):
    '''Manages "chunks" options for main dimensions.'

    --- CAUTION ---
        derivative results different with/without chunks, unless deriv_before_slice=True.
        All other results should be the same with/without chunking.

    Here is the idea behind chunks:
        For data with large arrays along maindims, sometimes arrays are
        so large that loading the whole thing into memory slows down computations a lot.
        Also, some computations are local, depending on no other cells or only nearby cells.

        In those cases, it is probably much faster to break up the maindims into chunks,
        load one chunk at a time, and combine results appropriately.

    The internal implementation goes something like this:
        - chunks will be a behavior attribute (set it via self.chunks=val, or sel(var, chunks=val)).
        - if chunks is set, hijack the self.__call__ to instead do chunking.
        - doing chunking means:
            - determining a list of appropriate slices along maindims, based on chunks
            - getting results for each slice (while getting those results, disable chunking)
            - combining the results appropriately

    Currently, chunks can only apply to vars which keep the chunked axis or axes.
    So, for example, if maindims=('x', 'y', 'z') and 'u' has spatial extent:
        - self('mod_u') can be chunked
        - self('mean_u') cannot be chunked;
            mean_var takes the mean across maindims, so result has no spatial extent.
        - self('fft_u') cannot be chunked;
            fft_var converts dims (default: fft_dims=maindims) into frequency dims.
    Internally, there is no check ahead of time to see if a var can be chunked.
    Instead, chunking will be attempted, and then will check if the chunked axes remain;
        if not, will crash with ChunkDimensionalityError.
    '''

    # # # CHUNKS DEFINITIONS # # #
    chunker_cls = Chunker

    cls_behavior_attrs.register('chunks', default={})
    chunks = simple_property('_chunks', setdefault=dict,
        doc='''chunks for maindims when loading arrays.
        E.g. chunks = dict(x=10, y_size=50)
            --> break x into 10 chunks total, and break y into chunks of size 50.
                (if y not divisible by 50, the last chunk will be smaller.)
        Can also provide a list of indexers for more low-level control,
        E.g. chunks = dict(x=[slice(0,50), 100, slice(800,900)])
            --> repeat calculation with x from 0 to 50, then 100th value, then 800 to 900.
            result will be a well-labeled xarray with x at those coordinates,
            here equivalent to result from slices=dict(x=[0,1,...,49,100,800,801,...,899])
        See help(self.chunker_cls) or help(self.chunker()) for more details.

        Notes:
            - mutually exclusive with providing slices of the same dims.
                E.g., cannot provide x chunk and x slice at the same time.
            - if provided, internally hijacks self.__call__ to do chunking instead,
                and inside of that will loop through calls to self with different slices.
            - only works for vars which keep the chunked axes.
                E.g., fails for self('mean_n') since mean result doesn't have maindims.''')
    @chunks.setter
    def chunks(self, value):
        if value is None:
            value = dict()
        self._chunks = value

    def chunker(self):
        '''gets Chunker based on the current values of self.chunks and self.slices.
        see help on MainDimensionsChunker, or help(type(self).chunks) for more details.
        '''
        return self.chunker_cls(self.chunks, self.maindims_full_sizes, self.slices)

    # # # CALLING SELF, WITH CHUNKS # # #
    def _call_hijacker(self, *args_super, **kw_super):
        '''returns False or name of hijacker method to use instead of self(var) call.
        Here, returns '_get_with_chunks' if self.chunks is nonempty, else super()._call_hijacker(...).
        '''
        if len(self.chunks) > 0:
            return '_get_with_chunks'
        else:
            return super()._call_hijacker(*args_super, **kw_super)

    def _get_with_chunks(self, var, *args_super, **kw_super):
        '''returns value of var, but chunking while loading, then merging chunked results.
        Result should be equivalent to self(var, chunks=None).
        CAUTION: derivatives may give wrong answers at chunk boundaries, unless deriv_before_slices=False
        '''
        chunker = self.chunker()
        chunkdims = chunker.chunk_dims
        if len(chunkdims) != 1:
            raise NotImplementedError('[TODO] implement chunking when len(chunk_dims) != 1')
        ## IMPLEMENTATION 1: naive (single-dim only, and no multiprocessing) implementation
        # (implementation 1 added in commit 55ac709. Removed since then.)

        ## IMPLEMENTATION 2: slightly less naive (single-dim only, but includes multiprocessing):
        chunkdim = chunkdims[0]
        slicers = chunker.slicers()
        task_inputs = [
            (self, (var, *args_super), {**kw_super, 'slices': slicer})
            for slicer in slicers
        ]
        task_inputs[0][2]['_chunks_check0'] = (chunkdim, self.call_depth + 1)  # will check if chunking allowed
        tasks = TaskList(*task_inputs)
        mp_strategy = dict(ncpu=self.ncpu,   # get multiprocessing preferences from self
                   timeout=self.timeout,
                   ncoarse=self.ncoarse,
                   print_freq=self.print_freq_explicit)
        with self.using(chunks=None, ncpu=1):  # internally, don't use chunking nor multiprocessing.
            results = tasks(**mp_strategy)
        return xr.concat(results, dim=chunkdim)

    # # # CHUNK CHECKS (BOOKKEEPING) # # #
    @property
    def _extra_kw_for_quantity_loader_call(self):
        '''extra kwargs which can be used to set attrs self during self.__call__.
        The implementation here returns ['_chunks_check0'] + any values from super().
        '''
        return ['_chunks_check0'] + super()._extra_kw_for_quantity_loader_call

    _chunks_check0 = simple_property('_chunks_check0_', default=None,
        doc='''None or (dim, call_depth) --> assert self(...) first chunk has dim at this call depth.
        E.g., _chunks_check0 = ('x', 3) --> assert 'x' in result.dims at call_depth=3,
            and if it isn't, raise ChunkDimensionalityError.''')

    def _call_postprocess(self, result, *, var, name=UNSET, item=UNSET):
        '''postprocess result from self.__call__. Called during self.__call__.
        (self.call_depth inside here will tell depth of the current call; depth=1 for top level.)

        result: any value, probably an xarray.DataArray
            result from self.__call__, before postprocessing.
        var, name, item: UNSET or value
            passed directly from self.__call__.

        The implementation here does the following (subclasses might override / add to this):
            (A) call super()._call_postprocess(), which probably (if parent is QuantityLoader) does:
                (1) if self.verbose >= 4, print a message about getting var.
                (2) result = self.attach_extra_coords(result).
                (3) set result.name = name, or result.attrs['name'] = name, if possible.
                (4) if self.assign_behavior_attrs at this call depth, do so now.
                (5) if self.call_depth == 1, call self._call_postprocess_toplevel.
                (6) if item, convert result to single value via result.item().
            (B) if self._chunks_check0 is not None, crash if chunkdim is missing from result.dims.
                else, return result, unchanged.
        '''
        result = super()._call_postprocess(result, var=var, name=name, item=item)
        if self._chunks_check0 is not None:
            x, depth = self._chunks_check0
            # depth is needed so that we don't do the checks internally (at larger depths),
            #   e.g. imagine var='5*n', we shouldn't do the check for '5', only for '5*n'.
            if self.call_depth == depth:
                if x not in getattr(result, 'dims', []):
                    errmsg = (f"Cannot chunk {var!r} along {x!r}. "
                              f"getting var {var!r} removes dim {x!r}; "
                              f"first chunk's result had dims={result.dims})")
                    raise ChunkDimensionalityError(errmsg)
        return result
