from __future__ import division, absolute_import, print_function

__copyright__ = "Copyright (C) 2017 - 2018 Xiaoyu Wei"

__license__ = """
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
"""

import numpy as np
import logging
import pyopencl as cl
import pyopencl.array
from pytools.obj_array import make_obj_array

# from pytools import memoize_method
from volumential.nearfield_potential_table import NearFieldInteractionTable
from volumential.expansion_wrangler_interface import (
        ExpansionWranglerInterface, ExpansionWranglerCodeContainerInterface)
from sumpy.fmm import SumpyExpansionWrangler, \
        SumpyTimingFuture, SumpyExpansionWranglerCodeContainer
from boxtree.pyfmmlib_integration import FMMLibExpansionWrangler

from sumpy.kernel import (
        LaplaceKernel, HelmholtzKernel, AxisTargetDerivative,
        DirectionalSourceDerivative)

logger = logging.getLogger(__name__)

def level_to_rscale(tree, level):
    return tree.root_extent * (2 ** -level)

def inverse_id_map(queue, mapped_ids):
    """Given a index mapping as its mapped ids, compute its inverse,
    and return the inverse by the inversely-mapped ids.
    """
    cl_array = False
    if isinstance(mapped_ids, cl.array.Array):
        cl_array = True
        mapped_ids = mapped_ids.get(queue)

    inv_ids = np.zeros_like(mapped_ids)
    inv_ids[mapped_ids] = np.arange(len(mapped_ids))

    if cl_array:
        inv_ids = cl.array.to_device(queue, inv_ids)

    return inv_ids

# {{{ sumpy backend

class FPNDSumpyExpansionWranglerCodeContainer(
        ExpansionWranglerCodeContainerInterface,
        SumpyExpansionWranglerCodeContainer):
    """Objects of this type serve as a place to keep the code needed
    for ExpansionWrangler if it is using sumpy to perform multipole
    expansion and manipulations.

    Since :class:`SumpyExpansionWrangler` necessarily must have a
    :class:`pyopencl.CommandQueue`, but this queue is allowed to be
    more ephemeral than the code, the code's lifetime
    is decoupled by storing it in this object.
    """
    get_wrangler = SumpyExpansionWranglerCodeContainer.get_wrangler

class FPNDSumpyExpansionWrangler(
        ExpansionWranglerInterface, SumpyExpansionWrangler):
    """This expansion wrangler uses "fpnd" strategy. That is, Far field is
    computed via Particle approximation and Near field is computed Directly.
    The FMM is performed using sumpy backend.

    .. attribute:: source_extra_kwargs

        Keyword arguments to be passed to interactions that involve
        the source field.

    .. attribute:: kernel_extra_kwargs

        Keyword arguments to be passed to interactions that involve
        expansions, but not the source field.

    .. attribute:: self_extra_kwargs

        Keyword arguments to be passed for handling
        self interactions (singular integrals)
    """

    # {{{ constructor

    def __init__(
        self,
        code_container,
        queue,
        tree,
        near_field_table,
        dtype,
        fmm_level_to_order,
        quad_order,
        potential_kind=1,
        source_extra_kwargs=None,
        kernel_extra_kwargs=None,
        self_extra_kwargs=None,
        list1_extra_kwargs=None,
    ):
        """
        near_field_table can either one of three things:
            1. a single table, when len(target_kernels) = 1 (single level)
            2. a list of tables, when len(target_kernels) = 1 (multiple levels)
            3. otherwise, a dictionary from kernel.__repr__() to a list of its tables
        """

        self.code = code_container
        self.queue = queue
        self.tree = tree

        self.near_field_table = {}
        # list of tables for a single out kernel
        if isinstance(near_field_table, list):
            assert len(self.code.target_kernels) == 1
            self.near_field_table[
                self.code.target_kernels[0].__repr__()
            ] = near_field_table
            self.n_tables = len(near_field_table)

        # single table
        elif isinstance(near_field_table, NearFieldInteractionTable):
            assert len(self.code.target_kernels) == 1
            self.near_field_table[self.code.target_kernels[0].__repr__()] = [
                near_field_table
            ]
            self.n_tables = 1

        # dictionary of lists of tables
        elif isinstance(near_field_table, dict):
            self.n_tables = dict()
            for out_knl in self.code.target_kernels:
                if repr(out_knl) not in near_field_table:
                    raise RuntimeError(
                            "Missing nearfield table for %s." % repr(out_knl))
                if isinstance(near_field_table[repr(out_knl)],
                        NearFieldInteractionTable):
                    near_field_table[repr(out_knl)] = [
                            near_field_table[repr(out_knl)]]
                else:
                    assert isinstance(near_field_table[repr(out_knl)], list)

                self.n_tables[repr(out_knl)] = len(near_field_table[repr(out_knl)])

            self.near_field_table = near_field_table
        else:
            raise RuntimeError("Table type unrecognized.")

        self.quad_order = quad_order
        self.potential_kind = potential_kind

        # TODO: make all parameters table-specific (allow using inhomogeneous tables)
        kname = repr(self.code.target_kernels[0])
        self.root_table_source_box_extent = (
                self.near_field_table[kname][0].source_box_extent)
        table_starting_level = np.round(
            np.log(self.tree.root_extent / self.root_table_source_box_extent)
            / np.log(2)
            )
        for kid in range(len(self.code.target_kernels)):
            kname = self.code.target_kernels[kid].__repr__()
            for lev, table in zip(
                    range(len(self.near_field_table[kname])),
                    self.near_field_table[kname]
                    ):
                assert table.quad_order == self.quad_order

                if not table.is_built:
                    raise RuntimeError(
                        "Near field interaction table needs to be built "
                        "prior to being used"
                    )

                table_root_extent = table.source_box_extent * 2 ** lev
                assert (
                    abs(self.root_table_source_box_extent - table_root_extent)
                    < 1e-15
                )

                # If the kernel cannot be scaled,
                # - tree_root_extent must be integral times of table_root_extent
                # - n_tables must be sufficient
                if not isinstance(self.n_tables, dict) and self.n_tables > 1:
                    if (
                        not abs(
                            int(self.tree.root_extent / table_root_extent)
                            * table_root_extent
                            - self.tree.root_extent
                        )
                        < 1e-15
                    ):
                        raise RuntimeError(
                            "Incompatible list of tables: the "
                            "source_box_extent of the root table must "
                            "divide the bounding box's extent by an integer."
                        )

            if not isinstance(self.n_tables, dict) and self.n_tables > 1:
                # this checks that the boxes at the highest level are covered
                if (
                    not tree.nlevels
                    <= len(self.near_field_table[kname]) + table_starting_level
                ):
                    raise RuntimeError(
                        "Insufficient list of tables: the "
                        "finest level mesh cells at level "
                        + str(tree.nlevels)
                        + " are not covered."
                    )

                # the check that the boxes at the coarsest level are covered is
                # deferred until trav.target_boxes is passed when invoking
                # eval_direct

        self.dtype = dtype

        if source_extra_kwargs is None:
            source_extra_kwargs = {}

        if kernel_extra_kwargs is None:
            kernel_extra_kwargs = {}

        if self_extra_kwargs is None:
            self_extra_kwargs = {}

        if list1_extra_kwargs is None:
            list1_extra_kwargs = {}

        if not callable(fmm_level_to_order):
            raise TypeError("fmm_level_to_order not passed")

        base_kernel = code_container.get_base_kernel()
        kernel_arg_set = frozenset(kernel_extra_kwargs.items())
        self.level_orders = [
            fmm_level_to_order(base_kernel, kernel_arg_set, tree, lev)
            for lev in range(tree.nlevels)
        ]

        # print("Multipole order = ",self.level_orders)

        self.source_extra_kwargs = source_extra_kwargs
        self.kernel_extra_kwargs = kernel_extra_kwargs
        self.self_extra_kwargs = self_extra_kwargs
        self.list1_extra_kwargs = list1_extra_kwargs

        self.extra_kwargs = source_extra_kwargs.copy()
        self.extra_kwargs.update(self.kernel_extra_kwargs)

    # }}} End constructor

    # {{{ data vector utilities

    def multipole_expansion_zeros(self):
        return SumpyExpansionWrangler.multipole_expansion_zeros(self)

    def local_expansion_zeros(self):
        return SumpyExpansionWrangler.local_expansion_zeros(self)

    def output_zeros(self):
        return SumpyExpansionWrangler.output_zeros(self)

    def reorder_sources(self, source_array):
        return SumpyExpansionWrangler.reorder_sources(self, source_array)

    def reorder_targets(self, target_array):
        if not hasattr(self.tree, 'user_target_ids'):
            self.tree.user_target_ids = inverse_id_map(
                self.queue, self.tree.sorted_target_ids)
        return target_array.with_queue(self.queue)[self.tree.user_target_ids]

    def reorder_potentials(self, potentials):
        return SumpyExpansionWrangler.reorder_potentials(self, potentials)

    def finalize_potentials(self, potentials):
        # return potentials
        return SumpyExpansionWrangler.finalize_potentials(self, potentials)

    # }}} End data vector utilities

    # {{{ formation & coarsening of multipoles

    def form_multipoles(self, level_start_source_box_nrs, source_boxes, src_weights):
        return SumpyExpansionWrangler.form_multipoles(
            self, level_start_source_box_nrs, source_boxes, src_weights
        )

    def coarsen_multipoles(
        self, level_start_source_parent_box_nrs, source_parent_boxes, mpoles
    ):
        return SumpyExpansionWrangler.coarsen_multipoles(
            self, level_start_source_parent_box_nrs, source_parent_boxes, mpoles
        )

    # }}} End formation & coarsening of multipoles

    # {{{ direct evaluation of near field interactions

    def eval_direct_single_out_kernel(
        self,
        out_pot,
        out_kernel,
        target_boxes,
        neighbor_source_boxes_starts,
        neighbor_source_boxes_lists,
        mode_coefs,
    ):

        # NOTE: mode_coefs are similar to source_weights BUT
        # do not include quadrature weights (purely function
        # expansiona coefficients)

        if 0:
            print("Returns range for list1")
            out_pot[:] = cl.array.to_device(self.queue, np.arange(len(out_pot)))
            return out_pot, None

        kname = out_kernel.__repr__()

        if isinstance(self.n_tables, int) and self.n_tables > 1:
            use_multilevel_tables = True
        elif isinstance(self.n_tables, dict) and self.n_tables[kname] > 1:
            use_multilevel_tables = True
        else:
            use_multilevel_tables = False

        if use_multilevel_tables:
            # this checks that the boxes at the coarsest level
            # and allows for some round-off error
            min_lev = np.min(
                self.tree.box_levels.get(self.queue)[target_boxes.get(self.queue)]
            )
            largest_cell_extent = self.tree.root_extent * 0.5 ** min_lev
            if not self.near_field_table[kname][0].source_box_extent >= (
                largest_cell_extent - 1e-15
            ):
                raise RuntimeError(
                    "Insufficient list of tables: the "
                    "coarsest level mesh cells at level "
                    + str(min_lev)
                    + " are not covered."
                )

        # table.case_encode
        distinct_numbers = set()
        for vec in self.near_field_table[kname][0].interaction_case_vecs:
            for cvc in vec:
                distinct_numbers.add(cvc)
        base = len(range(min(distinct_numbers), max(distinct_numbers) + 1))
        shift = -min(distinct_numbers)

        case_indices_dev = cl.array.to_device(
            self.queue, self.near_field_table[kname][0].case_indices
        )

        # table.data
        table_data_combined = np.zeros(
            (
                len(self.near_field_table[kname]),
                len(self.near_field_table[kname][0].data),
            ), dtype=self.near_field_table[kname][0].data.dtype
        )
        mode_nmlz_combined = np.zeros(
            (
                len(self.near_field_table[kname]),
                len(self.near_field_table[kname][0].mode_normalizers),
            ), dtype=self.near_field_table[kname][0].mode_normalizers.dtype
        )
        exterior_mode_nmlz_combined = np.zeros(
            (
                len(self.near_field_table[kname]),
                len(self.near_field_table[kname][0].kernel_exterior_normalizers),
            ),
            dtype=self.near_field_table[kname][0].kernel_exterior_normalizers.dtype
        )
        for lev in range(len(self.near_field_table[kname])):
            table_data_combined[lev, :] = self.near_field_table[kname][lev].data
            mode_nmlz_combined[lev, :] = \
                self.near_field_table[kname][lev].mode_normalizers
            exterior_mode_nmlz_combined[lev, :] = \
                self.near_field_table[kname][lev].kernel_exterior_normalizers

        self.queue.finish()
        logger.info(
                "table data for kernel "
                + out_kernel.__repr__() + " congregated")

        # The loop domain needs to know some info about the tables being used
        table_data_shapes = {
            "n_tables": len(self.near_field_table[kname]),
            "n_q_points": self.near_field_table[kname][0].n_q_points,
            "n_table_entries": len(self.near_field_table[kname][0].data),
        }
        assert table_data_shapes["n_q_points"] == len(
            self.near_field_table[kname][0].mode_normalizers
        )

        from volumential.list1 import NearFieldFromCSR

        near_field = NearFieldFromCSR(out_kernel, table_data_shapes,
            potential_kind=self.potential_kind,
            **self.list1_extra_kwargs)

        table_data_combined = cl.array.to_device(self.queue,
                table_data_combined)
        mode_nmlz_combined = cl.array.to_device(self.queue,
                mode_nmlz_combined)
        exterior_mode_nmlz_combined = cl.array.to_device(self.queue,
            exterior_mode_nmlz_combined)
        self.queue.finish()
        logger.info("sent table data to device")

        # NOTE: box_sources for this evaluation should be "box_targets".
        # This is due to the special features of how box-FMM works.

        res, evt = near_field(
            self.queue,
            result=out_pot,
            box_centers=self.tree.box_centers,
            box_levels=self.tree.box_levels,
            box_source_counts_cumul=self.tree.box_target_counts_cumul,
            box_source_starts=self.tree.box_target_starts,
            box_target_counts_cumul=self.tree.box_target_counts_cumul,
            box_target_starts=self.tree.box_target_starts,
            case_indices=case_indices_dev,
            encoding_base=base,
            encoding_shift=shift,
            mode_nmlz_combined=mode_nmlz_combined,
            exterior_mode_nmlz_combined=exterior_mode_nmlz_combined,
            neighbor_source_boxes_starts=neighbor_source_boxes_starts,
            root_extent=self.tree.root_extent,
            neighbor_source_boxes_lists=neighbor_source_boxes_lists,
            mode_coefs=mode_coefs,
            table_data_combined=table_data_combined,
            target_boxes=target_boxes,
            table_root_extent=self.root_table_source_box_extent,
        )

        # print(near_field.get_kernel())
        # import pudb; pu.db

        assert res is out_pot

        # sorted_target_ids=self.tree.user_source_ids,
        # user_source_ids=self.tree.user_source_ids)

        # FIXME: lazy evaluation sometimes returns incorrect results
        res.finish()

        return out_pot, evt

    def eval_direct(
        self,
        target_boxes,
        neighbor_source_boxes_starts,
        neighbor_source_boxes_lists,
        mode_coefs,
    ):
        pot = self.output_zeros()
        events = []
        for i in range(len(self.code.target_kernels)):
            # print("processing near-field of out_kernel", i)
            pot[i], evt = self.eval_direct_single_out_kernel(
                pot[i],
                self.code.target_kernels[i],
                target_boxes,
                neighbor_source_boxes_starts,
                neighbor_source_boxes_lists,
                mode_coefs,
            )
            events.append(evt)

        for out_pot in pot:
            out_pot.finish()

        return (pot, SumpyTimingFuture(self.queue, events))

    # }}} End direct evaluation of near field interactions

    # {{{ downward pass of fmm

    def multipole_to_local(
        self,
        level_start_target_box_nrs,
        target_boxes,
        src_box_starts,
        src_box_lists,
        mpole_exps,
    ):
        return SumpyExpansionWrangler.multipole_to_local(
            self,
            level_start_target_box_nrs,
            target_boxes,
            src_box_starts,
            src_box_lists,
            mpole_exps,
        )

    def eval_multipoles(
        self, target_boxes_by_source_level, source_boxes_by_level, mpole_exps
    ):
        return SumpyExpansionWrangler.eval_multipoles(
            self, target_boxes_by_source_level, source_boxes_by_level, mpole_exps
        )

    def form_locals(
        self,
        level_start_target_or_target_parent_box_nrs,
        target_or_target_parent_boxes,
        starts,
        lists,
        src_weights,
    ):
        return SumpyExpansionWrangler.form_locals(
            self,
            level_start_target_or_target_parent_box_nrs,
            target_or_target_parent_boxes,
            starts,
            lists,
            src_weights,
        )

    def refine_locals(
        self,
        level_start_target_or_target_parent_box_nrs,
        target_or_target_parent_boxes,
        local_exps,
    ):
        return SumpyExpansionWrangler.refine_locals(
            self,
            level_start_target_or_target_parent_box_nrs,
            target_or_target_parent_boxes,
            local_exps,
        )

    def eval_locals(self, level_start_target_box_nrs, target_boxes, local_exps):
        return SumpyExpansionWrangler.eval_locals(
            self, level_start_target_box_nrs, target_boxes, local_exps
        )

    # }}} End downward pass of fmm

    # {{{ direct evaluation of p2p (discrete) interactions

    def eval_direct_p2p(
        self, target_boxes, source_box_starts, source_box_lists, src_weights
    ):
        return SumpyExpansionWrangler.eval_direct(
            self, target_boxes, source_box_starts, source_box_lists, src_weights
        )

    # }}} End direct evaluation of p2p interactions

# }}} End sumpy backend

# {{{ fmmlib backend (for laplace, helmholtz)

class FPNDFMMLibExpansionWranglerCodeContainer(
        ExpansionWranglerCodeContainerInterface,
        ):
    """Objects of this type serve as a place to keep the code needed
    for ExpansionWrangler if it is using fmmlib to perform multipole
    expansion and manipulations.

    The interface is augmented with unecessary arguments acting as
    placeholders, such that it can be a drop-in replacement of sumpy
    backend.
    """
    def __init__(self, cl_context,
            multipole_expansion_factory, local_expansion_factory,
            target_kernels, exclude_self=True, *args, **kwargs):
        self.cl_context = cl_context
        self.multipole_expansion_factory = multipole_expansion_factory
        self.local_expansion_factory = local_expansion_factory

        self.target_kernels = target_kernels
        self.exclude_self = True

    def get_wrangler(self, queue, tree, dtype, fmm_level_to_order,
            source_extra_kwargs={}, kernel_extra_kwargs=None,
            *args, **kwargs):
        return FPNDFMMLibExpansionWrangler(self, queue, tree,
                dtype, fmm_level_to_order,
                source_extra_kwargs, kernel_extra_kwargs,
                *args, **kwargs)

class FPNDFMMLibExpansionWrangler(
        ExpansionWranglerInterface, FMMLibExpansionWrangler):
    """This expansion wrangler uses "fpnd" strategy. That is, Far field is
    computed via Particle approximation and Near field is computed Directly.
    The FMM is performed using FMMLib backend.

    .. attribute:: source_extra_kwargs

        Keyword arguments to be passed to interactions that involve
        the source field.

    .. attribute:: kernel_extra_kwargs

        Keyword arguments to be passed to interactions that involve
        expansions, but not the source field.

    Much of this class is borrowed from pytential.qbx.fmmlib.
    """
    # {{{ constructor

    def __init__(self, code_container, queue, tree,
            near_field_table, dtype,
            fmm_level_to_order,
            quad_order,
            potential_kind=1,
            source_extra_kwargs=None,
            kernel_extra_kwargs=None,
            self_extra_kwargs=None,
            list1_extra_kwargs=None,
            *args, **kwargs):
        self.code = code_container
        self.queue = queue

        tree = tree.get(queue)
        self.tree = tree

        self.dtype = dtype
        self.quad_order = quad_order
        self.potential_kind = potential_kind

        # {{{ digest target_kernels

        ifgrad = False
        outputs = []
        source_deriv_names = []
        k_names = []

        for out_knl in self.code.target_kernels:

            if self.is_supported_helmknl(out_knl):
                outputs.append(())
                no_target_deriv_knl = out_knl

            elif (isinstance(out_knl, AxisTargetDerivative)
                    and self.is_supported_helmknl(out_knl.inner_kernel)):
                outputs.append((out_knl.axis,))
                ifgrad = True
                no_target_deriv_knl = out_knl.inner_kernel

            else:
                raise ValueError(
                        "only the 2/3D Laplace and Helmholtz kernel "
                        "and their derivatives are supported")

            source_deriv_names.append(no_target_deriv_knl.dir_vec_name
                    if isinstance(no_target_deriv_knl, DirectionalSourceDerivative)
                    else None)

            base_knl = out_knl.get_base_kernel()
            k_names.append(base_knl.helmholtz_k_name
                    if isinstance(base_knl, HelmholtzKernel)
                    else None)

        self.outputs = outputs

        from pytools import is_single_valued

        if not is_single_valued(source_deriv_names):
            raise ValueError("not all kernels passed are the same in "
                    "whether they represent a source derivative")

        source_deriv_name = source_deriv_names[0]

        if not is_single_valued(k_names):
            raise ValueError("not all kernels passed have the same "
                    "Helmholtz parameter")

        k_name = k_names[0]

        if k_name is None:
            helmholtz_k = 0
        else:
            helmholtz_k = kernel_extra_kwargs[k_name]

        # }}}

        # {{{ table setup
        # TODO put this part into the inteferce class

        self.near_field_table = {}
        # list of tables for a single out kernel
        if isinstance(near_field_table, list):
            assert len(self.code.target_kernels) == 1
            self.near_field_table[
                self.code.target_kernels[0].__repr__()
            ] = near_field_table
            self.n_tables = len(near_field_table)

        # single table
        elif isinstance(near_field_table, NearFieldInteractionTable):
            assert len(self.code.target_kernels) == 1
            self.near_field_table[self.code.target_kernels[0].__repr__()] = [
                near_field_table
            ]
            self.n_tables = 1

        # dictionary of lists of tables
        elif isinstance(near_field_table, dict):
            self.n_tables = dict()
            for out_knl in self.code.target_kernels:
                if repr(out_knl) not in near_field_table:
                    raise RuntimeError(
                            "Missing nearfield table for %s." % repr(out_knl))
                if isinstance(near_field_table[repr(out_knl)],
                        NearFieldInteractionTable):
                    near_field_table[repr(out_knl)] = [
                            near_field_table[repr(out_knl)]]
                else:
                    assert isinstance(near_field_table[repr(out_knl)], list)

                self.n_tables[repr(out_knl)] = len(near_field_table[repr(out_knl)])

            self.near_field_table = near_field_table
        else:
            raise RuntimeError("Table type unrecognized.")

        # TODO: make all parameters table-specific (allow using inhomogeneous tables)
        kname = repr(self.code.target_kernels[0])
        self.root_table_source_box_extent = (
                self.near_field_table[kname][0].source_box_extent)
        table_starting_level = np.round(
            np.log(self.tree.root_extent / self.root_table_source_box_extent)
            / np.log(2)
            )
        for kid in range(len(self.code.target_kernels)):
            kname = self.code.target_kernels[kid].__repr__()
            for lev, table in zip(
                    range(len(self.near_field_table[kname])),
                    self.near_field_table[kname]
                    ):
                assert table.quad_order == self.quad_order

                if not table.is_built:
                    raise RuntimeError(
                        "Near field interaction table needs to be built "
                        "prior to being used"
                    )

                table_root_extent = table.source_box_extent * 2 ** lev
                assert (
                    abs(self.root_table_source_box_extent - table_root_extent)
                    < 1e-15
                )

                # If the kernel cannot be scaled,
                # - tree_root_extent must be integral times of table_root_extent
                # - n_tables must be sufficient
                if not isinstance(self.n_tables, dict) and self.n_tables > 1:
                    if (
                        not abs(
                            int(self.tree.root_extent / table_root_extent)
                            * table_root_extent
                            - self.tree.root_extent
                        )
                        < 1e-15
                    ):
                        raise RuntimeError(
                            "Incompatible list of tables: the "
                            "source_box_extent of the root table must "
                            "divide the bounding box's extent by an integer."
                        )

            if not isinstance(self.n_tables, dict) and self.n_tables > 1:
                # this checks that the boxes at the highest level are covered
                if (
                    not tree.nlevels
                    <= len(self.near_field_table[kname]) + table_starting_level
                ):
                    raise RuntimeError(
                        "Insufficient list of tables: the "
                        "finest level mesh cells at level "
                        + str(tree.nlevels)
                        + " are not covered."
                    )

                # the check that the boxes at the coarsest level are covered is
                # deferred until trav.target_boxes is passed when invoking
                # eval_direct

        if source_extra_kwargs is None:
            source_extra_kwargs = {}

        if kernel_extra_kwargs is None:
            kernel_extra_kwargs = {}

        if self_extra_kwargs is None:
            self_extra_kwargs = {}

        if list1_extra_kwargs is None:
            list1_extra_kwargs = {}

        self.list1_extra_kwargs = list1_extra_kwargs

        # }}} End table setup

        if not callable(fmm_level_to_order):
            raise TypeError("fmm_level_to_order not passed")

        dipole_vec = None
        if source_deriv_name is not None:
            dipole_vec = np.array([
                    d_i.get(queue=queue)
                    for d_i in source_extra_kwargs[source_deriv_name]],
                    order="F")

        def inner_fmm_level_to_nterms(tree, level):
            if helmholtz_k == 0:
                return fmm_level_to_order(
                        LaplaceKernel(tree.dimensions),
                        frozenset(), tree, level)
            else:
                return fmm_level_to_order(
                        HelmholtzKernel(tree.dimensions),
                        frozenset([("k", helmholtz_k)]), tree, level)

        rotation_data = None
        if 'traversal' in kwargs:
            # add rotation data if traversal is passed as a keyword argument
            from boxtree.pyfmmlib_integration import FMMLibRotationData
            rotation_data = FMMLibRotationData(self.queue, kwargs['traversal'])
        else:
            logger.warning("Rotation data is not utilized since traversal is "
                           "not known to FPNDFMMLibExpansionWrangler.")

        FMMLibExpansionWrangler.__init__(
                self, tree,

                helmholtz_k=helmholtz_k,
                dipole_vec=dipole_vec,
                dipoles_already_reordered=True,

                fmm_level_to_nterms=inner_fmm_level_to_nterms,
                rotation_data=rotation_data,

                ifgrad=ifgrad)

    # }}} End constructor

# {{{ scale factor for fmmlib

    def get_scale_factor(self):
        if self.eqn_letter == "l" and self.dim == 2:
            scale_factor = -1/(2*np.pi)
        elif self.eqn_letter == "h" and self.dim == 2:
            scale_factor = 1
        elif self.eqn_letter in ["l", "h"] and self.dim == 3:
            scale_factor = 1/(4*np.pi)
        else:
            raise NotImplementedError(
                    "scale factor for pyfmmlib %s for %d dimensions" % (
                        self.eqn_letter,
                        self.dim))

        return scale_factor

# }}} End scale factor for fmmlib

    # {{{ data vector utilities

    def multipole_expansion_zeros(self):
        return FMMLibExpansionWrangler.multipole_expansion_zeros(self)

    def local_expansion_zeros(self):
        return FMMLibExpansionWrangler.local_expansion_zeros(self)

    def output_zeros(self):
        return FMMLibExpansionWrangler.output_zeros(self)

    def reorder_sources(self, source_array):
        return FMMLibExpansionWrangler.reorder_sources(self, source_array)

    def reorder_targets(self, target_array):
        if not hasattr(self.tree, 'user_target_ids'):
            self.tree.user_target_ids = inverse_id_map(
                self.queue, self.tree.sorted_target_ids)
        return target_array[self.tree.user_target_ids]

    def reorder_potentials(self, potentials):
        return FMMLibExpansionWrangler.reorder_potentials(self, potentials)

    def finalize_potentials(self, potentials):
        # return potentials
        return FMMLibExpansionWrangler.finalize_potentials(self, potentials)

    # }}} End data vector utilities

    # {{{ formation & coarsening of multipoles

    def form_multipoles(self, level_start_source_box_nrs, source_boxes, src_weights):
        return FMMLibExpansionWrangler.form_multipoles(
            self, level_start_source_box_nrs, source_boxes, src_weights
        )

    def coarsen_multipoles(
        self, level_start_source_parent_box_nrs, source_parent_boxes, mpoles
    ):
        return FMMLibExpansionWrangler.coarsen_multipoles(
            self, level_start_source_parent_box_nrs, source_parent_boxes, mpoles
        )

    # }}} End formation & coarsening of multipoles

    # {{{ direct evaluation of near field interactions

    def eval_direct_single_out_kernel(
        self,
        out_pot,
        out_kernel,
        target_boxes,
        neighbor_source_boxes_starts,
        neighbor_source_boxes_lists,
        mode_coefs,
    ):

        # NOTE: mode_coefs are similar to source_weights BUT
        # do not include quadrature weights (purely function
        # expansiona coefficients)

        if 0:
            print("Returns range for list1")
            out_pot[:] = np.arange(len(out_pot))
            return out_pot, None

        kname = out_kernel.__repr__()

        if isinstance(self.n_tables, int) and self.n_tables > 1:
            use_multilevel_tables = True
        elif isinstance(self.n_tables, dict) and self.n_tables[kname] > 1:
            use_multilevel_tables = True
        else:
            use_multilevel_tables = False

        if use_multilevel_tables:
            # this checks that the boxes at the coarsest level
            # and allows for some round-off error
            min_lev = np.min(
                self.tree.box_levels.get(self.queue)[target_boxes.get(self.queue)]
            )
            largest_cell_extent = self.tree.root_extent * 0.5 ** min_lev
            if not self.near_field_table[kname][0].source_box_extent >= (
                largest_cell_extent - 1e-15
            ):
                raise RuntimeError(
                    "Insufficient list of tables: the "
                    "coarsest level mesh cells at level "
                    + str(min_lev)
                    + " are not covered."
                )

        # table.case_encode
        distinct_numbers = set()
        for vec in self.near_field_table[kname][0].interaction_case_vecs:
            for cvc in vec:
                distinct_numbers.add(cvc)
        base = len(range(min(distinct_numbers), max(distinct_numbers) + 1))
        shift = -min(distinct_numbers)

        case_indices_dev = cl.array.to_device(
            self.queue, self.near_field_table[kname][0].case_indices
        )

        # table.data
        table_data_combined = np.zeros(
            (
                len(self.near_field_table[kname]),
                len(self.near_field_table[kname][0].data),
            ), dtype=self.near_field_table[kname][0].data.dtype
        )
        mode_nmlz_combined = np.zeros(
            (
                len(self.near_field_table[kname]),
                len(self.near_field_table[kname][0].mode_normalizers),
            ), dtype=self.near_field_table[kname][0].mode_normalizers.dtype
        )
        for lev in range(len(self.near_field_table[kname])):
            table_data_combined[lev, :] = self.near_field_table[kname][lev].data
            mode_nmlz_combined[lev, :] = self.near_field_table[kname][
                lev
            ].mode_normalizers
        exterior_mode_nmlz_combined = np.zeros(
            (
                len(self.near_field_table[kname]),
                len(self.near_field_table[kname][0].kernel_exterior_normalizers),
            ),
            dtype=self.near_field_table[kname][0].kernel_exterior_normalizers.dtype
        )
        for lev in range(len(self.near_field_table[kname])):
            table_data_combined[lev, :] = self.near_field_table[kname][lev].data
            mode_nmlz_combined[lev, :] = \
                self.near_field_table[kname][lev].mode_normalizers
            exterior_mode_nmlz_combined[lev, :] = \
                self.near_field_table[kname][lev].kernel_exterior_normalizers

        logger.info(
                "Table data for kernel "
                + out_kernel.__repr__() + " congregated")

        # The loop domain needs to know some info about the tables being used
        table_data_shapes = {
            "n_tables": len(self.near_field_table[kname]),
            "n_q_points": self.near_field_table[kname][0].n_q_points,
            "n_table_entries": len(self.near_field_table[kname][0].data),
        }
        assert table_data_shapes["n_q_points"] == len(
            self.near_field_table[kname][0].mode_normalizers
        )

        from volumential.list1 import NearFieldFromCSR

        near_field = NearFieldFromCSR(out_kernel, table_data_shapes,
            potential_kind=self.potential_kind,
            **self.list1_extra_kwargs)

        res, evt = near_field(
            self.queue,
            result=out_pot,
            box_centers=self.tree.box_centers,
            box_levels=self.tree.box_levels,
            box_source_counts_cumul=self.tree.box_target_counts_cumul,
            box_source_starts=self.tree.box_target_starts,
            box_target_counts_cumul=self.tree.box_target_counts_cumul,
            box_target_starts=self.tree.box_target_starts,
            case_indices=case_indices_dev,
            encoding_base=base,
            encoding_shift=shift,
            mode_nmlz_combined=mode_nmlz_combined,
            exterior_mode_nmlz_combined=exterior_mode_nmlz_combined,
            neighbor_source_boxes_starts=neighbor_source_boxes_starts,
            root_extent=self.tree.root_extent,
            neighbor_source_boxes_lists=neighbor_source_boxes_lists,
            mode_coefs=mode_coefs,
            table_data_combined=table_data_combined,
            target_boxes=target_boxes,
            table_root_extent=self.root_table_source_box_extent,
        )

        if isinstance(out_pot, cl.array.Array):
            assert res is out_pot
            # FIXME: lazy evaluation sometimes returns incorrect results
            res.finish()
        else:
            assert isinstance(out_pot, np.ndarray)
            out_pot = res

        # sorted_target_ids=self.tree.user_source_ids,
        # user_source_ids=self.tree.user_source_ids)

        scale_factor = self.get_scale_factor()
        return out_pot / scale_factor, evt

    def eval_direct(
        self,
        target_boxes,
        neighbor_source_boxes_starts,
        neighbor_source_boxes_lists,
        mode_coefs,
    ):
        pot = self.output_zeros()
        if pot.dtype != np.object:
            pot = make_obj_array([pot, ])
        events = []
        for i in range(len(self.code.target_kernels)):
            # print("processing near-field of out_kernel", i)
            pot[i], evt = self.eval_direct_single_out_kernel(
                pot[i],
                self.code.target_kernels[i],
                target_boxes,
                neighbor_source_boxes_starts,
                neighbor_source_boxes_lists,
                mode_coefs,
            )
            events.append(evt)

        for out_pot in pot:
            if isinstance(out_pot, cl.array.Array):
                out_pot.finish()

        # boxtree.pyfmmlib_integration handles things diffferently
        # when target_kernels has only one element
        if len(pot) == 1:
            pot = pot[0]

        return (pot, SumpyTimingFuture(self.queue, events))

    # }}} End direct evaluation of near field interactions

    # {{{ downward pass of fmm

    def multipole_to_local(
        self,
        level_start_target_box_nrs,
        target_boxes,
        src_box_starts,
        src_box_lists,
        mpole_exps,
    ):
        return FMMLibExpansionWrangler.multipole_to_local(
            self,
            level_start_target_box_nrs,
            target_boxes,
            src_box_starts,
            src_box_lists,
            mpole_exps,
        )

    def eval_multipoles(
        self, target_boxes_by_source_level, source_boxes_by_level, mpole_exps
    ):
        return FMMLibExpansionWrangler.eval_multipoles(
            self, target_boxes_by_source_level, source_boxes_by_level, mpole_exps
        )

    def form_locals(
        self,
        level_start_target_or_target_parent_box_nrs,
        target_or_target_parent_boxes,
        starts,
        lists,
        src_weights,
    ):
        return FMMLibExpansionWrangler.form_locals(
            self,
            level_start_target_or_target_parent_box_nrs,
            target_or_target_parent_boxes,
            starts,
            lists,
            src_weights,
        )

    def refine_locals(
        self,
        level_start_target_or_target_parent_box_nrs,
        target_or_target_parent_boxes,
        local_exps,
    ):
        return FMMLibExpansionWrangler.refine_locals(
            self,
            level_start_target_or_target_parent_box_nrs,
            target_or_target_parent_boxes,
            local_exps,
        )

    def eval_locals(self, level_start_target_box_nrs, target_boxes, local_exps):
        return FMMLibExpansionWrangler.eval_locals(
            self, level_start_target_box_nrs, target_boxes, local_exps
        )

    # }}} End downward pass of fmm

    # {{{ direct evaluation of p2p (discrete) interactions

    def eval_direct_p2p(
        self, target_boxes, source_box_starts, source_box_lists, src_weights
    ):
        return FMMLibExpansionWrangler.eval_direct(
            self, target_boxes, source_box_starts, source_box_lists, src_weights
        )

    # }}} End direct evaluation of p2p interactions

    @staticmethod
    def is_supported_helmknl(knl):
        if isinstance(knl, DirectionalSourceDerivative):
            knl = knl.inner_kernel

        return (isinstance(knl, (LaplaceKernel, HelmholtzKernel))
                and knl.dim in (2, 3))

# }}} End fmmlib backend (for laplace, helmholtz)

class FPNDExpansionWranglerCodeContainer(FPNDSumpyExpansionWranglerCodeContainer):
    """The default code container.
    """

class FPNDExpansionWrangler(FPNDSumpyExpansionWrangler):
    """The default wrangler class.
    """

# vim: filetype=pyopencl:foldmethod=marker


r"""Functions for spectator scattering corrections to $B\to V\ell^+\ell^-$ decays.

This includes weak annihilation, chromomagnetic contributions, and light
quark-loop spectator scattering.
"""

import flavio
import numpy as np
from flavio.classes import AuxiliaryQuantity, Implementation
from flavio.physics.bdecays.common import meson_quark
from flavio.physics.bdecays.wilsoncoefficients import wctot_dict
from flavio.physics.common import conjugate_par, conjugate_wc, add_dict
from flavio.config import config

# Auxiliary quantities and implementations

# function needed for the QCD factorization implementation (see qcdf.py)
def ha_qcdf_function(B, V):
    scale = config['renormalization scale']['bvll']
    label = meson_quark[(B,V)] + 'ee' # the lepton flavour is irrelevant here
                                      # as only dipole and 4-quark operators contribute!
    def function(wc_obj, par_dict, q2, cp_conjugate):
        par = par_dict.copy()
        if cp_conjugate:
            par = conjugate_par(par)
        wc = wctot_dict(wc_obj, label, scale, par)
        if cp_conjugate:
            wc = conjugate_wc(wc)
        return flavio.physics.bdecays.bvll.qcdf.helicity_amps_qcdf(q2, wc, par, B, V)
    return function

# ... and the same for the interpolated version (see qcdf_interpolate.py)
def ha_qcdf_interpolate_function(B, V, contribution='all'):
    scale = config['renormalization scale']['bvll']
    def function(wc_obj, par_dict, q2, cp_conjugate):
        return flavio.physics.bdecays.bvll.qcdf_interpolate.helicity_amps_qcdf(q2, par_dict, B, V, cp_conjugate, contribution)
    return function

# loop over hadronic transitions and lepton flavours
# BTW, it is not necessary to loop over tau: for tautau final states, the minimum
# q2=4*mtau**2 is so high that QCDF is not valid anymore anyway!
for had in [('B0','K*0'), ('B+','K*+'), ('B0','rho0'), ('B+','rho+'), ('Bs','phi'), ]:
    process = had[0] + '->' + had[1] + 'll' # e.g. B0->K*0mumu
    quantity = process + ' spectator scattering'
    a = AuxiliaryQuantity(name=quantity, arguments=['q2', 'cp_conjugate'])
    a.description = ('Contribution to ' + process + ' helicity amplitudes from'
                    ' non-factorizable spectator scattering.')

    # Implementation: QCD factorization
    iname = process + ' QCDF'
    i = Implementation(name=iname, quantity=quantity,
                   function=ha_qcdf_function(B=had[0], V=had[1]))
    i.set_description("QCD factorization")

    # Implementation: interpolated QCD factorization
    iname = process + ' QCDF interpolated'
    i = Implementation(name=iname, quantity=quantity,
                   function=ha_qcdf_interpolate_function(B=had[0], V=had[1]))
    i.set_description("Interpolated version of QCD factorization")


# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Visualization of the results 3D VTK form"""

import os
from pyevtk.hl import gridToVTK
import numpy as np

def vtk_structure(grid_tensor, eh_tensor, path_res):
    r"""
    Generates 3D vtk file for visualizaiton.

    Args:
        grid_tensor (numpy.ndarray): grid data (shape: (dim_t, dim_x, dim_y, dim_z, 4)).
        eh_tensor (numpy.ndarray): electric and magnetic data (np.array, shape: (dim_t, dim_x, dim_y, dim_z, 6)).
        path_res (str): save path for the output vtk file.

    Supported Platforms:
        ``Ascend``

    Examples:
        >>> import numpy as np
        >>> from mindelec.vision import vtk_structure
        >>> grid_tensor = np.random.rand(20, 10, 10, 10, 4).astype(np.float32)
        >>> eh_tensor = np.random.rand(20, 10, 10, 10, 6).astype(np.float32)
        >>> path_res = './result_vtk'
        >>> vtk_structure(grid_tensor, eh_tensor, path_res)
    """
    if not isinstance(grid_tensor, np.ndarray):
        raise TypeError("The type of grid_tensor should be numpy array, but get {}".format(type(grid_tensor)))

    if not isinstance(eh_tensor, np.ndarray):
        raise TypeError("The type of eh_tensor should be numpy array, but get {}".format(type(eh_tensor)))

    if not isinstance(path_res, str):
        raise TypeError("The type of path_res should be str, but get {}".format(type(path_res)))
    if not os.path.exists(path_res):
        os.makedirs(path_res)

    input_grid = grid_tensor
    output_grid = eh_tensor

    shape_grid = input_grid.shape
    shape_eh = output_grid.shape

    if len(shape_grid) != 5 or shape_grid[-1] != 4:
        raise ValueError("grid_tensor shape should be (dim_t, dim_x, dim_y, dim_z, 4), but get {}"
                         .format(shape_grid))

    if len(shape_eh) != 5 or shape_eh[-1] != 6:
        raise ValueError("eh_tensor shape should be (dim_t, dim_x, dim_y, dim_z, 6), but get {}"
                         .format(shape_eh))

    if shape_grid[:4] != shape_eh[:4]:
        raise ValueError("grid_tensor and eh_tensor should have the same dimension except the last axis, "
                         "but get grid_tensor shape {} and eh_tensor shape{}".format(shape_grid, shape_eh))

    (dim_t, dim_x, dim_y, dim_z, d) = input_grid.shape
    input_grid = np.reshape(input_grid, (dim_t * dim_x * dim_y * dim_z, d))
    x_min, x_max = np.min(input_grid[:, 0]), np.max(input_grid[:, 0])
    y_min, y_max = np.min(input_grid[:, 1]), np.max(input_grid[:, 1])
    z_min, z_max = np.min(input_grid[:, 2]), np.max(input_grid[:, 2])

    x_all = np.linspace(x_min, x_max, dim_x, endpoint=True, dtype='float64')
    y_all = np.linspace(y_min, y_max, dim_y, endpoint=True, dtype='float64')
    z_all = np.linspace(z_min, z_max, dim_z, endpoint=True, dtype='float64')

    x = np.zeros((dim_x, dim_y, dim_z))
    y = np.zeros((dim_x, dim_y, dim_z))
    z = np.zeros((dim_x, dim_y, dim_z))

    for i in range(dim_x):
        for j in range(dim_y):
            for k in range(dim_z):
                x[i, j, k] = x_all[i]
                y[i, j, k] = y_all[j]
                z[i, j, k] = z_all[k]

    for t in range(dim_t):
        output_grid_show = output_grid[t]
        ex, ey, ez = output_grid_show[:, :, :, 0], output_grid_show[:, :, :, 1], output_grid_show[:, :, :, 2]
        hx, hy, hz = output_grid_show[:, :, :, 3], output_grid_show[:, :, :, 4], output_grid_show[:, :, :, 5]
        ex, ey, ez = ex.astype(np.float64), ey.astype(np.float64), ez.astype(np.float64)
        hx, hy, hz = hx.astype(np.float64), hy.astype(np.float64), hz.astype(np.float64)
        gridToVTK(os.path.join(path_res, 'eh_t' + str(t)),
                  x, y, z,
                  pointData={"Ex": ex, "Ey": ey, "Ez": ez, "Hx": hx, "Hy": hy, "Hz": hz})


"""
run_crack_lotf.py
Script to run LOTF molecular dynamics for a crack slab,
incrementing the load in small steps until fracture starts.
James Kermode <james.kermode@kcl.ac.uk>
February 2013
"""
import numpy as np

from ase.constraints import FixAtoms
from ase.md.verlet import VelocityVerlet
from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
import ase.units as units

from quippy import set_fortran_indexing
from quippy.atoms import Atoms
from quippy.potential import Potential
from quippy.io import AtomsWriter

from quippy.crack import (get_strain,
                          get_energy_release_rate,
                          ConstantStrainRate,
                          find_crack_tip_stress_field)

# additional requirements for the QM/MM simulation:
from quippy.potential import ForceMixingPotential
from quippy.lotf import LOTFDynamics, update_hysteretic_qm_region

# ******* Start of parameters ***********

input_file = 'crack.xyz'         # File from which to read crack slab structure
sim_T = 300.0*units.kB           # Simulation temperature
nsteps = 1000000                   # Total number of timesteps to run for
timestep = 1.0*units.fs          # Timestep (NB: time base units are not fs!)
cutoff_skin = 2.0*units.Ang      # Amount by which potential cutoff is increased
                                 # for neighbour calculations
tip_move_tol = 10.0              # Distance tip has to move before crack
                                 # is taken to be running
strain_rate = 1e-5*(1/units.fs)  # Strain rate
traj_file = 'traj.xyz'           # Trajectory output file (NetCDF or XYZ format)
traj_interval = 10               # Number of time steps between
                                 # writing output frames
param_file = 'params.xml'        # Filename of XML file containing
                                 # potential parameters
mm_init_args = 'IP SW'           # Initialisation arguments for
                                 # classical potential

# additional parameters for the QM/MM simulation:
qm_init_args = 'TB DFTB'         # Initialisation arguments for QM potential
qm_inner_radius = 6.0*units.Ang  # Inner hysteretic radius for QM region
qm_outer_radius = 8.0*units.Ang  # Inner hysteretic radius for QM region
extrapolate_steps = 10           # Number of steps for predictor-corrector
                                 # interpolation and extrapolation

# ******* End of parameters *************

set_fortran_indexing(False)

# ********** Read input file ************

print 'Loading atoms from file %s' % input_file
atoms = Atoms(input_file)

orig_height = atoms.info['OrigHeight']
orig_crack_pos = atoms.info['CrackPos'].copy()

# ***** Setup constraints *******

top = atoms.positions[:, 1].max()
bottom = atoms.positions[:, 1].min()
left = atoms.positions[:, 0].min()
right = atoms.positions[:, 0].max()

# fix atoms in the top and bottom rows
fixed_mask = ((abs(atoms.positions[:, 1] - top) < 1.0) |
              (abs(atoms.positions[:, 1] - bottom) < 1.0))
fix_atoms = FixAtoms(mask=fixed_mask)
print('Fixed %d atoms\n' % fixed_mask.sum())
atoms.set_constraint([fix_atoms])

# Increase epsilon_yy applied to all atoms at constant strain rate
strain_atoms = ConstantStrainRate(orig_height, strain_rate*timestep)

# ******* Set up potentials and calculators ********

mm_pot = Potential(mm_init_args,
                   param_filename=param_file,
                   cutoff_skin=cutoff_skin)

# Density functional tight binding (DFTB) potential
qm_pot = Potential(qm_init_args,
                   param_filename=param_file)

# Construct the QM/MM potential, which mixes QM and MM forces.
# The qm_args_str parameters control how the QM calculation is carried out:
# we use a single cluster, periodic in the z direction and terminated
# with hydrogen atoms. The positions of the outer layer of buffer atoms
# are not randomised.
qmmm_pot = ForceMixingPotential(pot1=mm_pot,
                                pot2=qm_pot,
                                qm_args_str='single_cluster cluster_periodic_z carve_cluster '+
                                            'terminate cluster_hopping=F randomise_buffer=F',
                                fit_hops=4,
                                lotf_spring_hops=3,
                                hysteretic_buffer=True,
                                hysteretic_buffer_inner_radius=7.0,
                                hysteretic_buffer_outer_radius=9.0,
                                cluster_hopping_nneighb_only=False,
                                min_images_only=True)

# Use the force mixing potential as the Atoms' calculator
atoms.set_calculator(qmmm_pot)
qmmm_pot.atoms = atoms

# *** Set up the initial QM region ****

qm_list = update_hysteretic_qm_region(atoms, [], orig_crack_pos,
                                      qm_inner_radius, qm_outer_radius)

# ********* Setup and run MD ***********

# Set the initial temperature to 2*simT: it will then equilibriate to
# simT, by the virial theorem
MaxwellBoltzmannDistribution(atoms, 2.0*sim_T)

# Initialise the dynamical system
dynamics = LOTFDynamics(atoms, timestep, extrapolate_steps)

# array to store time averaged stress field
avg_sigma = np.zeros((len(atoms), 3, 3))

# Print some information every time step
def printstatus():
    if dynamics.nsteps == 1:
        print """
State      Time/fs    Temp/K     Strain      G/(J/m^2)  CrackPos/A D(CrackPos)/A
---------------------------------------------------------------------------------"""

    log_format = ('%(label)-4s%(time)12.1f%(temperature)12.6f'+
                  '%(strain)12.5f%(G)12.4f%(crack_pos_x)12.2f    (%(d_crack_pos_x)+5.2f)')

    atoms.info['label'] = dynamics.state_label  # Label for the status line
    atoms.info['time'] = dynamics.get_time()/units.fs
    atoms.info['temperature'] = (atoms.get_kinetic_energy() /
                                 (1.5*units.kB*len(atoms)))
    atoms.info['strain'] = get_strain(atoms)
    atoms.info['G'] = get_energy_release_rate(atoms)/(units.J/units.m**2)

    crack_pos = find_crack_tip_stress_field(atoms, calc=mm_pot,
                                            avg_sigma=avg_sigma)
    atoms.info['crack_pos_x'] = crack_pos[0]
    atoms.info['d_crack_pos_x'] = crack_pos[0] - orig_crack_pos[0]

    print log_format % atoms.info

dynamics.attach(printstatus)

def atom_straining(atoms):
    crack_pos = find_crack_tip_stress_field(atoms, calc=mm_pot,
                                            avg_sigma=avg_sigma)
    # keep straining until the crack tip has advanced to tip_move_tol
    if not atoms.info['is_cracked'] and (crack_pos[0] - orig_crack_pos[0]) < tip_move_tol:
      strain_atoms.apply_strain(atoms)
    elif not atoms.info['is_cracked']:
      atoms.info['is_cracked'] = True

dynamics.attach(atom_straining, 1, dynamics.atoms)

# Function to update the QM region at the beginning of each extrapolation cycle
def update_qm_region(atoms):
   crack_pos = find_crack_tip_stress_field(atoms, calc=mm_pot,
                                           avg_sigma=avg_sigma)
   qm_list = qmmm_pot.get_qm_atoms(atoms)
   qm_list = update_hysteretic_qm_region(atoms, qm_list, crack_pos,
                                         qm_inner_radius, qm_outer_radius)
   qmmm_pot.set_qm_atoms(qm_list, atoms)
   #assert (atoms.hybrid == 1).sum() == len(qm_list)

dynamics.set_qm_update_func(update_qm_region)

# Save frames to the trajectory every `traj_interval` time steps
# but only when interpolating
trajectory = AtomsWriter(traj_file)

def traj_writer(dynamics):
   if dynamics.state == LOTFDynamics.Interpolation:
      # copy time-averaged stress into Atoms so it will be written to trajectory file
      dynamics.atoms.set_array('avg_sigma', avg_sigma.reshape((len(atoms), 9)))
      trajectory.write(dynamics.atoms)

dynamics.attach(traj_writer, traj_interval, dynamics)

# Start running!
dynamics.run(nsteps)


# Copyright 2016-2020 Blue Marble Analytics LLC.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import csv
import os.path
from pyomo.environ import Param, Set, NonNegativeReals, PercentFraction, Expression

def generic_add_model_components(
    m,
    d,
    reserve_zone_set,
    reserve_requirement_tmp_param,
    reserve_requirement_percent_param,
    reserve_zone_load_zone_set,
    ba_prj_req_contribution_set,
    prj_power_param,
    prj_capacity_param,
    reserve_requirement_expression,
):
    """
    :param m:
    :param d:
    :param reserve_zone_set:
    :param reserve_requirement_tmp_param:
    :param reserve_requirement_percent_param:
    :param reserve_zone_load_zone_set:
    :param ba_prj_req_contribution_set:
    :param prj_power_param:
    :param prj_capacity_param:
    :param reserve_requirement_expression:
    :return:

    Generic treatment of reserves. This function creates model components
    related to a particular reserve requirement, including
    1) the reserve requirement by zone and timepoint, if any
    2) the reserve requirement as a percent of load and map for which load
    zones' load to consider
    3) the contributions to the reserve requirement from projects: there are two
    types of these contributions, those based on the power output in the timepoint
    and those based on the project capacity.
    """

    # Magnitude of the requirement by reserve zone and timepoint
    # If not specified for a reserve zone - timepoint combination,
    # will default to 0
    setattr(
        m,
        reserve_requirement_tmp_param,
        Param(getattr(m, reserve_zone_set), m.TMPS, within=NonNegativeReals, default=0),
    )

    # Requirement as percentage of load
    setattr(
        m,
        reserve_requirement_percent_param,
        Param(getattr(m, reserve_zone_set), within=PercentFraction, default=0),
    )

    # Load zones included in the reserve percentage requirement
    setattr(
        m,
        reserve_zone_load_zone_set,
        Set(dimen=2, within=getattr(m, reserve_zone_set) * m.LOAD_ZONES),
    )

    # Projects contributing to BA requirement based on power output in the timepoint
    # and on capacity in the period
    setattr(
        m,
        ba_prj_req_contribution_set,
        Set(dimen=2, within=getattr(m, reserve_zone_set) * m.PROJECTS),
    )

    setattr(
        m,
        prj_power_param,
        Param(
            getattr(m, ba_prj_req_contribution_set), within=PercentFraction, default=0
        ),
    )

    setattr(
        m,
        prj_capacity_param,
        Param(
            getattr(m, ba_prj_req_contribution_set), within=PercentFraction, default=0
        ),
    )

    def reserve_requirement_rule(mod, reserve_zone, tmp):
        # If we have a map of reserve zones to load zones, apply the percentage
        # target; if no map provided, the percentage_target is 0
        if getattr(mod, reserve_zone_load_zone_set):
            percentage_target = sum(
                getattr(mod, reserve_requirement_percent_param)[reserve_zone]
                * mod.static_load_mw[lz, tmp]
                for (_reserve_zone, lz) in getattr(mod, reserve_zone_load_zone_set)
                if _reserve_zone == reserve_zone
            )
        else:
            percentage_target = 0

        # Project contributions, if any projects in the respective set
        if getattr(mod, ba_prj_req_contribution_set):
            # Project contributions to requirement based on power output
            prj_pwr_contribution = sum(
                getattr(mod, prj_power_param)[reserve_zone, prj]
                * mod.Power_Provision_MW[prj, tmp]
                for (_reserve_zone, prj) in getattr(mod, ba_prj_req_contribution_set)
                if _reserve_zone == reserve_zone
                if (prj, tmp) in mod.PRJ_OPR_TMPS
            )

            # Project contributions to requirement based on (available) capacity
            # We are not holding the extra reserves when projects are unavailable
            prj_cap_contribution = sum(
                getattr(mod, prj_capacity_param)[reserve_zone, prj]
                * mod.Capacity_MW[prj, mod.period[tmp]]
                * mod.Availability_Derate[prj, tmp]
                for (_reserve_zone, prj) in getattr(mod, ba_prj_req_contribution_set)
                if _reserve_zone == reserve_zone
                if (prj, tmp) in mod.PRJ_OPR_TMPS
            )
        else:
            prj_pwr_contribution = 0
            prj_cap_contribution = 0

        return (
            getattr(mod, reserve_requirement_tmp_param)[reserve_zone, tmp]
            + percentage_target
            + prj_pwr_contribution
            + prj_cap_contribution
        )

    setattr(
        m,
        reserve_requirement_expression,
        Expression(
            getattr(m, reserve_zone_set) * m.TMPS, rule=reserve_requirement_rule
        ),
    )

def generic_load_model_data(
    m,
    d,
    data_portal,
    scenario_directory,
    subproblem,
    stage,
    reserve_requirement_param,
    reserve_zone_load_zone_set,
    reserve_requirement_percent_param,
    ba_prj_req_contribution_set,
    prj_power_param,
    prj_capacity_param,
    reserve_type,
):
    """

    :param m:
    :param d:
    :param data_portal:
    :param scenario_directory:
    :param subproblem:
    :param stage:
    :param reserve_requirement_param:
    :param reserve_zone_load_zone_set:
    :param reserve_requirement_percent_param
    :param ba_prj_req_contribution_set
    :param prj_power_param
    :param prj_capacity_param
    :param reserve_type:
    :return:
    """
    input_dir = os.path.join(scenario_directory, str(subproblem), str(stage), "inputs")

    # Load by-tmp requriement if input file was written
    by_tmp_req_filename = os.path.join(
        input_dir, "{}_tmp_requirement.tab".format(reserve_type)
    )
    if os.path.exists(by_tmp_req_filename):
        tmp_params_to_load = (
            (
                getattr(m, reserve_requirement_param),
                m.frequency_response_requirement_partial_mw,
            )
            if reserve_type == "frequency_response"
            else getattr(m, reserve_requirement_param)
        )
        data_portal.load(filename=by_tmp_req_filename, param=tmp_params_to_load)

    # If we have a RPS zone to load zone map input file, load it and the
    # percent requirement; otherwise, initialize the set as an empty list (
    # the param defaults to 0)
    map_filename = os.path.join(input_dir, "{}_percent_map.tab".format(reserve_type))
    if os.path.exists(map_filename):
        data_portal.load(
            filename=map_filename, set=getattr(m, reserve_zone_load_zone_set)
        )
        data_portal.load(
            filename=os.path.join(
                input_dir, "{}_percent_requirement.tab".format(reserve_type)
            ),
            param=getattr(m, reserve_requirement_percent_param),
        )
    else:
        data_portal.data()[reserve_zone_load_zone_set] = {None: []}

    # If we have a project contributions file, load it into the respective
    prj_contr_filename = os.path.join(
        input_dir, "{}_requirement_project_contributions.tab".format(reserve_type)
    )
    if os.path.exists(prj_contr_filename):
        data_portal.load(
            filename=prj_contr_filename,
            index=getattr(m, ba_prj_req_contribution_set),
            param=(getattr(m, prj_power_param), getattr(m, prj_capacity_param)),
        )
    else:
        data_portal.data()[ba_prj_req_contribution_set] = {None: []}

def generic_get_inputs_from_database(
    scenario_id,
    subscenarios,
    subproblem,
    stage,
    conn,
    reserve_type,
    reserve_type_ba_subscenario_id,
    reserve_type_req_subscenario_id,
):
    """
    :param subscenarios:
    :param subproblem:
    :param stage:
    :param conn:
    :param reserve_type:
    :param reserve_type_ba_subscenario_id:
    :param reserve_type_req_subscenario_id:
    :return:
    """
    subproblem = 1 if subproblem == "" else subproblem
    stage = 1 if stage == "" else stage
    c = conn.cursor()

    partial_freq_resp_extra_column = (
        ", frequency_response_partial_mw"
        if reserve_type == "frequency_response"
        else ""
    )

    tmp_req = c.execute(
        """SELECT {}_ba, timepoint, {}_mw{}
        FROM inputs_system_{}
        INNER JOIN
        (SELECT timepoint
        FROM inputs_temporal
        WHERE temporal_scenario_id = {}
        AND subproblem_id = {}
        AND stage_id = {}) as relevant_timepoints
        USING (timepoint)
        INNER JOIN
        (SELECT {}_ba
        FROM inputs_geography_{}_bas
        WHERE {}_ba_scenario_id = {}) as relevant_bas
        USING ({}_ba)
        WHERE {}_scenario_id = {}
        AND stage_id = {}
        """.format(
            reserve_type,
            reserve_type,
            partial_freq_resp_extra_column,
            reserve_type,
            subscenarios.TEMPORAL_SCENARIO_ID,
            subproblem,
            stage,
            reserve_type,
            reserve_type,
            reserve_type,
            reserve_type_ba_subscenario_id,
            reserve_type,
            reserve_type,
            reserve_type_req_subscenario_id,
            stage,
        )
    )

    c2 = conn.cursor()
    # Get any percentage requirement
    percentage_req = c2.execute(
        """
        SELECT {}_ba, percent_load_req
        FROM inputs_system_{}_percent
        WHERE {}_scenario_id = {}
        """.format(
            reserve_type, reserve_type, reserve_type, reserve_type_req_subscenario_id
        )
    )

    # Get any reserve zone to load zone mapping for the percent target
    c3 = conn.cursor()
    lz_mapping = c3.execute(
        """
        SELECT {}_ba, load_zone
        FROM inputs_system_{}_percent_lz_map
        JOIN
        (SELECT {}_ba
        FROM inputs_geography_{}_bas
        WHERE {}_ba_scenario_id = {}) as relevant_bas
        USING ({}_ba)
        WHERE {}_scenario_id = {}
        """.format(
            reserve_type,
            reserve_type,
            reserve_type,
            reserve_type,
            reserve_type,
            reserve_type_ba_subscenario_id,
            reserve_type,
            reserve_type,
            reserve_type_req_subscenario_id,
        )
    )

    # Get any project contributions to the magnitude of the reserve requirement
    c4 = conn.cursor()
    project_contributions = c4.execute(
        """
        SELECT {reserve_type}_ba, project, percent_power_req, percent_capacity_req
        FROM inputs_system_{reserve_type}_project
        JOIN (
        SELECT {reserve_type}_ba
        FROM inputs_geography_{reserve_type}_bas
        WHERE {reserve_type}_ba_scenario_id = {reserve_type_ba_subscenario_id}
        ) as relevant_bas
        USING ({reserve_type}_ba)
        JOIN (
        SELECT project
        FROM inputs_project_portfolios
        WHERE project_portfolio_scenario_id = (
                SELECT project_portfolio_scenario_id
                FROM scenarios
                WHERE scenario_id = {scenario_id}
            )
        ) as relevant_prj
        USING (project)
        WHERE {reserve_type}_scenario_id = {reserve_type_req_subscenario_id}
        """.format(
            reserve_type=reserve_type,
            reserve_type_ba_subscenario_id=reserve_type_ba_subscenario_id,
            scenario_id=scenario_id,
            reserve_type_req_subscenario_id=reserve_type_req_subscenario_id,
        )
    )

    return tmp_req, percentage_req, lz_mapping, project_contributions

def generic_write_model_inputs(
    scenario_directory,
    subproblem,
    stage,
    timepoint_req,
    percent_req,
    percent_map,
    project_contributions,
    reserve_type,
):
    """
    Get inputs from database and write out the model input
    lf_reserves_down_requirement.tab file.
    :param scenario_directory: string, the scenario directory
    :param subproblem:
    :param stage:
    :param timepoint_req:
    :param percent_req:
    :param percent_map:
    :param project_contributions:
    :param reserve_type:
    :return:
    """
    inputs_dir = os.path.join(scenario_directory, str(subproblem), str(stage), "inputs")

    # Write the by-timepoint requirement file if by-tmp requirement specified
    timepoint_req = timepoint_req.fetchall()
    if timepoint_req:
        with open(
            os.path.join(inputs_dir, "{}_tmp_requirement.tab".format(reserve_type)),
            "w",
            newline="",
        ) as tmp_req_file:
            writer = csv.writer(tmp_req_file, delimiter="\t", lineterminator="\n")

            # Write header
            extra_column = (
                ["partial_requirement"] if reserve_type == "frequency_response" else []
            )
            writer.writerow(["ba", "timepoint", "requirement"] + extra_column)

            for row in timepoint_req:
                writer.writerow(row)

    # Write the percent requirement files only if there's a mapping
    ba_lz_map_list = [row for row in percent_map]

    if ba_lz_map_list:
        with open(
            os.path.join(inputs_dir, "{}_percent_requirement.tab".format(reserve_type)),
            "w",
            newline="",
        ) as percent_req_file:
            writer = csv.writer(percent_req_file, delimiter="\t", lineterminator="\n")

            # Write header
            writer.writerow(["ba", "percent_requirement"])

            for row in percent_req:
                writer.writerow(row)

        with open(
            os.path.join(inputs_dir, "{}_percent_map.tab".format(reserve_type)),
            "w",
            newline="",
        ) as percent_map_file:
            writer = csv.writer(percent_map_file, delimiter="\t", lineterminator="\n")

            # Write header
            writer.writerow(["ba", "load_zone"])

            for row in ba_lz_map_list:
                writer.writerow(row)
    else:
        pass

    # Project contributions to the magnitude requirement
    project_contributions = project_contributions.fetchall()

    prj_contributions = False
    for (ba, prj, pwr, cap) in project_contributions:
        if pwr is not None or cap is not None:
            prj_contributions = True

    if prj_contributions:
        with open(
            os.path.join(
                inputs_dir,
                "{}_requirement_project_contributions.tab".format(reserve_type),
            ),
            "w",
            newline="",
        ) as prj_file:
            writer = csv.writer(prj_file, delimiter="\t", lineterminator="\n")

            # Write header
            writer.writerow(
                ["ba", "project", "percent_power_req", "percent_capacity_req"]
            )
            for (ba, prj, pwr, cap) in project_contributions:
                if pwr is None:
                    pwr = "."
                if cap is None:
                    cap = "."
                writer.writerow([ba, prj, pwr, cap])


"""
  Copyright (c) 2015 Semafor Informatik & Energie AG

"""
import sys
import numpy as np
import copy
import operator as op

class Individual:
    def __init__(self, n_dim, f_dim):
        self.cur_f = [0]*f_dim
        self.cur_x = [0]*n_dim
        self.cur_v = []
        self.cur_c = []

        self.rank = 0
        self.crowd_d = 0
        self.idx = 0

    def __str__(self):
        return "f {} x {} rank {} crowd_d {}".format(self.cur_f, self.cur_x,
                                                     self.rank, self.crowd_d)

class Population:

    def __init__(self, probl, size, seed=None):
        self.individuals = []
        self.problem = probl
        self.champion = None
        self.dom_count = []
        self.dom_list = []
        np.random.seed(seed)
        for s in range(size):
            self.append([np.random.uniform(lb, ub)
                         for lb, ub in zip(self.problem.lower,
                                           self.problem.upper)])
#        kmax=int(np.sqrt(size))+1
#        f=[]
#        for i in range(kmax):
#            for k in range(kmax):
#                f.append((i*0.8/kmax +0.1, k*0.8/kmax +0.1))
#        for k,i in enumerate(self.individuals):
#            i.cur_f = f[k]

    def size(self):
        return len(self.individuals)

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

    def append(self, child):
        self.individuals.append(Individual(self.problem.dimension,
                                           self.problem.f_dim))
        self.individuals[-1].cur_x = child
        self.individuals[-1].idx = len(self.individuals)-1
        self.init_velocity()

    def eval(self):
        for k, i in enumerate(self.individuals):
            i.cur_f = self.problem.objfun(i.cur_x)
        self.update()

    def best_idx(self):
        return [i.idx for i in sorted(self.individuals,
                                      key=op.attrgetter('rank',
                                                        'crowd_d'))]

    def merge(self, pop):
        "sort by rank and crowding distance (proximity)"
        self.individuals += pop.individuals
        self.update()
        for i in self.individuals:
            if abs(i.crowd_d) > sys.float_info.epsilon:
                i.crowd_d = 1/i.crowd_d
            else:
                i.crowd_d = sys.float_info.max

        best = sorted(self.individuals,
                      key=op.attrgetter('rank',
                                        'crowd_d'))
        self.individuals = best[:pop.size()]
        self.update()

    def init_velocity(self):
        for i, j in enumerate(self.individuals[-1].cur_x):
            w = (self.problem.upper[i] - self.problem.lower[i])/2
            self.individuals[-1].cur_v.append(np.random.uniform(-w, w))

    def populate(self, x, f, s):
        """populate with decision and objective values multiplied by sign
        x: vector of decision values
        f: matrix of objective values
        s: vector of signs (1, -1)"""
        for k, i in enumerate(self.individuals):
            i.cur_x = x[k]
            i.cur_f = [v*sign
                       for v, sign in zip(f[:, k], s)]
        self.update()

    def get_ranked_decisions(self):
        px = dict()
        for i in self.individuals:
            k = i.rank
            if k in px:
                px[k].append(i.cur_x)
            else:
                px[k] = [i.cur_x]
        return px

    def get_ranked_objectives(self, s):
        po = dict()
        for i in self.individuals:
            cur_f = [v*sign
                     for v, sign in zip(i.cur_f, s)]
            k = i.rank
            if k in po:
                po[k].append(cur_f)
            else:
                po[k] = [cur_f]
        return po

    def update(self):
        size = len(self.individuals)
        self.dom_count = []
        self.dom_list = [[] for s in range(size)]
        self.champion = None
        for s in range(size):
            self.dom_count.append(0)
            self.update_dom(s)
            self.update_champion(s)
        self.update_pareto_information()

    def update_dom(self, n):
        """Loop over the population (j) and construct
        dom_list[n] and dom_count """

        for i, x in enumerate(self.individuals):
            if i != n:
                # check if individual in position i
                # dominates the one in position n
                if self.problem.compare_fc(x.cur_f, x.cur_c,
                                           self.individuals[n].cur_f,
                                           self.individuals[n].cur_c):
                    self.dom_count[n] += 1
                    self.dom_list[i].append(n)

    def update_champion(self, idx):
        if self.champion is None or \
            self.problem.compare_fc(self.individuals[idx].cur_f,
                                    self.individuals[idx].cur_c,
                                    self.champion['f'], self.champion['c']):
            self.champion = dict(x=self.individuals[idx].cur_x,
                                 f=self.individuals[idx].cur_f,
                                 c=self.individuals[idx].cur_c)

    def update_crowding(self, F):
        # sort along the fitness dimension
        for i in range(self.problem.f_dim):
            I = sorted(F, key=lambda k: self.individuals[k].cur_f[i],
                       reverse=True)
            self.individuals[I[0]].crowd_d = sys.float_info.max
            self.individuals[I[-1]].crowd_d = sys.float_info.max
            df = (self.individuals[I[-1]].cur_f[i] -
                  self.individuals[I[0]].cur_f[i])
            for j in range(1, len(F)-1):
                if abs(df) > sys.float_info.epsilon:
                    self.individuals[I[j]].crowd_d += (
                        self.individuals[I[j+1]].cur_f[i] -
                        self.individuals[I[j-1]].cur_f[i])/df

    def update_pareto_information(self):
        size = len(self.individuals)
        self.pareto_rank = [0]*size
        F = [i for i, c in enumerate(self.dom_count) if c == 0]
        irank = 1
        dom_count_copy = list(self.dom_count)
        while True:
            self.update_crowding(F)
            S = []
            for i, f in enumerate(F):
                for j, k in enumerate(self.dom_list[f]):
                    dom_count_copy[k] -= 1
                    if dom_count_copy[k] == 0:
                        S.append(k)
                        self.pareto_rank[k] = irank
                        self.individuals[k].rank = irank
            if S:
                F = list(S)
            else:
                return

            irank += 1

    def compute_pareto_fronts(self):
        self.update_pareto_information()
        retval = [[] for s in range(max(self.pareto_rank)+1)]
        for i, j in enumerate(self.individuals):
            retval[self.pareto_rank[i]].append(i)
        return retval

    def compute_ideal(self):
        return [min(x)
                for x in zip(*[i.cur_f
                               for i in self.individuals if i.rank == 0])]

    def compute_nadir(self):
        return [max(x)
                for x in zip(*[i.cur_f
                               for i in self.individuals if i.rank == 0])]

    def compute_worst(self):
        return [max(x)
                for x in zip(*[i.cur_f
                               for i in self.individuals])]

    def compute_norm_dist(self):
        zi = np.array(self.compute_ideal())
        zw = np.array(self.compute_worst())
        znad = np.array(self.compute_nadir())
        return np.sqrt(((znad-zi)**2).sum() / ((zw-zi)**2).sum())

    def plot_pareto_fronts(
            self,
            rgb=(
                0,
                0,
                0),
            comp=[
                0,
                1],
            symbol='o',
            size=6,
            fronts=[]):
        """
        Plots the population pareto front in a 2-D graph

        USAGE: pop.plot_pareto_front(comp = [0,1], rgb=(0,1,0))

        * comp: components of the fitness function to plot in the 2-D window
        * rgb: specify the color of the 1st front (use strong colors here)
        * symbol: marker for the individual
        * size: size of the markersymbol
        * fronts: list of fronts to be plotted (use [0] to only show the first)
        """
        from numpy import linspace
        import matplotlib.pyplot as plt

        if len(comp) != 2:
            raise ValueError(
                'Invalid components of the objective function selected for plot')

        p_dim = self.problem.f_dimension

        if p_dim == 1:
            raise ValueError(
                'Pareto fronts of a 1-dimensional problem cannot be plotted')

        if not all([c in range(0, p_dim) for c in comp]):
            raise ValueError(
                'You need to select valid components of the objective function')

        p_list = self.compute_pareto_fronts()
        if (len(fronts) > 0):
            n = len(p_list)
            consistent = [d < n for d in fronts]
            if consistent.count(False) > 0:
                raise ValueError(
                    'Check your fronts list, there seem to be not enough fronts')
            p_list = [p_list[idx] for idx in fronts]

        cl = list(zip(linspace(0.9 if rgb[0] else 0.1, 0.9, len(p_list)),
                      linspace(0.9 if rgb[1] else 0.1, 0.9, len(p_list)),
                      linspace(0.9 if rgb[2] else 0.1, 0.9, len(p_list))))

        for id_f, f in enumerate(p_list):
            for ind in f:
                plt.plot([ind.cur_f[comp[0]]],
                         [ind.cur_f[comp[1]]],
                         symbol,
                         color=cl[id_f], markersize=size)
            x = [ind.cur_f[comp[0]] for ind in f]
            y = [ind.cur_f[comp[1]] for ind in f]
            tmp = [(a, b) for a, b in zip(x, y)]
            tmp = sorted(tmp, key=lambda k: k[0])
            plt.step([c[0] for c in tmp], [c[1]
                     for c in tmp], color=cl[id_f], where='post')
        return plt.gca()


from __future__ import annotations

from typing import Any, Generic, Tuple, TypeVar

import jax.numpy as jnp
from jax import vjp
from jax.tree_util import tree_map

from ..annotations import PyTree
from ..dataclasses import dataclass
from ..shims import custom_vjp
from .augmented import State
from .comparing import ComparingIteratedFunction, ComparingState
from .iterated_function import (Comparand, IteratedFunction, Parameters, TheAugmentedState,
                                Trajectory)

__all__ = ['IteratedFunctionWithCombinator', 'ComparingIteratedFunctionWithCombinator']

Differentiand = TypeVar('Differentiand', bound=PyTree)

@dataclass
class _ZResiduals(Generic[Parameters, State, Comparand, Differentiand, TheAugmentedState]):
    outer_iterated_function: IteratedFunctionWithCombinator[Parameters, State, Comparand,
                                                            Differentiand, Any, TheAugmentedState]
    outer_theta: Parameters
    x_star: State

@dataclass
class _ZParameters(Generic[Parameters, State, Differentiand]):
    outer_theta: Parameters
    x_star: State
    x_star_differentiand: Differentiand
    x_star_bar_differentiand: Differentiand

def _ffp_fwd(outer_iterated_function: IteratedFunctionWithCombinator[Parameters, State, Comparand,
                                                                     Differentiand, Any,
                                                                     TheAugmentedState],
             theta: Parameters,
             initial_state: State) -> Tuple[TheAugmentedState, _ZResiduals[Parameters, State,
                                                                           Comparand, Differentiand,
                                                                           TheAugmentedState]]:
    """
    Args:
        theta: The parameters for which gradients can be calculated.
        initial_state: An initial guess of the final state.
    Returns:
        x_star: the result of the minimization.
        residuals: residuals used in _ffp_bwd.
    """
    augmented: TheAugmentedState = outer_iterated_function.find_fixed_point(
        theta, initial_state)
    return augmented, _ZResiduals(outer_iterated_function, theta, augmented.current_state)

def _ffp_bwd(residuals: _ZResiduals[Parameters, State, Comparand, Differentiand, TheAugmentedState],
             augmented_star_bar: TheAugmentedState) -> Tuple[None, Parameters, None]:
    """
    Args:
        residuals: residuals produced by _ffp_fwd.
        augmented_star_bar: cotangents
    Returns:
        theta_bar: cotangents for theta
        zeroed_xs: cotangents for initial_state
    """
    outer_iterated_function = residuals.outer_iterated_function
    outer_theta = residuals.outer_theta
    x_star = residuals.x_star
    x_star_differentiand = outer_iterated_function.extract_differentiand(outer_theta, x_star)
    x_star_bar = augmented_star_bar.current_state
    x_star_bar_differentiand = outer_iterated_function.extract_differentiand(outer_theta,
                                                                             x_star_bar)

    def f_of_theta(some_theta: Parameters) -> Differentiand:
        state = outer_iterated_function.expected_state(some_theta, x_star)
        return outer_iterated_function.extract_differentiand(outer_theta, state)

    z_iterator = _ZIterate(minimum_iterations=outer_iterated_function.z_minimum_iterations,
                           maximum_iterations=outer_iterated_function.z_maximum_iterations,
                           iterated_function=outer_iterated_function)
    z_parameters = _ZParameters(residuals.outer_theta, x_star, x_star_differentiand,
                                x_star_bar_differentiand)
    augmented = z_iterator.find_fixed_point(z_parameters, x_star_bar_differentiand)
    z_star_differentiand: State = augmented.current_state

    _, df_by_dtheta = vjp(f_of_theta, residuals.outer_theta)
    theta_bar, = df_by_dtheta(z_star_differentiand)
    return None, theta_bar, None

@dataclass
class IteratedFunctionWithCombinator(
        IteratedFunction[Parameters, State, Comparand, Trajectory, TheAugmentedState],
        Generic[Parameters, State, Comparand, Differentiand, Trajectory, TheAugmentedState]):
    """
    An IteratedFunctionWithCombinator is an IteratedFunction that invokes a combinator so that
    differentiation works through the fixed point.  Besides inheriting from this class, no other
    action is necessary to get this capability.

    It is a generic class with all of the parameters of IteratedFunction, and Differentiand, which
    is the type of the *portion of the state* with respect to which derivatives at the fixed point
    are calculated.

    Attributes:
        z_maximum_iterations:
            The maximum number of iterations to use to evaluate the adjoint's fixed point.
    """
    z_minimum_iterations: int = 11
    z_maximum_iterations: int = 1000

    # Overridden methods ---------------------------------------------------------------------------
    def find_fixed_point(self,
                         theta: Parameters,
                         initial_state: State) -> TheAugmentedState:
        """
        Args:
            theta: The parameters for which gradients can be calculated.
            initial_state: An initial guess of the final state.
        Returns: The augmented state at the fixed point.
        """
        return super().find_fixed_point(theta, initial_state)

    find_fixed_point = custom_vjp(find_fixed_point)  # type: ignore

    # Abstract methods -----------------------------------------------------------------------------
    def extract_differentiand(self, theta: Parameters, state: State) -> Differentiand:
        """
        Returns: The differentiable values in the state.  It is used by the combinator to find
            cotangents.
        """
        raise NotImplementedError

    def implant_differentiand(self,
                              theta: Parameters,
                              state: State,
                              differentiand: Differentiand) -> State:
        """
        Args:
            state: A state that will provide nondifferentiable values.
            differentiand: A differentiand that will provide differentiable values.
        Returns: A state containing differentiable from the differentiand and nondifferentiable
            values from the inputted state.
        """
        raise NotImplementedError

    # Apply vjp ------------------------------------------------------------------------------------
    find_fixed_point.defvjp(_ffp_fwd, _ffp_bwd)  # type: ignore

class ComparingIteratedFunctionWithCombinator(
        IteratedFunctionWithCombinator[Parameters, State, Comparand, Differentiand, Trajectory,
                                       ComparingState[State, Comparand]],
        ComparingIteratedFunction[Parameters, State, Comparand, Trajectory],
        Generic[Parameters, State, Comparand, Differentiand, Trajectory]):
    pass

@dataclass
class _ZIterate(ComparingIteratedFunctionWithCombinator[
        _ZParameters[Parameters, State, Differentiand],
        Differentiand,
        Differentiand,
        Differentiand,
        None],
        Generic[Parameters, State, Comparand, TheAugmentedState, Differentiand]):
    """
    The state of _ZIterate is the differentiand of the outer iterated function.
    """
    iterated_function: IteratedFunctionWithCombinator[
        Parameters, State, Comparand, Differentiand, Any, TheAugmentedState]

    # Implemented methods --------------------------------------------------------------------------
    def expected_state(self,
                       theta: _ZParameters[Parameters, State, Differentiand],
                       state: Differentiand) -> Differentiand:
        return self.sampled_state(theta, state)

    def sampled_state(self,
                      theta: _ZParameters[Parameters, State, Differentiand],
                      state: Differentiand) -> Differentiand:
        # The state should be called z, but we can't change the interface because of Liskov's
        # substitution principle.
        z = state
        del state

        def f_of_x(x_differentiand: Differentiand) -> Differentiand:
            x = self.iterated_function.implant_differentiand(theta.outer_theta, theta.x_star,
                                                             x_differentiand)
            state = self.iterated_function.expected_state(theta.outer_theta, x)
            return self.iterated_function.extract_differentiand(theta.outer_theta, state)

        _, df_by_dx = vjp(f_of_x, theta.x_star_differentiand)
        df_by_dx_times_z, = df_by_dx(z)
        return tree_map(jnp.add, theta.x_star_bar_differentiand, df_by_dx_times_z)

    def extract_comparand(self, state: Differentiand) -> Differentiand:
        return state

    def extract_differentiand(self,
                              theta: _ZParameters[Parameters, State, Differentiand],
                              state: Differentiand) -> Differentiand:
        return state

    def implant_differentiand(self,
                              theta: _ZParameters[Parameters, State, Differentiand],
                              state: Differentiand,
                              differentiand: Differentiand) -> Differentiand:
        return differentiand