# -*- coding: utf-8 -*-
import os
import fnmatch
import copy
import datetime
import pickle
import pathlib

import numpy as np
from numpy.lib.format import open_memmap
import mpmath
import PIL
import PIL.PngImagePlugin
import numba

import fractalshades as fs
import fractalshades.numpy_utils.xrange as fsx
import fractalshades.settings as fssettings
import fractalshades.utils as fsutils

from fractalshades.mprocessing import Multiprocess_filler

class _Pillow_figure:
    def __init__(self, img, pnginfo):
        """
        This class is a wrapper that can be used to redirect a Fractal_plotter
        output, for instance when generating the documentation.
        """
        self.img = img
        self.pnginfo = pnginfo

    def save_png(self, im_path):
        """
        Saves as a png with Lanczos antialiasing if exceeds the max width
        """
        im = self.img
        width, height = im.size
        max_width = fs.settings.output_context["doc_max_width"]

        if width > max_width:
            ratio = float(width) / float(max_width)
            new_height = int(height / ratio)
            im = im.resize((max_width, new_height), PIL.Image.LANCZOS)
        im.save(im_path, format="png", pnginfo=self.pnginfo)

class Fractal_plotter:
    def  __init__(self, postproc_batch):
        """
        The base plotting class.

        A Fractal plotter :

        - points to a single Fractal
        - can hold several `fractalshades.postproc.Postproc_batch` each one
          point to a single
          (`Fracal`, ``calculation``) couple
        - can hold several post-processing layers

        Parameters
        ----------
        postproc_batch
            A single `fractalshades.postproc.Postproc_batch` or a list of
            these

        Notes
        -----

        .. warning::
            When passed a list of `fractalshades.postproc.Postproc_batch`
            objects, each one shall point to
            the same `Fractal`.
        """
        # postproc_batchescan be single or an enumeration
        postproc_batches = postproc_batch
        if isinstance(postproc_batch, fs.postproc.Postproc_batch):
            postproc_batches = [postproc_batch]
        for i, postproc_batch in enumerate(postproc_batches):
            if i == 0:
                self.postproc_batches = [postproc_batch]
                self.posts = copy.copy(postproc_batch.posts)
                self.postnames_2d = postproc_batch.postnames_2d
                # iterator :
                self.fractal = postproc_batch.fractal
            else:
                self.add_postproc_batch(postproc_batch)

        self.chunk_slices = self.fractal.chunk_slices
        # layer data
        self.layers = []
        self.scalings = None # to be computed !
        # Plotting directory
        self.plot_dir = self.fractal.directory

    @property
    def postnames(self):
        return self.posts.keys()

    @property
    def size(self):
        f = self.fractal
        return (f.nx, f.ny)

    def add_postproc_batch(self, postproc_batch):
        """
        Adds another post-processing batch & register the associated postproc
        names. `postproc_batch` shall map to the same fractal.
        Note : several postproc_batches are needed whenever different
        calculations need to be combined in an output plot, as a postproc
        batch can only map to a unique calc_name.
        :meta private:
        """
        if postproc_batch.fractal != self.fractal:
            raise ValueError("Attempt to add a postproc_batch from a different"
                             "fractal: {} from {}".format(
                postproc_batch.fractal, postproc_batch.calc_name))
        self.postproc_batches += [postproc_batch]
        self.posts.update(postproc_batch.posts)
        self.postnames_2d += postproc_batch.postnames_2d

    def add_layer(self, layer):
        """
        Adds a layer field to allow subsequent graphical operations or outputs.

        Parameters
        ----------
        layer : `fractalshades.colors.layers.Virtual_layer` or a derived class
            The layer to add

        Notes
        -----

        .. warning::
            Layer `postname` shall have been already registered in one of the
            plotter batches.

        .. warning::
            When a layer is added, a link layer -> Fractal_plotter is
            created ; a layer can only point to a single `Fractal_plotter`.
        """
        postname = layer.postname
        if postname not in (list(self.postnames) + self.postnames_2d):
            raise ValueError("Layer `{}` shall be registered in "
                             "Fractal_plotter postproc_batches: {}".format(
            postname, list(self.postnames) + self.postnames_2d))
        self.layers += [layer]
        layer.link_plotter(self)

    def __getitem__(self, layer_name):
        """ Get the layer by its postname
        """
        for layer in self.layers:
            if layer.postname == layer_name:
                return layer
        raise KeyError("Layer {} not in available layers: {}".format(
                layer_name, list(l.postname for l in self.layers)))

    def plot(self):
        """
        The base method to produce images.

        When called, it will got through all the instance-registered layers
        and plot each layer for which the `output` attribute is set to `True`.
        """
        self.store_postprocs()
        self.compute_scalings()
        self.write_postproc_report()
        self.open_images()
        self.push_layers_to_images()
        self.save_images()
        self.clean_up()

    def store_postprocs(self):
        """ Computes and stores posprocessed data in a temporary mmap
        :meta private:
        """
        if fs.settings.optimize_RAM:
            self.has_memmap = True
            self.open_temporary_mmap()
        else:
            self.has_memmap = False
            self.open_RAM_data()

        inc_posproc_rank = 0
        self.postname_rank = dict()

        for batch in self.postproc_batches:
            if self.has_memmap:
                self.store_temporary_mmap(
                        chunk_slice=None,
                        batch=batch,
                        inc_posproc_rank=inc_posproc_rank
                )
            else:
                self.store_data(
                        chunk_slice=None,
                        batch=batch,
                        inc_posproc_rank=inc_posproc_rank
                )
            for i, postname in enumerate(batch.postnames):
                self.postname_rank[postname] = inc_posproc_rank + i
            inc_posproc_rank += len(batch.posts)

    def temporary_mmap_path(self):
        """ Path to the temporary memmap used to stored plotting arrays"""
        # from tempfile import mkdtemp
        return os.path.join(
            self.fractal.directory, "data", "_plotter.tpm")

    def open_temporary_mmap(self):
        """
        Creates the memory mappings for postprocessed arrays
        Note: We expand to 2d format
        """
        f = self.fractal
        nx, ny = (f.nx, f.ny)
        n_pp = len(self.posts)
        mmap = open_memmap(
            filename=self.temporary_mmap_path(),
            mode='w+',
            dtype=f.float_postproc_type,
            shape=(n_pp, nx, ny),
            fortran_order=False,
            version=None)
        del mmap

    def open_RAM_data(self):
        """
        Same as open_temporary_mmap but in RAM
        """
        f = self.fractal
        nx, ny = (f.nx, f.ny)
        n_pp = len(self.posts)
        self._RAM_data = np.zeros(
            shape=(n_pp, nx, ny),
            dtype=f.float_postproc_type
        )

    @Multiprocess_filler(iterable_attr="chunk_slices",
        iter_kwargs="chunk_slice", veto_multiprocess=False)
    def store_temporary_mmap(self, chunk_slice, batch, inc_posproc_rank):
        """ Compute & store temporary arrays for this postproc batch
            Note : inc_posproc_rank rank shift to take into account potential
            other batches for this plotter.
            (memory mapping version)
        """
        f = self.fractal
        inc = inc_posproc_rank
        post_array, chunk_mask = self.fractal.postproc(batch, chunk_slice)
        arr_2d = f.reshape2d(post_array, chunk_mask, chunk_slice)
        n_posts, cx, cy = arr_2d.shape
        (ix, ixx, iy, iyy) = chunk_slice
        mmap = open_memmap(filename=self.temporary_mmap_path(), mode='r+')
        mmap[inc:inc+n_posts, ix:ixx, iy:iyy] = arr_2d

    @Multiprocess_filler(iterable_attr="chunk_slices",
        iter_kwargs="chunk_slice", veto_multiprocess=True)
    def store_data(self, chunk_slice, batch, inc_posproc_rank):
        """ Compute & store temporary arrays for this postproc batch
            (in-RAM version -> shall not use multiprocessing)
        """
        f = self.fractal
        inc = inc_posproc_rank
        post_array, chunk_mask = self.fractal.postproc(batch, chunk_slice)
        arr_2d = f.reshape2d(post_array, chunk_mask, chunk_slice)
        n_posts, cx, cy = arr_2d.shape
        (ix, ixx, iy, iyy) = chunk_slice
        self._RAM_data[inc:inc+n_posts, ix:ixx, iy:iyy] = arr_2d

    # All methods needed for plotting
    def compute_scalings(self):
        """ Compute the scaling for all layer field
        (needed for mapping to color) """
        for layer in self.layers:
            self.compute_layer_scaling(chunk_slice=None, layer=layer)

    def write_postproc_report(self):
        report_path = os.path.join(
            self.fractal.directory, type(self).__name__ + ".txt")

        def write_layer_report(i, layer, report):
            report.write(" - Layer #{} :\n".format(i))
            postname = layer.postname
            report.write("post-processing: `{}`\n".format(postname))
            report.write("kind: {}\n".format(type(layer).__name__))
            report.write("func: {}\n".format(layer._func_arg))

            mask = layer.mask
            if mask is None:
                mask_str = "None"
            else:
                mask_str = "{} `{}` with mask color: {}".format(
                    type(mask[0]).__name__,
                    mask[0].postname,
                    mask[1])
            report.write("mask: {}\n".format(mask_str))

            report.write("output: {}\n".format(layer.output))
            report.write("min: {}\n".format(layer.min))
            report.write("max: {}\n\n".format(layer.max))

        with open(report_path, 'w', encoding='utf-8') as report:
            for i, layer in enumerate(self.layers):
                write_layer_report(i, layer, report)

    @Multiprocess_filler(iterable_attr="chunk_slices",
        iter_kwargs="chunk_slice", veto_multiprocess=True)
    def compute_layer_scaling(self, chunk_slice, layer):
        """ Compute the scaling for this layer """
        layer.update_scaling(chunk_slice)

    def open_images(self):
        self._im = []
        for layer in self.layers:
            if layer.output:
                self._im += [PIL.Image.new(mode=layer.mode, size=self.size)]
            else:
                self._im += [None]

    def save_images(self):
        for i, layer in enumerate(self.layers):
            if not(layer.output):
                continue
            file_name = "{}_{}".format(type(layer).__name__, layer.postname)
            base_img_path = os.path.join(self.plot_dir, file_name + ".png")
            self.save_tagged(self._im[i], base_img_path, self.fractal.params)

    def save_tagged(self, img, img_path, tag_dict):
        """
        Saves *img* to png format at *path*, tagging with *tag_dict*.
        https://dev.exiv2.org/projects/exiv2/wiki/The_Metadata_in_PNG_files
        """
        pnginfo = PIL.PngImagePlugin.PngInfo()
        for k, v in tag_dict.items():
            pnginfo.add_text(k, str(v))
        if (fssettings.output_context["doc"]
            and not(fssettings.output_context["gui_iter"] > 0)):
            fssettings.add_figure(_Pillow_figure(img, pnginfo))
        else:
            img.save(img_path, pnginfo=pnginfo)

    def push_layers_to_images(self):
        for i, layer in enumerate(self.layers):
            if not(layer.output):
                continue
            self.push_cropped(chunk_slice=None, layer=layer, im=self._im[i])

    @Multiprocess_filler(iterable_attr="chunk_slices",
        iter_kwargs="chunk_slice", veto_multiprocess=True)
    def push_cropped(self, chunk_slice, layer, im):
        """ push "cropped image" from layer for this chunk to the image"""
        (ix, ixx, iy, iyy) = chunk_slice
        ny = self.fractal.ny
        crop_slice = (ix, ny-iyy, ixx, ny-iy)
        paste_crop = layer.crop(chunk_slice)
        im.paste(paste_crop, box=crop_slice)

    def clean_up(self):
        if self.has_memmap:
            os.unlink(self.temporary_mmap_path())
        else:
            del self._RAM_data

class Fractal:

    REPORT_ITEMS = [
        "chunk1d_begin",
        "chunk1d_end",
        "iref",
        "glitch_max_attempt",
        "chunk_pts",
        "total-glitched",
        "dyn-glitched"]

    # Note : chunk_mask is pre-computed and saved also but not at the same
    # stage (at begining of calculation)
    SAVE_ARRS = [
        "Z",
        "U",
        "stop_reason",
        "stop_iter"]

    def __init__(self, directory):
        """
The base class for all escape-time fractals calculations.

Derived class should implement the actual calculation methods used in the
innner loop. This class provides the outer looping (calculation is run on
successive tiles), enables multiprocessing, and manage raw-result storing
and retrieving.

Parameters
----------
directory : str
    Path for the working base directory

Attributes
----------
directory : str
    the working directory
iref : int
    the reference point index, when using perturbation technique (if
    not, shall be None)
glitch_max_attempt : int
    The maximal number of glitch correction loops
subset
    A boolean array-like of the size of the image, when False the
    calculation is skipped for this point. It is usually the result
    of a previous calculation that can be passed via a Fractal_array
    wrapper.
    If None (default) all points will be calculated.
glitch_stop_index : int
    A calculation can exit for several reasons which are tracked in
    `stop_reason` int-array (see below the description of the raw data
    arrays). Values of `stop_reason` above `glitch_stop_index` indicate
    glitched pixels, for which the calculation is deemed invalid.
complex_type :
    the datatype used for Z output arrays
codes :
    the string identifier codes for the saved arrays:
    (`complex_codes`, `int_codes`, `termination_codes`)

Notes
-----

These notes describe implementation details and should be useful mostly to
advanced users when subclassing.

.. note::

    **Special methods**

    this class and its subclasses may define several methods decorated with
    specific tags:

    `fractalshades.zoom_options`
        decorates the methods used to define the zoom
    `fractalshades.calc_options`
        decorates the methods defining the calculation inner-loop
    `fractalshades.interactive_options`
        decorates the methods that can be called
        interactively from the GUI (right-click then context menu selection).
        The coordinates of the click are passed to the called method.

.. note::

    **Calculation parameters**

    To lanch a calculation, call `~fractalshades.Fractal.run`. The parameters
    from the last
    @ `fractalshades.zoom_options` call and last
    @ `fractalshades.calc_options` call will be used.
    They are stored as class attributes, above a list of such attributes and
    their  meaning (non-exhaustive as derived class cmay define their own).
    Note that they should normally not be directly accessed but defined in
    derived class through zoom and calc methods.

.. note::

    **Saved data**

        The calculation results (raw output of the inner loop at exit) are
        saved to disk and internally accessed during plotting phase through
        memory-mapping. These are:

        chunk_mask
            boolean - alias for `subset`
            Saved to disk as ``calc_name``\_Z.arr in ``data`` folder
        Z
            Complex fields, several fields can be defined and accessed through
            a field string identifier.
            Saved to disk as ``calc_name``\_Z.arr in ``data`` folder
        U
            Integer fields, several fields can be defined and accessed through
            a field string identifier.
            Saved to disk as ``calc_name``\_U.arr in ``data`` folder
        stop_reason
            Byte codes: the reasons for loop exit (max iteration reached ?
            overflow ? other ?) A string identifier
            Saved to disk as ``calc_name``\_stop_reason.arr in ``data`` folder
        stop_iter
            Integer: iterations count at loop exit
            Saved to disk as ``calc_name``\_stop_iter.arr in ``data`` folder

        The string identifiers are stored in ``codes`` attributes.
"""
        self.directory = directory
        self.iref = None # None when no reference point used / needed
        self.glitch_max_attempt = 0
        self.subset = None
        self.glitch_stop_index = None

    def init_data_types(self, complex_type):
        if type(complex_type) is tuple:
            type_modifier, _ = complex_type
            if type_modifier != "Xrange":
                raise ValueError(type_modifier)
        self.complex_type = complex_type
        self.float_postproc_type = np.float32
        self.termination_type = np.int8
        self.int_type = np.int32

    @fsutils.zoom_options
    def zoom(self, *,
             x: float,
             y: float,
             dx: float,
             nx: int,
             xy_ratio: float,
             theta_deg: float,
             projection: str="cartesian",
             antialiasing: bool=False):
        """
        Define and stores as class-attributes the zoom parameters for the next
        calculation.

        Parameters
        ----------
        x : float
            x-coordinate of the central point
        y : float
            y-coordinate of the central point
        dx : float
            span of the view rectangle along the x-axis
        nx : int
            number of pixels of the image along the x-axis
        xy_ratio: float
            ratio of dx / dy and nx / ny
        theta_deg : float
            Pre-rotation of the calculation domain, in degree
        projection : "cartesian" | "spherical" | "exp_map"
            Kind of projection used (default to cartesian)
        antialiasing : bool
            If True, some degree of randomization is applied
        """
        # We're all set, the job is done by `zoom_options` wrapper...

    def run(self):
        """
        Lauch a full calculation.

        The parameters from the last
        @ `fractalshades.zoom_options`\-tagged method call and last
        @ `fractalshades.calc_options`\-tagged method call will be used.

        If calculation results are already there, the parameters will be
        compared and if identical, the calculation will be skipped. This is
        done for each tile and each glitch correction iteration, so i enables
        calculation to restart from an unfinished status.
        """
        if not(self.res_available()):
            # We write the param file and initialize the
            # memmaps for progress reports and calc arrays
            # It is not process safe so we dot it before entering multi-processing
            # loop
            fsutils.mkdir_p(os.path.join(self.directory, "data"))
            self.open_report_mmap()
            self.open_data_mmaps()
            self.save_params()

        # Lazzy compilation of subset boolean array chunk-span
        self._mask_beg_end = None

        # JIT-compiled function
        # self.jitted_numba_cycles = numba.njit(numba_cycles)
        self._iterate = self.iterate()
        self.cycles()

        # Export to human-readable format
        if fs.settings.inspect_calc:
            self.inspect_calc()

    @property
    def interrupt_path(self):
        return os.path.join(self.directory, "data", "_interrupt.lck")

    def raise_interruption(self):
        pathlib.Path(self.interrupt_path).touch()

    def lower_interruption(self):
        if os.path.isfile(self.interrupt_path):
            os.unlink(self.interrupt_path)

    def is_interrupted(self):
        """ Either programmatically 'interrupted' (from the GUI) or by the user
        in batch mode through fs.settings.skip_calc """
        return (os.path.isfile(self.interrupt_path)
                or fs.settings.skip_calc)

    @property
    def ny(self):
        return int(self.nx / self.xy_ratio + 0.5)

    @property
    def dy(self):
        return self.dx / self.xy_ratio

    @property
    def px(self):
        if not(hasattr(self, "_px")): # is None:
            self._px = self.dx / self.nx
        return self._px

    @property
    def multiprocess_dir(self):
        """ Directory used for multiprocess stdout stderr streams redirection
        :meta private:
        """
        return os.path.join(self.directory, "multiproc_calc")

    @property
    def Xrange_complex_type(self):
        """ Return True if the data type is a xrange array
        :meta private:
        """
        if type(self.complex_type) is tuple:
            type_modifier, _ = self.complex_type
            return type_modifier == "Xrange"
        return False

    @property
    def base_complex_type(self):
        complex_type = self.complex_type
        if type(complex_type) is tuple:
            _, complex_type = complex_type
        return complex_type

    @property
    def base_float_type(self):
        select = {np.dtype(np.complex64): np.float32,
                  np.dtype(np.complex128): np.float64}
        return select[np.dtype(self.base_complex_type)]

    @property
    def params(self):
        """ Used to tag an output image or check if data is already computed
        and stored
        :meta private:
        """
        software_params = {
                "Software": "fractalshades " + fs.__version__,
                "fractal_type": type(self).__name__,
                # "debug": ("1234567890" * 10), # tested 10000 chars ok
                "datetime": datetime.datetime.today().strftime(
                        '%Y-%m-%d_%H:%M:%S')}
        zoom_params = self.zoom_options
        calc_function = self.calc_options_lastcall # TODO rename to calc_callable
        calc_params = self.calc_options

        res = dict(software_params)
        res.update(zoom_params)
        res["calc-function"] = calc_function
        res.update({"calc-param_" + k: v for (k, v) in calc_params.items()})

        return res

    def clean_up(self, calc_name):
        """
        Deletes all saved data files associated with a given ``calc_name``.

        Parameters
        ----------
        calc_name : str
            The string identifying the calculation run for which we want to
            delete the files.
        """
        for pattern in [
                calc_name + "_*.arr",
                calc_name + ".params",
                calc_name + ".report",
                calc_name + "_pt*.ref",
                calc_name + "_pt*.sa"
        ]:
            data_dir = os.path.join(self.directory, "data")
            if not os.path.isdir(data_dir):
                return
            with os.scandir(data_dir) as it:
                for entry in it:
                    if (fnmatch.fnmatch(entry.name, pattern)):
                        os.unlink(entry.path)

    @property
    def pts_count(self):
        """ Return the total number of points for the current calculation
        taking into account the `subset` parameter
        :meta private:
        """
        if self.subset is not None:
#            print("in pts_count", self.subset)
#            print("in pts_count", np.sum(self.subset[None]),
#                  np.count_nonzero(self.subset[None]))
            return np.count_nonzero(self.subset[None])
        else:
            return self.nx * self.ny

    # The various method associated with chunk mgmt ===========================
    def chunk_slices(self): #, chunk_size=None):
        """
        Generator function
        Yields the chunks spans (ix, ixx, iy, iyy)
        with each chunk of size chunk_size x chunk_size
        """
        # if chunk_size is None:
        chunk_size = fssettings.chunk_size

        for ix in range(0, self.nx, chunk_size):
            ixx = min(ix + chunk_size, self.nx)
            for iy in range(0, self.ny, chunk_size):
                iyy = min(iy + chunk_size, self.ny)
                yield  (ix, ixx, iy, iyy)

    @property
    def chunks_count(self):
        """
        Return the total number of chunks (tiles) for the current image
        """
        chunk_size = fssettings.chunk_size
        (cx, r) = divmod(self.nx, chunk_size)
        if r != 0:
            cx += 1
        (cy, r) = divmod(self.ny, chunk_size)
        if r != 0: cy += 1
        return cx * cy

    def chunk_rank(self, chunk_slice):
        """
        Return the generator yield index for chunk_slice
        """
        chunk_size = fssettings.chunk_size
        (ix, _, iy, _) = chunk_slice
        chunk_item_x = ix // chunk_size
        chunk_item_y = iy // chunk_size
        (cy, r) = divmod(self.ny, chunk_size)
        if r != 0: cy += 1
        return chunk_item_x * cy + chunk_item_y

    def chunk_from_rank(self, rank):
        """
        Return the chunk_slice from the generator yield index
        """
        chunk_size = fssettings.chunk_size
        (cy, r) = divmod(self.ny, chunk_size)
        if r != 0: cy += 1
        chunk_item_x, chunk_item_y = divmod(rank, cy)
        ix = chunk_item_x * chunk_size
        iy = chunk_item_y * chunk_size
        ixx = min(ix + chunk_size, self.nx)
        iyy = min(iy + chunk_size, self.ny)
        return (ix, ixx, iy, iyy)

    def mask_beg_end(self, rank):
        """ Return the span for the boolean mask index for the chunk index
        `rank` """
        if self._mask_beg_end is None:
            arr = np.empty((self.chunks_count + 1,), dtype=np.int32)
            arr[0] = 0
            for i, chunk_slice in enumerate(self.chunk_slices()):
                (ix, ixx, iy, iyy) = chunk_slice
                arr[i + 1] = arr[i] + (ixx - ix) * (iyy - iy)
            self._mask_beg_end = arr
        return self._mask_beg_end[rank: rank + 2]

    def chunk_pts(self, chunk_slice):
        """
        Return the number of compressed 1d points for this chunk_slice
        (taking into account 2d subset bool if available)
        """
        subset = self.subset
        (ix, ixx, iy, iyy) = chunk_slice
        if subset is not None:
            subset_pts = np.count_nonzero(subset[chunk_slice]) # TODO : test this
            return subset_pts
        else:
            return (ixx - ix) * (iyy - iy)

    @property
    def chunk_mask(self):#, chunk_slice):
        """ Legacy - simple alias """
        return self.subset
#        subset = self.subset
#        if subset is not None:
#            return ~subset # np.ravel(subset[chunk_slice])
#        else:
#            return None

    def c_chunk(self, chunk_slice):
        """
        Returns a chunk of c_vec for the calculation
        Parameters
         - chunk_span
         - data_type: expected one of np.float64, np.longdouble

        Returns:
        c_vec : [chunk_size x chunk_size] 2d-vec of type datatype

        Projection availables cases :
            - cartesian : standard cartesisan
            - spherical : uses a spherical projection
            - exp_map : uses an exponential map projection
            - mixed_exp_map :  a mix of cartesian, exponential

        Note : return type is always standard prec - standard range
        """

        (x, y)  = (self.x, self.y)

        offset = self.chunk_offset(chunk_slice)
        return (x + offset[0]) + (y + offset[1]) * 1j # TODO test this

    def chunk_offset(self, chunk_slice, ensure_Xr=False):
        """
        Only computes the delta around ref central point for different projections
        Note : return type is always standard prec - standard or extended range

        ensure_Xr : enforce extended range if True
        """
#        select = {np.complex256: np.float128,
#                  np.complex128: np.float64}
        data_type = self.base_float_type # select[self.base_complex_type]

        (xy_ratio, theta_deg)  = (self.xy_ratio, self.theta_deg)
        (nx, ny, dx, dy) = (self.nx, self.ny, self.dx, self.dy)

        if self.Xrange_complex_type or ensure_Xr:
            dx_m, dx_exp = mpmath.frexp(dx)
            dx_m = np.array(dx_m, data_type)
            dx = fsx.Xrange_array(dx_m, dx_exp)
            dy_m, dy_exp = mpmath.frexp(dy)
            dy_m = np.array(dy_m, data_type)
            dy = fsx.Xrange_array(dy_m, dy_exp)
        else:
            dx = float(dx)
            dy = float(dy)

        (ix, ixx, iy, iyy) = chunk_slice
        theta = theta_deg / 180. * np.pi

        dx_grid = np.linspace(-0.5, 0.5, num=nx, dtype=data_type)
        dy_grid = np.linspace(-0.5, 0.5, num=ny, dtype=data_type)
        dy_vec, dx_vec  = np.meshgrid(dy_grid[iy:iyy], dx_grid[ix:ixx])

        if self.antialiasing:
            rg = np.random.default_rng(0)
            dx_vec += (0.5 - rg.random(dx_vec.shape, dtype=data_type)) * 0.5 / nx
            dy_vec += (0.5 - rg.random(dy_vec.shape, dtype=data_type)) * 0.5 / ny

        dx_vec = dx * dx_vec
        dy_vec = dy * dy_vec

        if self.projection == "cartesian":
            offset = [(dx_vec * np.cos(theta)) - (dy_vec * np.sin(theta)),
                      (dx_vec * np.sin(theta)) + (dy_vec * np.cos(theta))]

        elif self.projection == "spherical":
            dr_sc = np.sqrt(dx_vec**2 + dy_vec**2) / max(dx, dy) * np.pi
            k = np.where(dr_sc >= np.pi * 0.5, np.nan,  # outside circle
                         np.where(dr_sc < 1.e-12, 1., np.tan(dr_sc) / dr_sc))
            dx_vec *= k
            dy_vec *= k
            offset = [(dx_vec * np.cos(theta)) - (dy_vec * np.sin(theta)),
                      (dx_vec * np.sin(theta)) + (dy_vec * np.cos(theta))]

        elif self.projection == "mixed_exp_map":
            # square + exp. map
            h_max = 2. * np.pi * xy_ratio # max h reached on the picture
            xbar = (dx_vec + 0.5 * dx - dy) / dx * h_max # 0 .. hmax
            ybar = dy_vec / dy * 2. * np.pi              # -pi .. +pi
            rho = dx * 0.5 * np.where(xbar > 0., np.exp(xbar), 0.)
            phi = ybar + theta
            dx_vec = (dx_vec + 0.5 * dx - 0.5 * dy) * xy_ratio
            dy_vec = dy_vec * xy_ratio
            offset = [np.where(xbar <= 0.,
                          (dx_vec * np.cos(theta)) - (dy_vec * np.sin(theta)),
                          rho * np.cos(phi)),
                      np.where(xbar <= 0.,
                          (dx_vec * np.sin(theta)) + (dy_vec * np.cos(theta)),
                          rho * np.sin(phi))]

        elif self.projection == "exp_map":
            # only exp. map
            h_max = 2. * np.pi * xy_ratio # max h reached on the picture
            xbar = (dx_vec + 0.5 * dx - dy) / dx * h_max # 0 .. hmax
            ybar = dy_vec / dy * 2. * np.pi              # -pi .. +pi
            rho = dx * 0.5 * np.exp(xbar)
            phi = ybar + theta
            offset = [rho * np.cos(phi), rho * np.sin(phi)]

        else:
            raise ValueError("Projection not implemented: {}".format(
                              self.projection))
        return offset

    def px_chunk(self, chunk_slice):
        """
        Local size of pixel for different projections
        """
        data_type = self.base_float_type

        xy_ratio  = self.xy_ratio
        (nx, ny, dx, dy) = (self.nx, self.ny, self.dx, self.dy)
        (ix, ixx, iy, iyy) = chunk_slice

        if not(self.Xrange_complex_type):
            dx = float(dx)
            dy = float(dy)
        else:
            dx_m, dx_exp = mpmath.frexp(dx)
            dx_m = np.array(dx_m, data_type)
            dx = fsx.Xrange_array(dx_m, dx_exp)
            dy_m, dy_exp = mpmath.frexp(dy)
            dy_m = np.array(dy_m, data_type)
            dy = fsx.Xrange_array(dy_m, dy_exp)

        dx_grid = dx * np.linspace(-0.5, 0.5, num=nx, dtype=data_type)
        dy_grid = dy * np.linspace(-0.5, 0.5, num=ny, dtype=data_type)
        dy_vec, dx_vec  = np.meshgrid(dy_grid[iy:iyy], dx_grid[ix:ixx])

        if self.projection == "cartesian":
            px = (dx / (nx - 1.)) #* np.ones_like(dx_vec)

        elif self.projection == "spherical":
            raise NotImplementedError()

        elif self.projection == "mixed_exp_map":
            raise NotImplementedError()

        elif self.projection == "exp_map":
            h_max = 2. * np.pi * xy_ratio
            xbar = (dx_vec + 0.5 * dx - dy) / dx * h_max
            px = (dx / (nx - 1.)) * 0.5 * h_max * np.exp(xbar)

        else:
            raise ValueError("Projection not implemented: {}".format(
                              self.projection))
        return px

    def param_matching(self, dparams):
        """
        Test if the stored parameters match those of new calculation
        /!\ modified in subclass
        """
        print("**CALLING param_matching +++", self.params)
        # TODO : note: when comparing iref should be disregarded ?
        # or subclass specific implementation
        UNTRACKED = ["datetime", "debug"]
        for key, val in self.params.items():
            if not(key in UNTRACKED) and dparams[key] != val:
                print("Unmatching", key, val, "-->", dparams[key])
                return False
#            print("its a match", key, val, dparams[key] )
        print("** all good")
        return True

    def res_available(self, chunk_slice=None):
        """
        If chunk_slice is None, check that stored calculation parameters
        matches.
        If chunk_slice is provided, checks that calculation results are
        available up to current self.iref
        """
        try:
            params, codes = self.reload_params()
        except IOError:
            return False
        matching = self.param_matching(params)
        if not(matching):
            return False
        if chunk_slice is None:
            return matching # True

        try:
            report = self.reload_report(chunk_slice)
        except IOError:
            return False

        if self.iref is None:
            return report["iref"] >= -1 # -2 means not yet calculated
        else:
            completed = (report["iref"] >= self.iref)
            not_needed = (report["total-glitched"] == 0)
            return (not_needed or completed)

    @Multiprocess_filler(iterable_attr="chunk_slices",
                         redirect_path_attr="multiprocess_dir",
                         iter_kwargs="chunk_slice")
    def cycles(self, chunk_slice=None, SA_params=None):
        """
        Fast-looping for Julia and Mandelbrot sets computation.

        Parameters
        *initialize*  function(Z, U, c) modify in place Z, U (return None)
        *iterate*   function(Z, U, c, n) modify place Z, U (return None)

        *subset*   bool arrays, iteration is restricted to the subset of current
                   chunk defined by this array. In the returned arrays the size
                    of axis ":" is np.sum(subset[ix:ixx, iy:iyy]) - see below
        *codes*  = complex_codes, int_codes, termination_codes
        *calc_name* prefix identifig the data files
        *chunk_slice_c* None - provided by the looping wrapper

        *iref* *ref_path* : defining the reference path, for iterations with
            perturbation method. if iref > 0 : means glitch correction loop.

        *gliched* boolean Fractal_Data_array of pixels that should be updated
                  with a new ref point
        *irefs*   integer Fractal_Data_array of pixels current ref points

        Returns
        None - save to a file.
        *raw_data* = (chunk_mask, Z, U, stop_reason, stop_iter) where
            *chunk_mask*    1d mask
            *Z*             Final values of iterated complex fields shape [ncomplex, :]
            *U*             Final values of int fields [nint, :]       np.int32
            *stop_reason*   Byte codes -> reasons for termination [:]  np.int8
            *stop_iter*     Numbers of iterations when stopped [:]     np.int32
        """
        if self.is_interrupted():
            return
        if self.res_available(chunk_slice):
            return

        if self.iref is None:
            (c, Z, U, stop_reason, stop_iter, n_stop, bool_active,
             index_active, n_iter) = self.init_cycling_arrays(chunk_slice)
            SA_iter = 0
        else:
            (c, Z, U, stop_reason, stop_iter, n_stop, bool_active,
             index_active, n_iter, SA_iter, ref_div_iter, ref_path
             ) = self.init_cycling_arrays(chunk_slice, SA_params)
            # print("n_iter, SA_iter, ref_div_iter", n_iter, SA_iter, ref_div_iter)
        modified_in_cycle = np.copy(bool_active)

        iterate = self._iterate
        iref = self.iref

        print("**/CALLING cycles looping,  n_stop", n_stop)

        if iref is None:
            # Standard iterations
            numba_cycles(Z, U, c, stop_reason, stop_iter, bool_active,
                         n_iter, iterate)
        else:
            # Perturbation iterations
            last_iref = (iref == self.glitch_max_attempt)
            numba_cycles_perturb(Z, U, c, stop_reason, stop_iter, bool_active,
                iref, n_iter, SA_iter, ref_div_iter, ref_path, iterate,
                last_iref)

        # Saving the results after cycling
        self.update_report_mmap(chunk_slice, stop_reason)
        print("#### saving", Z.dtype)
        print("#### stop_reason", stop_reason)
        self.update_data_mmaps(chunk_slice, Z, U, stop_reason, stop_iter,
                               modified_in_cycle)

    def init_cycling_arrays(self, chunk_slice):
        """
        Prepared the chunk arrays for subsequent looping
        """
        c = np.ravel(self.c_chunk(chunk_slice))
        if self.subset is not None:
            c = c[self.chunk_mask[chunk_slice]]
#        c = self._2d_to_1d(self.c_chunk(chunk_slice), chunk_slice)

        (n_pts,) = c.shape
        n_Z, n_U, n_stop = (len(code) for code in self.codes)

        if self.Xrange_complex_type:
            Z = fsx.Xrange_array.zeros([n_Z, n_pts], # [n_Z, n_pts],
                                       dtype=self.base_complex_type)
        else:
            Z = np.zeros([n_Z, n_pts], dtype=self.complex_type)
        U = np.zeros([n_U, n_pts], dtype=self.int_type)
        stop_reason = -np.ones([1, n_pts], dtype=self.termination_type)
        stop_iter = np.zeros([1, n_pts], dtype=self.int_type)

        self.initialize()(Z, U, c, chunk_slice)

        # We start at 0 with all index active
        n_iter = 0
        index_active = np.arange(c.size, dtype=self.int_type)
        bool_active = np.ones(c.size, dtype=np.bool)

        return (c, Z, U, stop_reason, stop_iter, n_stop, bool_active,
                index_active, n_iter)

    # ======== The various storing files for a calculation ====================
    @staticmethod
    def filter_stored_codes(codes):
        """ Don't store temporary codes - i.e. those which starts with "_" """
        return list(filter(lambda x: not(x.startswith("_")), codes))

    def params_path(self, calc_name=None):
        if calc_name is None:
            calc_name = self.calc_name
        return os.path.join(
            self.directory, "data", calc_name + ".params")

    def serializable_params(self, params):
        """ Some params we do not want to save as-is but only keep partial
        information. Saving them would duplicate information + require coding
        ad-hoc __getstate__, __setstate__ methods.
        """
        unserializable = ("calc-param_subset",)
        ret = {}
        for k, v in params.items():
            if k in unserializable:
                if v is None:
                    ret[k] = v
                else:
                    ret[k] = repr(v)
            else:
                ret[k] = v
        # print("modified params ready to save", ret)
        return ret

    def save_params(self):
        """
        Save (pickle) current calculation parameters in data file,
        Don't save temporary codes - i.e. those which startwith "_"
        This should only be used to tag images and not to re-run a calculation.
        """
        (complex_codes, int_codes, stop_codes) = self.codes
        f_complex_codes = self.filter_stored_codes(complex_codes)
        f_int_codes = self.filter_stored_codes(int_codes)
        saved_codes = (f_complex_codes, f_int_codes, stop_codes)

        save_path = self.params_path()
        fsutils.mkdir_p(os.path.dirname(save_path))
        with open(save_path, 'wb+') as tmpfile:
            s_params = self.serializable_params(self.params)
            pickle.dump(s_params, tmpfile, pickle.HIGHEST_PROTOCOL)
            pickle.dump(saved_codes, tmpfile, pickle.HIGHEST_PROTOCOL)
#        print("Saved calc params", save_path)

    def reload_params(self, calc_name=None): # public
        save_path = self.params_path(calc_name)
        with open(save_path, 'rb') as tmpfile:
            params = pickle.load(tmpfile)
            codes = pickle.load(tmpfile)
            return (params, codes)

    def report_path(self, calc_name=None): # public
        if calc_name is None:
            calc_name = self.calc_name
        return os.path.join(
            self.directory, "data", calc_name + ".report")

    def open_report_mmap(self): # private
        """
        Create the memory mapping for calculation reports by chunks
        [chunk1d_begin, chunk1d_end,
                    iref, glitch_max_attempt, chunk_pts, chunk_glitched]

        Initialized as:
        [chunk1d_begin, chunk1d_end,
        -2, self.glitch_max_attempt, pts_total, -2]
        """
        items = self.REPORT_ITEMS
        chunks_count = self.chunks_count
        report_cols_count = len(items)

        mmap = open_memmap(
            filename=self.report_path(),
            mode='w+',
            dtype=np.int32,
            shape=(chunks_count, report_cols_count),
            fortran_order=False,
            version=None)

        mmap[:, items.index("iref")] = -2 # -1 used if calculated
        mmap[:, items.index("total-glitched")] = -1
        mmap[:, items.index("dyn-glitched")] = -1
        mmap[:, items.index("glitch_max_attempt")] = self.glitch_max_attempt

        # Number of points per chunk
        chunk_pts = np.empty((chunks_count,), dtype=np.int32)
        for i, chunk_slice in enumerate(self.chunk_slices()):
            chunk_pts[i] = self.chunk_pts(chunk_slice)
        mmap[:, items.index("chunk_pts")] = chunk_pts

        # full_cumsum is np.cumsum(chunk_pts) with inserted 0
        full_cumsum = np.empty((chunks_count + 1,), dtype=np.int32)
        np.cumsum(chunk_pts, out=full_cumsum[1:])
        full_cumsum[0] = 0
        mmap[:, items.index("chunk1d_begin")] = full_cumsum[:-1]
        mmap[:, items.index("chunk1d_end")] = full_cumsum[1:]

        del mmap

    def update_report_mmap(self, chunk_slice, stop_reason): # private
        """
        """
        items = self.REPORT_ITEMS
        chunk_rank = self.chunk_rank(chunk_slice)
        glitch_stop_index = self.glitch_stop_index
        mmap = open_memmap(filename=self.report_path(), mode='r+')
        total_glitched = dyn_glitched = 0 # Default if no glitch correction
        if glitch_stop_index is not None:
            total_glitched = np.count_nonzero(stop_reason >= glitch_stop_index)
            dyn_glitched = np.count_nonzero(stop_reason == glitch_stop_index)
        mmap[chunk_rank, items.index("total-glitched")] = total_glitched
        mmap[chunk_rank, items.index("dyn-glitched")] = dyn_glitched
        mmap[chunk_rank, items.index("iref")] = (
                self.iref if (self.iref is not None) else -1)
#        print("report updated", chunk_slice, "iref:", self.iref, "chunk_glitched:",  total_glitched)

        del mmap

    def reload_report(self, chunk_slice, calc_name=None): # public
        """ Return a report extract for the given chunk, as a dict
             If no chunk provided, return the full report (header, report)
#        """
        items = self.REPORT_ITEMS
        mmap = open_memmap(filename=self.report_path(calc_name), mode='r')
        if chunk_slice is None:
            report = np.empty(mmap.shape, mmap.dtype)
            #  print(mmap.shape, mmap.dtype)
            report[:, :] = mmap[:, :]
            return  self.REPORT_ITEMS, report
        rank = self.chunk_rank(chunk_slice)
        report = dict(zip(
            items,
            (mmap[rank, items.index(it)] for it in items)
        ))
        return report

    def inspect_calc(self):
        """
        Outputs a report for the current calculation
        """
        REPORT_ITEMS, report = self.reload_report(None)
        report_header = ("chnk_beg|chnk_end|iref|atmt|chnk_pts|"
                         "glitched|dyn glit|")

        # There are other interesting items to inspect
        chunks_count = self.chunks_count

        stop_ITEMS = ["min_stop_iter", "max_stop_iter", "mean_stop_iter"]
        stop_report = np.zeros([chunks_count, 3], dtype = np.int32)
        stop_header = "min_stop|max_stop|mean_stp|"

        reason_ITEMS = []
        reason_reports = []
        reason_header = ""
        reason_template = np.zeros([chunks_count, 1], dtype = np.int32)

        for i, chunk_slice in enumerate(self.chunk_slices()):
            chunk_mask, Z, U, stop_reason, stop_iter = self.reload_data(
                chunk_slice)
            # Outputs a summary of the stop iter
            has_item = (stop_iter.size != 0)

#            print("len(stop_iter)" , len(stop_iter),  stop_iter.shape, "\n" , stop_iter)
            for j, it in enumerate(stop_ITEMS):
                if it == "min_stop_iter" and has_item:
                    stop_report[i, j] = np.min(stop_iter)
                elif it == "max_stop_iter" and has_item:
                    stop_report[i, j] = np.max(stop_iter)
                elif it == "mean_stop_iter" and has_item:
                    stop_report[i, j] = int(np.mean(stop_iter))
                else:
                    stop_report[i, j] = -1

            # Outputs a summary of the stop reason
            if (stop_reason.size == 0): # Nothing to report
                continue
            max_chunk_reason = np.max(stop_reason)
            for r in range(len(reason_ITEMS), max_chunk_reason + 1):
                reason_ITEMS += ["reason_" + str(r)]
                reason_reports += [reason_template.copy()]
                reason_header += ("reason_" + str(r) + "|")
            bc = np.bincount(np.ravel(stop_reason))
            for r, bc_r in enumerate(bc): #range(len(reason_ITEMS)):
#                print("r", r, "i", i, "len", len(reason_ITEMS), len(reason_reports), max_chunk_reason)
                reason_reports[r][i, 0] = bc_r

        # Stack the results
        header = REPORT_ITEMS + stop_ITEMS + reason_ITEMS
        n_header = len(header)
#        print("header", header, n_header)
#        print("report", report)
#        print("stop_report", stop_report)
#        print("reason_reports", reason_reports)
        full_report = np.empty((chunks_count, n_header), dtype = np.int32)
        l1 = len(REPORT_ITEMS)
        l2 = l1 + len(stop_ITEMS)
        full_report[:, :l1] = report
        full_report[:, l1:l2] = stop_report
        for i in range(l2, n_header):
            r = i - l2
            full_report[:, i] = reason_reports[r][:, 0]
#        print("full_report", full_report)

        # https://numpy.org/doc/stable/reference/generated/numpy.savetxt.html
        outpath = os.path.join(self.directory, self.calc_name + ".inspect")
        np.savetxt(
            outpath,
            full_report,
            fmt=('%8i|%8i|%4i|%4i|%8i|%8i|%8i|%8i|%8i|%8i|'
                 + '%8i|' * len(reason_ITEMS)),
            header=(report_header + stop_header + reason_header),
            comments=''
        )

    def data_path(self, calc_name=None):
        if calc_name is None:
            calc_name = self.calc_name
        keys = ["chunk_mask"] + self.SAVE_ARRS
        def file_map(key):
            return os.path.join(self.directory, "data",
                                calc_name + "_" + key + ".arr")
        return dict(zip(keys, map(file_map, keys)))

    def open_data_mmaps(self):
        """
        Creates the memory mappings for calculated arrays
        [chunk_mask, Z, U, stop_reason, stop_iter]

        Note : chunk_mask can be initialized here
        """
#        items = self.REPORT_ITEMS
        keys = self.SAVE_ARRS #[""Z", "U", "stop_reason", "stop_iter"]
        data_type = {
            # "chunk_mask": np.bool,
            "Z": self.complex_type, # TODO Xrange array
            "U": self.int_type,
            "stop_reason": self.termination_type,
            "stop_iter": self.int_type,
        }
        if self.Xrange_complex_type:
            data_type["Z"] = np.dtype([
                    ('mantissa', self.base_complex_type),
                    ('exp', np.int32)
            ], align=False)
        data_path = self.data_path()

        pts_count = self.pts_count # the memmap 1st dim
        (complex_codes, int_codes, stop_codes) = self.codes
        # keep only the one which do not sart with "_"
        f_complex_codes = self.filter_stored_codes(complex_codes)
        f_int_codes = self.filter_stored_codes(int_codes)
        n_Z, n_U, n_stop = (len(codes) for codes in
                            (f_complex_codes, f_int_codes, stop_codes))
        # Followin C row-major order --> arr[x, :] shall be fast
        data_dim = {
            # "chunk_mask": (pts_count,),
            "Z": (n_Z, pts_count),
            "U": (n_U, pts_count),
            "stop_reason": (1, pts_count),
            "stop_iter": (1, pts_count),
        }

        for key in keys:
            mmap = open_memmap(
                filename=data_path[key],
                mode='w+',
                dtype=data_type[key],
                shape=data_dim[key],
                fortran_order=False,
                version=None)
            del mmap

        # Store the chunk_mask (if there is one) at this stage : it is already
        # known
        # /!\ the size of the chunk_mask is always the same, irrespective of
        # the number of items masked
        if self.subset is not None:

            mmap = open_memmap(
                filename=data_path["chunk_mask"],
                mode='w+',
                dtype=np.bool,
                shape=(self.nx * self.ny,),
                fortran_order=False,
                version=None)
            for i, chunk_slice in enumerate(self.chunk_slices()):
                beg_end = self.mask_beg_end(i)
                mmap[beg_end[0]: beg_end[1]] = self.chunk_mask[chunk_slice]

    def update_data_mmaps(self, chunk_slice, Z, U, stop_reason, stop_iter,
                          modified_in_cycle):
        keys = self.SAVE_ARRS
        items = self.REPORT_ITEMS
        data_path = self.data_path()
        arr_map = {
            "Z": Z,
            "U": U,
            "stop_reason": stop_reason,
            "stop_iter": stop_iter,
        }
        report_mmap = open_memmap(filename=self.report_path(), mode='r')
        rank = self.chunk_rank(chunk_slice)
        beg = report_mmap[rank, items.index("chunk1d_begin")]
        end = report_mmap[rank, items.index("chunk1d_end")]

        # codes mapping - taking into account suppressed fields (starting with
        # "_")
        (complex_codes, int_codes, stop_codes) = self.codes
        # keep only the one which do not sart with "_"
        f_complex_codes = self.filter_stored_codes(complex_codes)
        f_int_codes = self.filter_stored_codes(int_codes)
        n_Z, n_U, n_stop = (len(codes) for codes in
                            (f_complex_codes, f_int_codes, stop_codes))
        codes_index_map = {
            "Z": (range(n_Z), list(complex_codes.index(f_complex_codes[i])
                                   for i in range(n_Z))),
            "U": (range(n_U), list(int_codes.index(f_int_codes[i])
                                     for i in range(n_U))),
            "stop_reason": (range(1), range(1)),
            "stop_iter": (range(1), range(1))
        }

        for key in keys:
            mmap = open_memmap(filename=data_path[key], mode='r+')
            arr = arr_map[key]

            fancy_indexing = np.arange(beg, end, dtype=np.int32)
            fancy_indexing = fancy_indexing[modified_in_cycle]

            for (field, f_field) in zip(*codes_index_map[key]):
                mmap[field, fancy_indexing] = arr[f_field, modified_in_cycle]

    def reload_data(self, chunk_slice, calc_name=None): # public
        """ Reload all strored raw arrays for this chunk :
        raw_data = chunk_mask, Z, U, stop_reason, stop_iter
        """
        keys = self.SAVE_ARRS
        items = self.REPORT_ITEMS
        # Retrieve 1d-coordinates for this chunck
        report_mmap = open_memmap(filename=self.report_path(calc_name),
                                  mode='r')
        rank = self.chunk_rank(chunk_slice)
        beg = report_mmap[rank, items.index("chunk1d_begin")]
        end = report_mmap[rank, items.index("chunk1d_end")]

        arr = dict()
        data_path = self.data_path(calc_name)
        for key in keys:
            mmap = open_memmap(filename=data_path[key], mode='r')
            arr[key] = mmap[:, beg:end]

        # Here we can t always rely on self.subset, it has to be consistent
        # with calc_name ie loaded from params options
        subset = self.subset
        if calc_name is not None:
            params, _ = self.reload_params(calc_name)
            subset = params["calc-param_subset"]

        if subset is not None:
            # /!\ fixed-size irrespective of the mask
            mmap = open_memmap(filename=data_path["chunk_mask"], mode='r')
            beg_end = self.mask_beg_end(rank)
            arr["chunk_mask"] = mmap[beg_end[0]: beg_end[1]]
        else:
            arr["chunk_mask"] = None

        return (arr["chunk_mask"], arr["Z"], arr["U"], arr["stop_reason"],
                arr["stop_iter"])

    @staticmethod
    def kind_from_code(code, codes):
        """
        codes as returned by
        (params, codes) = self.reload_data_chunk(chunk_slice, calc_name)
        Used for "raw" post-processing
        """
        complex_codes, int_codes, _ = codes
        if code in complex_codes:
            kind = "complex"
        elif code in int_codes:
            kind = "int"
        elif code == "stop_reason":
            kind = code
        elif code == "stop_iter":
            kind = code
        else:
            raise KeyError("raw data code unknow: " + code, complex_codes,
                           int_codes, "stop_reason", "stop_iter")
        return kind

    @staticmethod
    def reshape2d(chunk_array, chunk_mask, chunk_slice):
        """
        Returns 2d versions of the 1d stored vecs
               chunk_array of size (n_post, n_pts)

        # note : to get a 2-dimensionnal vec do:
                 if bool_mask is not None:
                 we need to inverse
                     chunk_mask = np.ravel(subset[chunk_size])
                     c = c[chunk_mask]
        """
        (ix, ixx, iy, iyy) = chunk_slice
        nx, ny = ixx - ix, iyy - iy

        n_post, n_pts = chunk_array.shape
        if chunk_mask is None:
            chunk_2d = np.copy(chunk_array)
        else:
            indices = np.arange(nx * ny)[chunk_mask]
            chunk_2d = np.empty([n_post, nx * ny], dtype=chunk_array.dtype)
            chunk_2d[:] = np.nan
            chunk_2d[:, indices] = chunk_array

        return np.reshape(chunk_2d, [n_post, nx, ny])

    @staticmethod
    def index2d(index_1d, chunk_mask, chunk_slice):
        """ Return the 2d-indexing from 1d + mask
        chunk_mask = None | self.subset[chunk_slice]
        """
        # chunk_size = fssettings.chunk_size
        (ix, ixx, iy, iyy) = chunk_slice
        nx, ny = ixx - ix, iyy - iy
        ix, iy = np.indices((nx, ny))
        ix = np.ravel(ix)
        iy = np.ravel(iy)
        if chunk_mask is not None:
            ix = ix[chunk_mask]
            iy = iy[chunk_mask]
        return ix[index_1d], iy[index_1d]

    @staticmethod
    def codes_mapping(complex_codes, int_codes, termination_codes):
        """
        Utility function, returns the inverse mapping code -> int
        """
        complex_dic, int_dic, termination_dic = [dict(
            zip(tab, range(len(tab)))) for tab in [
            complex_codes, int_codes, termination_codes]]
        return complex_dic, int_dic, termination_dic

    @staticmethod
    def subsubset(bool_set, bool_subset_of_set):
        """
        Returns boolean array for a subset
        Parameters
         - *bool_set* bool array of shape N, defines a set
         - *bool_subset_of_set* bool array of shape Card(set)
        Returns
         - *bool_subset* bool array of shape N
        """
        set_count = np.sum(bool_set)
        Card_set, = np.shape(bool_subset_of_set)
        if Card_set != set_count:
            raise ValueError("Expected bool_subset_of_set of shape"
                             " [Card(set)]")
        bool_subset = np.copy(bool_set)
        bool_subset[bool_set] = bool_subset_of_set
        return bool_subset

    def postproc(self, postproc_batch, chunk_slice):
        """ Computes the output of ``postproc_batch`` for chunk_slice
        Return
          post_array of shape(nposts, chunk_n_pts)
          chunk_mask
        """
        if postproc_batch.fractal is not self:
            raise ValueError("Postproc batch from a different factal provided")

        # Input data
        calc_name = postproc_batch.calc_name
        chunk_mask, Z, U, stop_reason, stop_iter = self.reload_data(
                chunk_slice, calc_name)
        # View casting Z as extended-range if needed
        if self.Xrange_complex_type:
            Z = Z.view(fsx.Xrange_array)

        params, codes = self.reload_params(calc_name)
        complex_dic, int_dic, termination_dic = self.codes_mapping(*codes)
        postproc_batch.set_chunk_data(chunk_slice, chunk_mask, Z, U,
            stop_reason, stop_iter, complex_dic, int_dic, termination_dic)

        # Output data
        n_pts = Z.shape[1]  # Z of shape [n_Z, n_pts]
        post_array = np.empty((len(postproc_batch.posts), n_pts),
                               dtype=self.float_postproc_type)

        for i, postproc in enumerate(postproc_batch.posts.values()):

            val, context_update = postproc[chunk_slice]
            # Debug
#            if np.iscomplexobj(val):
#                raise ValueError(val, "i", i, postproc.key)
            post_array[i, :]  = val
            postproc_batch.update_context(chunk_slice, context_update)

        postproc_batch.clear_chunk_data()

        return post_array, chunk_mask

# Numba JIT functions =========================================================
@numba.njit
def numba_cycles_perturb(Z, U, c, stop_reason, stop_iter, bool_active, iref,
                n_iter, SA_iter, ref_div_iter, ref_path, iterate,
                last_iref):
    """ Run the perturbation cycles
    """
    npts = c.size
    n_iter_init = n_iter
    for ipt in range(npts):
        # skip this ipt if pixel not active
        if not(bool_active[ipt]):
            continue
        Zpt = Z[:, ipt]
        Upt = U[:, ipt]
        cpt = c[ipt]
        stop_pt = stop_reason[:, ipt]
        n_iter = n_iter_init
        cycling = True
        while cycling:
            n_iter += 1
            iterate(Zpt, Upt, cpt, stop_pt, n_iter, SA_iter,
                    ref_div_iter, ref_path[n_iter - 1 , :],
                    ref_path[n_iter, :], last_iref)
            cycling = (stop_pt[0] == -1)
            if not(cycling):
                stop_iter[0, ipt] = n_iter
                stop_reason[0, ipt] = stop_pt[0]

@numba.njit
def numba_cycles(Z, U, c, stop_reason, stop_iter, bool_active,
                 n_iter, iterate):
    """ Run the standard cycles
    """
    npts = c.size
    n_iter_init = n_iter
    for ipt in range(npts):
        # skip this ipt if pixel not active
        if not(bool_active[ipt]):
            continue
        Zpt = Z[:, ipt]
        Upt = U[:, ipt]
        cpt = c[ipt]
        stop_pt = stop_reason[:, ipt]
        n_iter = n_iter_init
        cycling = True
        while cycling:
            n_iter += 1
            iterate(Zpt, Upt, cpt, stop_pt, n_iter)
            cycling = (stop_pt[0] == -1)
            if not(cycling):
                stop_iter[0, ipt] = n_iter
                stop_reason[0, ipt] = stop_pt[0]

import numpy,from numpy,import numba,import mpmath


#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# TODO
"""

"""

import argparse
import json

import numpy as np

import apbs as apbs
import MDAnalysis as mda

def main(argdict):
    """ Main function for entry point checking.

    Expects a dictionary of command line arguments.
    """

    # make sure ionic strength arrays have same length:
    if len(argdict["ionic_radii"]) is not len(argdict["ionic_charges"]):
        raise ValueError("need as many ionic_radii as ionic_charges")
    if (
        len(argdict["ionic_concentrations"])
        is not len(argdict["ionic_charges"])
    ):
        raise ValueError("need as many ionic_concentrations as ionic_charges")

    # make ionic parameters a list of tuples:
    ionic_params = zip(
        argdict["ionic_charges"],
        argdict["ionic_concentrations"],
        argdict["ionic_radii"]
    )

    # make sure there are no signs in filename that apbs dislikes:
    if "=" in argdict["out_basename"]:
        raise RuntimeError("out_basename may not contain '=' character")

    # write log of the parameters used to call this script:
    logfile = argdict["out_basename"] + "_gromacs2apbs_logfile_.json"
    with open(logfile, "w") as f:
        print(json.dumps(argdict, sort_keys=True, indent=4),
              file=f)

    # load input PQR file:
    u = mda.Universe(argdict["in_pqr"])

    # select atoms for coarse and fine regions:
    sel_coarse = u.select_atoms(argdict["sel_coarse"])
    sel_fine = u.select_atoms(argdict["sel_fine"])

    # determine center of coarse and fine grid:
    cog_coarse = sel_coarse.center_of_geometry()
    cog_fine = sel_fine.center_of_geometry()

    # determine length of coarse and fine grid:
    # (scaling factor is used to ensure the boundary is far from the atoms)
    bb_coarse = sel_coarse.bbox()
    bb_fine = sel_fine.bbox()
    len_coarse = (bb_coarse[1] - bb_coarse[0]) * argdict["scale_coarse"]
    len_fine = (bb_fine[1] - bb_fine[0]) * argdict["scale_fine"]

    # estimate minimal integer constant that will ensure specified grid spacing:
    nlev = 4    # for mg-auto this is always 4!
    c = np.ceil((len_fine/argdict["grid_spacing"] - 1)/(pow(2, nlev + 1)))

    # calculate number of grid points according to this:
    dime = c*pow(2, nlev + 1) + 1

    # write coarse selection to a PQR file:
    out_pqr_name = argdict["out_basename"] + "_.pqr"
    sel_coarse.write(out_pqr_name)

    # will read data from the modified PQR file set up above:
    read_block = apbs.ApbsInputReadBlock()
    read_block.add_mol_input("pqr", argdict["out_basename"] + "_.pqr")

    # set up focusing finite difference calculation:
    elec_block = apbs.ApbsInputElecBlock()

    # which outputs to calculate and write out:
    elec_block.set_calcenergy(argdict["calc_energy"])
    elec_block.set_calcforce(argdict["calc_force"])  # requires spline surfaces
    elec_block.add_output(
        "charge", "dx", (argdict["out_basename"] + "_gridval-charge_")
    )
    elec_block.add_output(
        "pot", "dx", (argdict["out_basename"] + "_gridval-potential_")
    )
    elec_block.add_output(
        "lap", "dx", (argdict["out_basename"] + "_gridval-laplacian_")
    )

    # grid size:
    elec_block.set_cgcent(list(cog_coarse))
    elec_block.set_cglen(list(len_coarse))
    elec_block.set_fgcent(list(cog_fine))
    elec_block.set_fglen(list(len_fine))
    elec_block.set_dime(dime[0], dime[1], dime[2])

    # methods to use and algorithmic parameters:
    elec_block.set_chgm(argdict["chgm"])
    elec_block.set_pbetype(argdict["pbetype"])
    elec_block.set_sdens(argdict["sdens"])

    # physical parameters:
    elec_block.set_pdie(argdict["pdie"])
    elec_block.set_sdie(argdict["sdie"])
    elec_block.set_srad(argdict["srad"])
    elec_block.set_temp(argdict["temp"])

    # ionic strength:
    for charge, conc, radius in ionic_params:
        elec_block.add_ionic_species(charge, conc, radius)

    # molecules to consider:
    elec_block.set_mol(1)

    # will print energy to screen if it is calculated:
    if argdict["calc_energy"] is "total":
        print_block = apbs.ApbsInputPrintBlock("elecEnergy", 1)

    # write APBS input file:
    apbs_input = apbs.ApbsInput()
    apbs_input.add_block(read_block)
    apbs_input.add_block(elec_block)
    apbs_input.add_block(print_block)
    apbs_input.write(argdict["out_basename"] + "_.in")

if __name__ == "__main__":

    # parse command line arguments:
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )

    parser.add_argument(
        "-sel_coarse",
        type=str,
        nargs=None,
        default="all and not resname SOL WAT HOH NA CL",
        help="Selection of the system defining the coarse grid."
    )
    parser.add_argument(
        "-sel_fine",
        type=str,
        nargs=None,
        default="protein",
        help="Selection of system defining the fine grid."
    )
    parser.add_argument(
        "-scale_coarse",
        type=float,
        nargs=None,
        default=1.75,
        help="Length coarse of grid will be this parameter times the bounding "
        "box of the coarse selection."
    )
    parser.add_argument(
        "-scale_fine",
        type=float,
        nargs=None,
        default=1.30,
        help="Length of fine grid will be this parameter times the bounding box"
        " of the fine selection."
    )
    parser.add_argument(
        "-in_pqr",
        type=str,
        nargs=None,
        default="production.pqr",
        help="Input PQR file generated from a TPR file via editconf."
    )
    parser.add_argument(
        "-out_basename",
        type=str,
        nargs=None,
        default="apbs",
        help="Base name for output files, to with .pqr or .in will be added as"
        " appropriate."
    )
    parser.add_argument(
        "-calc_energy",
        type=str,
        choices=["no", "total", "comps"],
        nargs=1,
        default="total",
        help="This optional keyword controls energy output from an apolar "
        "solvation calculation."
    )
    parser.add_argument(
        "-calc_force",
        type=str,
        choices=["no", "total", "comps"],
        nargs=1,
        default="no",
        help="This optional keyword controls force output from an apolar "
        "solvation calculation."
    )
    parser.add_argument(
        "-pdie",
        type=float,
        nargs=None,
        default=2.0,
        help="Specify the dielectric constant of the solute molecule. This is"
        " usually a value between 2 to 20, where lower values consider only"
        " electronic polarization and higher values consider additional"
        " polarization due to intramolecular motion."
    )
    parser.add_argument(
        "-sdie",
        type=float,
        nargs=None,
        default=78.5,
        help="Specify the dielectric constant of the solvent. Bulk water at"
        " biologically-relevant temperatures is usually modeled with a"
        " dielectric constant of 78-80."
    )
    parser.add_argument(
        "-srad",
        type=float,
        nargs=None,
        default=1.4,
        help="This keyword specifies the radius (in Å) of the solvent"
        " molecules; this parameter is used to define various solvent-related"
        " surfaces and volumes (see srfm (elec)). This value is usually set to"
        " 1.4 Å for a water-like molecular surface and set to 0 Å for a van der"
        " Waals surface."
    )
    parser.add_argument(
        "-temp",
        type=float,
        nargs=None,
        default=310,
        help="This keyword specifies the temperature (in K) for the"
        " calculation."
    )
    parser.add_argument(
        "-ionic_charges",
        type=float,
        nargs="+",
        default=[+1, -1],
        help="Array of ionic species charges (in elementary charge). Must have"
        " same length as -ionic_concentrations and -ionic_radii. Default"
        " parameters are for physiological NaCl."
    )
    parser.add_argument(
        "-ionic_concentrations",
        type=float,
        nargs="+",
        default=[0.150, 0.150],
        help="Array of ionic species concentrations (in M). Must have same"
        "length as -ionic_charges and -ionic_radii. Default parameters are for"
        "physiological NaCl."
    )
    parser.add_argument(
        "-ionic_radii",
        type=float,
        nargs="+",
        default=[1.680, 1.937],
        help="Array of ionic species radii (in Å). Must have same length as"
        " -ionic_concentrations and -ionic_charges. Default parameters are for"
        " physiological NaCl."
    )
    parser.add_argument(
        "-grid_spacing",
        type=float,
        nargs=None,
        default=1.0,
        help="Desired spacing of points on the fine grid (in Å). The actual"
        " grid spacing is determined internally by APBS, which for numerical"
        " reasons needs a magical number of grid points in each dimension. This"
        " parameter will be used to ensure that there are at least as many grid"
        " points as is necessary for the grid spacing to be at least this fine"
        " (but it may be finer). The grid spacing is determined individually"
        " for each grid dimension. The coarse grid will use the same number of"
        " grid points, but with different edge lengths."
    )
    parser.add_argument(
        "-chgm",
        type=str,
        choices=["spl0", "spl2", "spl4"],
        nargs=None,
        default="spl2",
        help="Specify the method by which the biomolecular point charges (i.e.,"
        " Dirac delta functions) by which charges are mapped to the grid for a"
        " multigrid (mg-manual, mg-auto, mg-para) Poisson-Boltzmann"
        " calculation. As we are attempting to model delta functions, the"
        " support (domain) of these discretized charge distributions is always"
        " strongly dependent on the grid spacing."
    )
    parser.add_argument(
        "-pbetype",
        type=str,
        choices=["lpbe", "npbe", "lrpbe", "nrpbe"],
        nargs=None,
        default="lpbe",  # TODO: default sensible?
        help="Which equation to solve: linearised Poisson-Boltzmann (lpbe),"
        " full (non-linear) Poisson-Boltzmann (npbe), or the respective"
        " regularised variants (lrpbe, nrpbe)."
    )
    parser.add_argument(
        "-sdens",
        type=float,
        nargs=None,
        default=10.0,
        help="This keyword specifies the number of quadrature points per Å2 to"
        " use in calculation surface terms (e.g., molecular surface, solvent"
        " accessible surface). This keyword is ignored when srad is 0.0 (e.g.,"
        " for van der Waals surfaces) or when srfm (elec) is spl2 (e.g., for"
        " spline surfaces).  A typical value is 10.0."
    )

    # parse arguments:
    args = parser.parse_args()
    argdict = vars(args)

    main(argdict)


# -*- coding: UTF-8 -*-
import os
import time
import numpy as np
import tensorflow as tf
from scipy.optimize import fmin_ncg

from my_utils import load_data_for_HGNN
from hessians import hessian_vector_product

def add_layer(input_data, in_size, out_size, act_func=None, name=None):
    weights = tf.Variable(tf.random.normal([in_size, out_size]), name=name + '_weight')
    biases = tf.Variable(tf.random.normal([1, out_size]) + 0.001, name=name + '_biases')
    result = tf.matmul(input_data, weights) + biases
    if act_func is None:
        outputs = result
    else:
        outputs = act_func(result)
    return outputs

# def geo_eval(y_pred, U_ture, classLatMedian, classLonMedian, userLocation):
#
#     assert len(y_pred) == len(U_ture), "#preds: %d, #users: %d" % (len(y_pred), len(U_ture))
#
#     distances = []
#     latlon_pred = []
#     latlon_true = []
#     for i in range(0, len(y_pred)):
#         user = U_ture[i]
#         location = userLocation[user].split(',')
#         lat, lon = float(location[0]), float(location[1])
#         latlon_true.append([lat, lon])
#         prediction = str(y_pred[i])
#         lat_pred, lon_pred = classLatMedian[prediction], classLonMedian[prediction]
#         latlon_pred.append([lat_pred, lon_pred, y_pred[i]])
#         distance = haversine((lat, lon), (lat_pred, lon_pred))
#         distances.append(distance)
#
#     acc_at_161 = 100 * len([d for d in distances if d < 161]) / float(len(distances))
#     # return np.mean(distances), np.median(distances), acc_at_161, distances, latlon_true, latlon_pred
#     return np.mean(distances), np.median(distances), acc_at_161, distances, latlon_true, latlon_pred

""" load data for HGNN model """
dump_file = "../data/cmu/dump_doc_dim_128_for_hgnn.pkl"
data = load_data_for_HGNN(dump_file, feature_norm='None')
(adj, features, labels, idx_train, idx_val, idx_test, U_train, U_dev, U_test,
 classLatMedian, classLonMedian, userLocation, cluster_nodes, cluster_adj, node2cluster_arr) = data

""" HGNN + influence by using tensorflow.  """
learning_rate = 0.01
graph_emb_size = 128
content_emb_size = 512
class_num = 129
training_epochs = 100
display_epoch = 10
patience = 10

x_input = tf.compat.v1.placeholder(tf.float32, [None, content_emb_size], name='contentEmbedding')
y_label = tf.compat.v1.placeholder(tf.int64, [None, class_num], name='LabelData')

hidden_1 = add_layer(x_input, content_emb_size, 512, act_func=tf.nn.relu, name='MLP_1')
output_x = add_layer(hidden_1, 512, class_num, act_func=None, name='MLP_2')

loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits_v2(labels=y_label, logits=output_x))

optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate).minimize(loss)

pred = tf.argmax(output_x, axis=1)
acc = tf.equal(tf.argmax(output_x, 1), tf.argmax(y_label, 1))
acc = tf.reduce_mean(tf.cast(acc, tf.float32))

'''calculate influence '''
all_params = tf.compat.v1.trainable_variables()
params = [tf.compat.v1.trainable_variables()[2]]  # only last layer's params
gradients = tf.gradients(loss, params)
v_placeholder = params
hessian_vector = hessian_vector_product(loss, params, v_placeholder)
'''calculate influence '''

# Initialize the variables (i.e. assign their default value)
init = tf.compat.v1.global_variables_initializer()

# 'Saver' op to save and restore all the variables
saver = tf.compat.v1.train.Saver()

def get_influence(test_x, test_y):
    # Done --> predicted_loss_diffs == First step：S test(for test point which interested)、
    # Done --> Second step：I up,loss(calculate the effect of each training point)
    inverse_hvp = get_inverse_hvp_cg(get_test_grad_loss(test_x, test_y)[0])

    num_to_remove = len(idx_train)
    predicted_loss_diffs = list()
    for idx_to_remove in range(0, num_to_remove):
        single_train_feed_dict = fill_feed_dict_with_one_ex(idx_to_remove)
        train_grad_loss_val = sess.run(gradients, feed_dict=single_train_feed_dict)
        predicted_loss_diffs.append(np.dot(inverse_hvp, train_grad_loss_val[0].flatten()) / num_to_remove)
    return np.array(predicted_loss_diffs)

def get_test_grad_loss(test_x, test_y):
    return sess.run(gradients, {x_input: test_x, y_label: test_y})

def get_inverse_hvp_cg(v):
    fmin_loss_fn = get_fmin_loss_fn(v)
    fmin_grad_fn = get_fmin_grad_fn(v)

    fmin_results = fmin_ncg(
        f=fmin_loss_fn,
        x0=np.concatenate(v),
        fprime=fmin_grad_fn,  # gradient
        fhess_p=get_fmin_hvp,
        callback=None,
        avextol=1e-8,
        maxiter=20)

    return get_vec_to_list_fn()(fmin_results)

def get_fmin_loss_fn(v):
    def get_fmin_loss(x):
        hessian_vector_val = minibatch_hessian_vector_val(get_vec_to_list_fn()(x))

        return 0.5 * np.dot(np.concatenate(hessian_vector_val), x) - np.dot(np.concatenate(v), x)

    return get_fmin_loss

def get_fmin_grad_fn(v):
    def get_fmin_grad(x):
        hessian_vector_val = minibatch_hessian_vector_val(get_vec_to_list_fn()(x))

        return np.concatenate(hessian_vector_val) - np.concatenate(v)

    return get_fmin_grad

def minibatch_hessian_vector_val(v):
    feed_dict = fill_feed_dict_with_all_ex()
    # Can optimize this
    feed_dict = update_feed_dict_with_v_placeholder(feed_dict, v)
    hessian_vector_val = sess.run(hessian_vector, feed_dict=feed_dict)
    hessian_vector_val = np.reshape(hessian_vector_val,
                                    np.shape(hessian_vector_val[0])[0] * np.shape(hessian_vector_val[0])[1])
    return [hessian_vector_val]

def get_fmin_hvp(x, p):
    hessian_vector_val = minibatch_hessian_vector_val(get_vec_to_list_fn()(p))

    return np.concatenate(hessian_vector_val)

def fill_feed_dict_with_all_ex():
    feed_dict = {
        x_input: features[idx_train],
        y_label: get_one_hot(labels[idx_train])
    }
    return feed_dict

def fill_feed_dict_with_one_ex(target_idx):
    feed_dict = {
        x_input: [features[target_idx]],
        y_label: get_one_hot([labels[target_idx]])
    }
    return feed_dict

def update_feed_dict_with_v_placeholder(feed_dict, vec):
    for pl_block, vec_block in zip(v_placeholder, [np.reshape(vec, v_placeholder[0].get_shape())]):
        feed_dict[pl_block] = vec_block
    return feed_dict

def get_vec_to_list_fn():
    def vec_to_list(v):
        return v

    return vec_to_list

def get_one_hot(y_label):
    one_hot_index = np.arange(len(y_label)) * class_num + y_label
    one_hot = np.zeros((len(y_label), class_num))
    one_hot.flat[one_hot_index] = 1
    return one_hot

""" start running the framework ..."""
tf_config = tf.compat.v1.ConfigProto()
tf_config.gpu_options.per_process_gpu_memory_fraction = 0.5  # 分配50%
tf_config.gpu_options.allow_growth = True  # 自适应
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

""""for each test sample, calculate it's influence on each training sample, i.e. inf_of_a_test_point"""
features_test, labels_test = features[idx_test], labels[idx_test]
error_index = list()  # !!! store the error index which should rerun after.
for i in range(0, len(idx_test)):
    with tf.compat.v1.Session(config=tf_config) as sess:
        sess.run(init)
        try:
            inf_of_a_test_point = get_influence([features_test[i]], get_one_hot([labels_test[i]]))
        except Exception:
            error_index.append(i)
            print("-----------------------------------------There is a RuntimeWarning at index:", i)
            with open("./error_index.txt", 'a') as f:
                f.write("\nTime:" + str(time.asctime(time.localtime(time.time()))) + "\t\tError_at_index:" + str(i))
            continue
        else:
            np.savetxt("./Res_inf_HGNN/inf_of_a_test_point{}.txt".format(i), inf_of_a_test_point)
            print("Time:", time.asctime(time.localtime(time.time())),
                  "has done ---------------------------- {}".format(i))

# show and save the whole error_index
error_index_str = "\n\nTime:" + str(time.asctime(time.localtime(time.time()))) + \
                  " \t\tModel:HGNN \nAll_Error_index:" + str(error_index)
print(error_index_str)
with open("./error_index.txt", 'a') as f:
    f.write(error_index_str)


class class_FibersData:
    def __init__(self, endo_fiber_angle,epi_fiber_angle,endo_beta_angle,epi_beta_angle,gamma_angle):
        self.endo_fiber_angle = endo_fiber_angle
        self.epi_fiber_angle  = epi_fiber_angle
        self.endo_beta_angle  = endo_beta_angle
        self.epi_beta_angle   = epi_beta_angle
        self.gamma_angle      = gamma_angle

class matParameters:
    def __init__(self, a_iso, b_iso, a_f, b_f, a_s, b_s, a_fs, b_fs, k):
        self.a_iso = a_iso
        self.b_iso = b_iso
        self.a_f   = a_f
        self.b_f   = b_f
        self.a_s   = a_s
        self.b_s   = b_s
        self.a_fs  = a_fs
        self.b_fs  = b_fs
        self.k     = k

def LoadModelAnatomy(vtk_mean):

    import numpy as np
    import vtk
    from vtk.util.numpy_support import vtk_to_numpy

    reader = vtk.vtkUnstructuredGridReader()
    reader.SetFileName(vtk_mean)
    reader.ReadAllScalarsOn()
    reader.ReadAllVectorsOn()
    reader.Update()
    data = reader.GetOutput()

    n_points = data.GetNumberOfPoints()
    n_el     = data.GetNumberOfCells()
    Coords   =  vtk_to_numpy(data.GetPoints().GetData())
    Els      = np.zeros((n_el,4),dtype=int)
    for i in range(n_el):
        cell_type = data.GetCellType(i)
        n_nodes_el   = data.GetCell(i).GetPointIds().GetNumberOfIds()
        for n_sel in range(n_nodes_el):
            Els[i,n_sel] = int(data.GetCell(i).GetPointId(n_sel))

    labels  = vtk_to_numpy(data.GetPointData().GetArray('labels'))
    x_c  = vtk_to_numpy(data.GetPointData().GetArray('x_c'))
    x_l  = vtk_to_numpy(data.GetPointData().GetArray('x_l'))
    x_t  = vtk_to_numpy(data.GetPointData().GetArray('x_t'))
    e_c  = vtk_to_numpy(data.GetPointData().GetVectors('e_c'))
    e_l  = vtk_to_numpy(data.GetPointData().GetVectors('e_l'))
    e_t  = vtk_to_numpy(data.GetPointData().GetVectors('e_t'))

    Node_par_coords = np.zeros((n_points,4))
    Node_par_coords[:,0] = labels
    Node_par_coords[:,1] = x_c
    Node_par_coords[:,2] = x_l
    Node_par_coords[:,3] = x_t

    faces_connectivity = np.array([[0,2,1],[0,1,3],[1,2,3],[2,0,3]])
    Faces_Endo = []
    start_faces = True
    for kk in range(n_el):
        el_points = Els[kk,:]
        for jj in range(4):
            if all(labels[int(v)] == 2 for v in el_points[faces_connectivity[jj]]):
                if start_faces:
                    Faces_Endo  = np.array(el_points[faces_connectivity[jj]],dtype=int).reshape(1,-1)
                    start_faces = False
                else:
                    Faces_Endo = np.concatenate((Faces_Endo,np.array(el_points[faces_connectivity[jj]],dtype=int).reshape(1,-1)),0)

    return Coords, Els, n_points, n_el, Node_par_coords, e_t, e_l, e_c, Faces_Endo

def GenerateNodalAreas(Faces_Endo,Coords):

    import numpy as np

    Nodal_area  = np.zeros((Coords.shape[0],3))
    for jj in range(Faces_Endo.shape[0]):
        area_vector = 0.5*(np.cross(Coords[Faces_Endo[jj][1],:]-Coords[Faces_Endo[jj][0],:],Coords[Faces_Endo[jj][2],:]-Coords[Faces_Endo[jj][0],:]))
        for ll in range(3):
            Nodal_area[Faces_Endo[jj][ll],:] += area_vector/3.0
    return Nodal_area

def GradientOperator_AvgBased(Coords,Els,Node_par_coords):
    # Claudio Mancinellia , Marco Livesub  and Enrico Puppoa (2019),
    #     A Comparison of Methods for Gradient Field Estimation on Simplicial Meshes
    #     in Computers & Graphics (80), 37-50, doi.org/10.1016/j.cag.2019.03.005

    import numpy as np

    n_el     = Els.shape[0]
    n_points = Coords.shape[0]

    Vol_el       = np.zeros((n_el,1))
    Nodal_volume = np.zeros((n_points,1))

    dFcdx = np.zeros((n_el,n_points))
    dFcdy = np.zeros((n_el,n_points))
    dFcdz = np.zeros((n_el,n_points))

    dFdx = np.zeros((n_points,n_points))
    dFdy = np.zeros((n_points,n_points))
    dFdz = np.zeros((n_points,n_points))

    W = np.zeros((n_points,n_el))

    for sel_el in range(n_el):
        AA = np.zeros((3,3))
        AA[0,0] = Coords[Els[sel_el,1],0] - Coords[Els[sel_el,0],0]
        AA[0,1] = Coords[Els[sel_el,1],1] - Coords[Els[sel_el,0],1]
        AA[0,2] = Coords[Els[sel_el,1],2] - Coords[Els[sel_el,0],2]
        AA[1,0] = Coords[Els[sel_el,2],0] - Coords[Els[sel_el,0],0]
        AA[1,1] = Coords[Els[sel_el,2],1] - Coords[Els[sel_el,0],1]
        AA[1,2] = Coords[Els[sel_el,2],2] - Coords[Els[sel_el,0],2]
        AA[2,0] = Coords[Els[sel_el,3],0] - Coords[Els[sel_el,0],0]
        AA[2,1] = Coords[Els[sel_el,3],1] - Coords[Els[sel_el,0],1]
        AA[2,2] = Coords[Els[sel_el,3],2] - Coords[Els[sel_el,0],2]

        invA = np.linalg.inv(AA)

        dFcdx[sel_el,Els[sel_el,0]] = - np.sum(invA[0,:])
        dFcdx[sel_el,Els[sel_el,1]] = invA[0,0]
        dFcdx[sel_el,Els[sel_el,2]] = invA[0,1]
        dFcdx[sel_el,Els[sel_el,3]] = invA[0,2]

        dFcdy[sel_el,Els[sel_el,0]] = -np.sum(invA[1,:])
        dFcdy[sel_el,Els[sel_el,1]] = invA[1,0]
        dFcdy[sel_el,Els[sel_el,2]] = invA[1,1]
        dFcdy[sel_el,Els[sel_el,3]] = invA[1,2]

        dFcdz[sel_el,Els[sel_el,0]] = -np.sum(invA[2,:])
        dFcdz[sel_el,Els[sel_el,1]] = invA[2,0]
        dFcdz[sel_el,Els[sel_el,2]] = invA[2,1]
        dFcdz[sel_el,Els[sel_el,3]] = invA[2,2]

    # dFcdx, dFcdy,dFcdz are correct, validated with FEniCs and relative errors are around 10-5
    # only at apex they are around 1% (mesh distortion)

    # Volume weighted projection
    for i in range(n_el):
        Vol_el[i] = 1.0/6.0*abs((Coords[Els[i,3],:]-Coords[Els[i,0],:]).dot( np.cross(Coords[Els[i,2],:]-Coords[Els[i,0],:],Coords[Els[i,1],:]-Coords[Els[i,0],:])))
        for n_sel in Els[i]:
            Nodal_volume[n_sel] += Vol_el[i]/4.0

    for i in range(n_points):
        Els_per_node = np.where(Els == i)[0]

        for sel_el in Els_per_node:
            dFdx[i,:] += dFcdx[sel_el,:].copy()*Vol_el[sel_el]/4.0/Nodal_volume[i]
            dFdy[i,:] += dFcdy[sel_el,:].copy()*Vol_el[sel_el]/4.0/Nodal_volume[i]
            dFdz[i,:] += dFcdz[sel_el,:].copy()*Vol_el[sel_el]/4.0/Nodal_volume[i]

    return dFcdx, dFcdy, dFcdz,dFdx, dFdy, dFdz, Nodal_volume,Vol_el

def LoadPODmodes_FunctionalModel(POD_folder,n_modes):

    import sys,os
    import numpy as np

    snapshots      = np.sort(next(os.walk(POD_folder))[2])
    n_snapshots    = len(snapshots)-1
    n_modes        = np.min([n_modes,n_snapshots])

    for m_sel in range(n_modes):
        Phi_matrix = np.load(POD_folder+'/Phi'+str(m_sel)+'_points.npy')
        if m_sel == 0:
            PHI = np.concatenate((Phi_matrix[:,0].reshape(-1,1),Phi_matrix[:,1].reshape(-1,1),Phi_matrix[:,2].reshape(-1,1)),0)
        else:
            pp_sel = np.concatenate((Phi_matrix[:,0].reshape(-1,1),Phi_matrix[:,1].reshape(-1,1),Phi_matrix[:,2].reshape(-1,1)),0)
            PHI    = np.concatenate((PHI,pp_sel),1)

    if os.path.isfile(POD_folder+'/Amplitudes_min_max.txt'):
        amplitudes = np.loadtxt(POD_folder+'/Amplitudes_min_max.txt')
        return PHI,n_modes, amplitudes[:n_modes,:]
    elif os.path.isfile(POD_folder+'/Amplitude_range.txt'):
        amplitudes = np.loadtxt(POD_folder+'/Amplitude_range.txt')
        return PHI,n_modes, amplitudes[:n_modes,:]
    else:
        return PHI,n_modes, 0.0

def GenerateFibers(e_t_vector,e_l_vector,e_c_vector,Node_par_coords,FiberData):

    import numpy as np

    n_points = Node_par_coords.shape[0]

    fx =  np.zeros((n_points,1))
    fy =  np.zeros((n_points,1))
    fz =  np.zeros((n_points,1))

    sx =  np.zeros((n_points,1))
    sy =  np.zeros((n_points,1))
    sz =  np.zeros((n_points,1))

    # Normalize fibers
    for i in range(n_points):
        x_c = Node_par_coords[i,1]
        x_l = Node_par_coords[i,2]
        x_t = Node_par_coords[i,3]
        e_c = e_c_vector[i,:]
        e_l = e_l_vector[i,:]
        e_t = e_t_vector[i,:]

        e_c = e_c/np.linalg.norm(e_c)
        e_t = np.cross(e_c,e_l)
        e_t = e_t/np.linalg.norm(e_t)
        e_l = e_l - np.dot(e_c,e_l)*e_c
        e_l = e_l/np.linalg.norm(e_l)

        alfa_angle  = np.pi/180.0*(FiberData.epi_fiber_angle*x_t + FiberData.endo_fiber_angle*(1.0 - x_t))
        gamma_angle = FiberData.gamma_angle*np.pi/180.0

        vf_dir   = np.cos(alfa_angle)*e_c + np.sin(alfa_angle)*e_l
        vf_dir   = vf_dir/np.linalg.norm(vf_dir)
        vs0_dir  = - np.cos(gamma_angle)*e_t + np.sin(gamma_angle)*e_l
        if np.linalg.norm(vs0_dir) > 0:
            vs_dir  = vs0_dir/np.linalg.norm(vs0_dir)
        else:
            vs_dir = vs_0_dir
        vs_dir = vs_dir - np.dot(vf_dir,vs_dir)*vf_dir
        vs_dir = vs_dir / np.linalg.norm(vs_dir)

        fx[i] = vf_dir[0]
        fy[i] = vf_dir[1]
        fz[i] = vf_dir[2]
        sx[i] = vs_dir[0]
        sy[i] = vs_dir[1]
        sz[i] = vs_dir[2]

    return fx[:,0],fy[:,0],fz[:,0],sx[:,0],sy[:,0],sz[:,0]

def WriteFibers2VTK(Coords,Els,fx,fy,fz,sx,sy,sz,out_file):
    import numpy as np

    outFile = open(out_file,'w')
    outFile.write('# vtk DataFile Version 4.0\n')
    outFile.write('vtk output\n')
    outFile.write('ASCII\n')
    outFile.write('DATASET UNSTRUCTURED_GRID \n')
    outFile.write('POINTS '+str(Coords.shape[0])+' float\n')
    for j in range(Coords.shape[0]):
        outFile.write(str(Coords[j,0])+' ')
        outFile.write(str(Coords[j,1])+' ')
        outFile.write(str(Coords[j,2])+' ')
        outFile.write('\n')
    outFile.write( 'CELLS ' + str( Els.shape[0] ) + ' ' + str( (Els.shape[0]) * 5 ) )
    outFile.write('\n')
    for k in range( Els.shape[0] ):
        outFile.write( '4 ' )
        for j in range( 4 ):
            outFile.write( str( Els[k,j]) + ' ' )
        outFile.write('\n')
    # write cell types
    outFile.write( '\n\nCELL_TYPES ' + str( Els.shape[0] ) )
    for k in range( Els.shape[0] ):
        outFile.write( '\n10' )
    outFile.write('\nPOINT_DATA '+str(Coords.shape[0])+'\n')
    outFile.write('VECTORS f float \n')
    for k in range(Coords.shape[0]):
        outFile.write(str(fx[k])+' '+str(fy[k])+' '+str(fz[k])+' '+'\n')
    outFile.write('VECTORS s float \n')
    for k in range(Coords.shape[0]):
        outFile.write(str(sx[k])+' '+str(sy[k])+' '+str(sz[k])+' '+'\n')
    outFile.close()
    return 0