__copyright__ = """
Copyright (C) 2020 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 itertools
import numpy as np
import pyopencl as cl
import loopy as lp
from pytools import memoize_method, ProcessLogger
from pytools.obj_array import make_obj_array
from boxtree.tools import DeviceDataRecord
from meshmode.array_context import PyOpenCLArrayContext
from meshmode.dof_array import unflatten, flatten, thaw
from volumential.volume_fmm import interpolate_volume_potential

import logging
logger = logging.getLogger(__name__)

__doc__ = r"""
.. currentmodule:: volumential

From :mod:`meshmode`
-------------------------

Interpolation from functions given by DoF vectors of :mod:`meshmode`.
The underlying mesh on the :mod:`meshmode` side must be discretizing the
same bounding box.

The intersection testing assumes a boundedness property for the supported
element types: each element is bounded inside the smallest :math:`l^\infty`
ball that is centered at the element's center and covers all its vertices.
This property might be broken, for example, by high order elements that
warp the element boundary too much.

.. autoclass:: ElementsToSourcesLookup

.. autoclass:: LeavesToNodesLookup

.. autofunction:: interpolate_from_meshmode

To :mod:`meshmode`
---------------------------

"""

# {{{ output

class ElementsToSourcesLookup(DeviceDataRecord):
    """
    .. attribute:: tree

        The :class:`boxtree.Tree` instance representing the box mesh.

    .. attribute:: discr

        The :class:`meshmode.discretization.Discretization` instance
        representing the external mesh and DoF distribution.

    .. attribute:: sources_in_element_starts

        Indices into :attr:`sources_in_element_lists`.

        .. code-block:: python

            sources_in_element_lists[
                sources_in_element_starts[global_iel]
                :sources_in_element_starts[global_iel] + 1
                ]

        contains the list of source nodes residing in the given element.

        .. note:: ``global_iel`` is the global element id in `meshmode`.
            ``global_iel = mesh.groups[igrp].element_nr_base + iel``.

    .. attribute:: sources_in_element_lists

        Indices into :attr:`tree.sources`.

    .. automethod:: get
    """

class LeavesToNodesLookup(DeviceDataRecord):
    """
    .. attribute:: trav

        The :class:`boxtree.FMMTraversalInfo` instance representing the
        box mesh with metadata needed for interpolation. It contains a
        reference to the underlying tree as `trav.tree`.

    .. attribute:: discr

        The :class:`meshmode.discretization.Discretization` instance
        representing the external mesh and DoF distribution.

    .. attribute:: nodes_in_leaf_starts

        Indices into :attr:`nodes_in_leaf_lists`.

        .. code-block:: python

            nodes_in_leaf_lists[
                nodes_in_leaf_starts[box_id]:nodes_in_leaf_starts[box_id] + 1]

        contains the list of discretization nodes residing in the given leaf box.

        .. note:: Only leaf boxes have non-empty entries in this table.
            Nonetheless, this list is indexed by the global box index.

    .. attribute:: nodes_in_leaf_lists

        Indices into :attr:`discr.nodes()`.

        .. note:: Unlike :class:`ElementsToSourcesLookup`, lists are not disjoint
            in the leaves-to-nodes lookup. :mod:`volumential` automatically computes
            the average contribution from overlapping boxes.

    .. automethod:: get
    """

# }}} End output

# {{{ elements-to-sources lookup builder

class ElementsToSourcesLookupBuilder:
    """Given a :mod:`meshmod` mesh and a :mod:`boxtree.Tree`, both discretizing
    the same bounding box, this class helps to build a look-up table from
    element to source nodes that are positioned inside the element.
    """

    def __init__(self, context, tree, discr):
        """
        :arg tree: a :class:`boxtree.Tree`
        :arg discr: a :class: `meshmode.discretization.Discretization`

        Boxes and elements can be non-aligned as long as the domains
        (bounding boxes) are the same.
        """
        assert tree.dimensions == discr.dim
        self.dim = discr.dim
        self.context = context
        self.tree = tree
        self.discr = discr

        from pyopencl.algorithm import KeyValueSorter
        self.key_value_sorter = KeyValueSorter(context)

        from boxtree.area_query import AreaQueryBuilder
        self.area_query_builder = AreaQueryBuilder(self.context)

    # {{{ kernel generation

    @memoize_method
    def codegen_get_dimension_specific_snippets(self):
        """Dimension-dependent code loopy instructions.
        """
        import sympy as sp
        axis_names = ["x", "y", "z"]
        axis_names = axis_names[:self.dim]

        # tolerance
        tol = -1e-12

        def make_sympy_vec(comp_names):
            comps = []
            for cn in comp_names:
                comps.append(sp.var(cn))
            return sp.Matrix(comps)

        def get_simplex_measure(vtx_names):
            mat0 = sp.ones(self.dim + 1)
            for iv, v in enumerate(vtx_names):
                vtx = sp.Matrix([sp.var(f"{v}{comp}") for comp in axis_names])
                mat0[iv, :-1] = vtx.T
            return str(mat0.det())

        if self.dim == 2:

            # {{{ 2d

            code_get_simplex = \
                """
                <> Ax = mesh_vertices_0[mesh_vertex_indices[iel, 0]]
                <> Ay = mesh_vertices_1[mesh_vertex_indices[iel, 0]]
                <> Bx = mesh_vertices_0[mesh_vertex_indices[iel, 1]]
                <> By = mesh_vertices_1[mesh_vertex_indices[iel, 1]]
                <> Cx = mesh_vertices_0[mesh_vertex_indices[iel, 2]]
                <> Cy = mesh_vertices_1[mesh_vertex_indices[iel, 2]]
                """
            code_get_point = \
                """
                <> Px = source_points_0[source_id]
                <> Py = source_points_1[source_id]
                """
            # simplex measures
            code_s0 = get_simplex_measure(["P", "B", "C"])
            code_s1 = get_simplex_measure(["A", "P", "C"])
            code_s2 = get_simplex_measure(["A", "B", "P"])
            code_compute_simplex_measures = \
                f"""
                <> s0 = {code_s0}
                <> s1 = {code_s1}
                <> s2 = {code_s2}
                """
            code_measures_have_common_sign = " and ".join([
                f"s{c1} * s{c2} >= {tol}"
                for c1, c2 in itertools.combinations(["0", "1"], 2)])

            # }}} End 2d

        elif self.dim == 3:

            # {{{ 3d

            code_get_simplex = \
                """
                <> Ax = mesh_vertices_0[mesh_vertex_indices[iel, 0]]
                <> Ay = mesh_vertices_1[mesh_vertex_indices[iel, 0]]
                <> Az = mesh_vertices_2[mesh_vertex_indices[iel, 0]]
                <> Bx = mesh_vertices_0[mesh_vertex_indices[iel, 1]]
                <> By = mesh_vertices_1[mesh_vertex_indices[iel, 1]]
                <> Bz = mesh_vertices_2[mesh_vertex_indices[iel, 1]]
                <> Cx = mesh_vertices_0[mesh_vertex_indices[iel, 2]]
                <> Cy = mesh_vertices_1[mesh_vertex_indices[iel, 2]]
                <> Cz = mesh_vertices_2[mesh_vertex_indices[iel, 2]]
                <> Dx = mesh_vertices_0[mesh_vertex_indices[iel, 3]]
                <> Dy = mesh_vertices_1[mesh_vertex_indices[iel, 3]]
                <> Dz = mesh_vertices_2[mesh_vertex_indices[iel, 3]]
                """
            code_get_point = \
                """
                <> Px = source_points_0[source_id]
                <> Py = source_points_1[source_id]
                <> Pz = source_points_2[source_id]
                """
            # simplex measures
            code_s0 = get_simplex_measure(["P", "B", "C", "D"])
            code_s1 = get_simplex_measure(["A", "P", "C", "D"])
            code_s2 = get_simplex_measure(["A", "B", "P", "D"])
            code_s3 = get_simplex_measure(["A", "B", "C", "P"])
            code_compute_simplex_measures = \
                f"""
                <> s0 = {code_s0}
                <> s1 = {code_s1}
                <> s2 = {code_s2}
                <> s3 = {code_s3}
                """
            code_measures_have_common_sign = " and ".join([
                f"s{c1} * s{c2} >= {tol}"
                for c1, c2 in itertools.combinations(["0", "1", "2"], 2)])

            # }}} End 3d

        else:
            raise NotImplementedError()

        return {"code_get_simplex": code_get_simplex,
                "code_get_point": code_get_point,
                "code_compute_simplex_measures": code_compute_simplex_measures,
                "code_measures_have_common_sign": code_measures_have_common_sign,
                }

    @memoize_method
    def get_simplex_lookup_kernel(self):
        """Returns a loopy kernel that computes a potential vector
        representing the (q_point --> element_id) relationship.
        When a source q_point lies on the element boundary, it will be
        assigned an element depending on code scheduling. This ensures
        that the resulting lookup lists are disjoint.

        The kernel assumes that the mesh uses one single group of simplex elements.
        Also, the test only works for affine elements.
        """
        logger.debug("start building elements-to-sources lookup kernel")

        snippets = self.codegen_get_dimension_specific_snippets()
        loopy_knl = lp.make_kernel(
            ["{ [ iel ]: 0 <= iel < nelements }",
             "{ [ ineighbor ]: nearby_leaves_beg <= ineighbor < nearby_leaves_end }",
             "{ [ isrc ]: 0 <= isrc < n_box_sources }"
             ],
            ["""
            for iel
                <> nearby_leaves_beg = leaves_near_ball_starts[iel]
                <> nearby_leaves_end = leaves_near_ball_starts[iel + 1]

                {code_get_simplex}

                for ineighbor
                    <> ileaf = leaves_near_ball_lists[ineighbor]
                    <> box_source_beg = box_source_starts[ileaf]
                    <> n_box_sources = box_source_counts_cumul[ileaf]

                    for isrc
                        <> source_id = box_source_beg + isrc

                        {code_get_point}
                        {code_compute_simplex_measures}

                        result[source_id] = if(
                            {code_measures_have_common_sign},
                            iel,
                            result[source_id])  {{atomic}}
                    end
                end
            end
            """.format(**snippets)],
            [lp.ValueArg("nelements, dim, nboxes, nsources", np.int32),
             lp.GlobalArg("mesh_vertex_indices", np.int32, "nelements, dim+1"),
             lp.GlobalArg("box_source_starts", np.int32, "nboxes"),
             lp.GlobalArg("box_source_counts_cumul", np.int32, "nboxes"),
             lp.GlobalArg("leaves_near_ball_lists", np.int32, None),
             lp.GlobalArg("result", np.int32, "nsources", for_atomic=True),
             "..."],
            name="build_sources_in_simplex_lookup",
            lang_version=(2018, 2),
        )

        logger.debug("done building elements-to-sources lookup kernel")
        return loopy_knl

    # }}} End kernel generation

    def compute_short_lists(self, actx, wait_for=None):
        """balls --> overlapping leaves
        """
        if not isinstance(actx, PyOpenCLArrayContext):
            if isinstance(actx, cl.CommandQueue):
                from warnings import warn
                warn("Command queue passed to the interpolator. "
                     "Supply an array context to enable proper caching.")
                actx = PyOpenCLArrayContext(actx)
            else:
                raise ValueError

        mesh = self.discr.mesh
        if len(mesh.groups) > 1:
            raise NotImplementedError("Mixed elements not supported")
        melgrp = mesh.groups[0]
        ball_centers_host = (np.max(melgrp.nodes, axis=2)
                             + np.min(melgrp.nodes, axis=2)) / 2
        ball_radii_host = np.max(
                np.max(melgrp.nodes, axis=2) - np.min(melgrp.nodes, axis=2),
                axis=0) / 2

        ball_centers = make_obj_array([
            cl.array.to_device(actx.queue, center_coord_comp)
            for center_coord_comp in ball_centers_host])
        ball_radii = cl.array.to_device(actx.queue, ball_radii_host)

        area_query_result, evt = self.area_query_builder(
            actx.queue, self.tree, ball_centers, ball_radii,
            peer_lists=None, wait_for=wait_for)
        return area_query_result, evt

    def __call__(self, actx, balls_to_leaves_lookup=None, wait_for=None):
        """
        :arg queue: a :class:`pyopencl.CommandQueue`
        """
        if not isinstance(actx, PyOpenCLArrayContext):
            if isinstance(actx, cl.CommandQueue):
                from warnings import warn
                warn("Command queue passed to the interpolator. "
                     "Supply an array context to enable proper caching.")
                actx = PyOpenCLArrayContext(actx)
            else:
                raise ValueError

        slk_plog = ProcessLogger(logger, "element-to-source lookup: run area query")

        if balls_to_leaves_lookup is None:
            balls_to_leaves_lookup, evt = \
                self.compute_short_lists(actx.queue, wait_for=wait_for)
            wait_for = [evt]

        # -----------------------------------------------------------------
        # Refine the area query using point-in-simplex test

        logger.debug("element-to-source lookup: refine starts")

        element_lookup_kernel = self.get_simplex_lookup_kernel()

        vertices_dev = make_obj_array([
            cl.array.to_device(actx.queue, verts)
            for verts in self.discr.mesh.vertices])

        mesh_vertices_kwargs = {
            f"mesh_vertices_{iaxis}": vertices_dev[iaxis]
            for iaxis in range(self.dim)}

        source_points_kwargs = {
            f"source_points_{iaxis}": self.tree.sources[iaxis]
            for iaxis in range(self.dim)}

        evt, res = element_lookup_kernel(
            actx.queue, dim=self.dim, nboxes=self.tree.nboxes,
            nelements=self.discr.mesh.nelements, nsources=self.tree.nsources,
            result=cl.array.zeros(actx.queue,
                                  self.tree.nsources, dtype=np.int32) - 1,
            mesh_vertex_indices=self.discr.mesh.groups[0].vertex_indices,
            box_source_starts=self.tree.box_source_starts,
            box_source_counts_cumul=self.tree.box_source_counts_cumul,
            leaves_near_ball_starts=balls_to_leaves_lookup.leaves_near_ball_starts,
            leaves_near_ball_lists=balls_to_leaves_lookup.leaves_near_ball_lists,
            wait_for=wait_for, **mesh_vertices_kwargs, **source_points_kwargs)

        source_to_element_lookup, = res

        wait_for = [evt]

        # elements = source_to_element_lookup.get()
        # for idx in [362,  365,  874,  877, 1386, 1389, 1898, 1901])

        # -----------------------------------------------------------------
        # Invert the source-to-element lookup by a key-value sort

        logger.debug("element-to-source lookup: key-value sort")

        sources_in_element_starts, sources_in_element_lists, evt = \
            self.key_value_sorter(
                actx.queue,
                keys=source_to_element_lookup,
                values=cl.array.arange(
                    actx.queue, self.tree.nsources, dtype=self.tree.box_id_dtype),
                nkeys=self.discr.mesh.nelements,
                starts_dtype=self.tree.box_id_dtype,
                wait_for=wait_for)

        slk_plog.done()

        return ElementsToSourcesLookup(
            tree=self.tree, discr=self.discr,
            sources_in_element_starts=sources_in_element_starts,
            sources_in_element_lists=sources_in_element_lists), evt

# }}} End elements-to-sources lookup builder

# {{{ leaves-to-nodes lookup builder

class LeavesToNodesLookupBuilder:
    """Given a :mod:`meshmod` mesh and a :mod:`boxtree.Tree`, both discretizing
    the same bounding box, this class helps to build a look-up table from
    leaf boxes to mesh nodes that are positioned inside the box.
    """

    def __init__(self, context, trav, discr):
        """
        :arg trav: a :class:`boxtree.FMMTraversalInfo`
        :arg discr: a :class: `meshmode.discretization.Discretization`

        Boxes and elements can be non-aligned as long as the domains
        (bounding boxes) are the same.
        """
        assert trav.tree.dimensions == discr.dim
        self.dim = discr.dim
        self.context = context
        self.trav = trav
        self.discr = discr

        from boxtree.area_query import LeavesToBallsLookupBuilder
        self.leaves_to_balls_lookup_builder = \
            LeavesToBallsLookupBuilder(self.context)

    def __call__(self, actx, tol=1e-12, wait_for=None):
        """
        :arg queue: a :class:`pyopencl.CommandQueue`
        :tol: nodes close enough to the boundary will be treated as
            lying on the boundary, whose interpolated values are averaged.
        """
        if not isinstance(actx, PyOpenCLArrayContext):
            if isinstance(actx, cl.CommandQueue):
                from warnings import warn
                warn("Command queue passed to the interpolator. "
                     "Supply an array context to enable proper caching.")
                actx = PyOpenCLArrayContext(actx)
            else:
                raise ValueError

        nodes = flatten(thaw(actx, self.discr.nodes()))
        radii = cl.array.zeros_like(nodes[0]) + tol

        lbl_lookup, evt = self.leaves_to_balls_lookup_builder(
            actx.queue, self.trav.tree, nodes, radii, wait_for=wait_for)

        return LeavesToNodesLookup(
            trav=self.trav, discr=self.discr,
            nodes_in_leaf_starts=lbl_lookup.balls_near_box_starts,
            nodes_in_leaf_lists=lbl_lookup.balls_near_box_lists), evt

# }}} End leaves-to-nodes lookup builder

# {{{ transform helper

def compute_affine_transform(source_simplex, target_simplex):
    """Computes A and b for the affine transform :math:`y = A x + b`
    that maps ``source_simplex`` to ``target_simplex``.

    :param source_simplex: a dim-by-(dim+1) :mod:`numpy` array
    :param target_simplex: a dim-by-(dim+1) :mod:`numpy: array
    """
    assert source_simplex.shape == target_simplex.shape
    dim = source_simplex.shape[0]

    if dim == 2:
        assert source_simplex.shape == (2, 3)
        mat = np.zeros([6, 6])
        mat[:3, :2] = source_simplex.T
        mat[-3:, 2:4] = source_simplex.T
        mat[:3, -2] = 1
        mat[-3:, -1] = 1
        rhs = target_simplex.reshape(-1)
        solu = np.linalg.solve(mat, rhs)
        return solu[:4].reshape(2, 2), solu[-2:]

    elif dim == 3:
        assert source_simplex.shape == (3, 4)
        mat = np.zeros([12, 12])
        mat[:4, :3] = source_simplex.T
        mat[4:8, 3:6] = source_simplex.T
        mat[-4:, 6:9] = source_simplex.T
        mat[:4, -3] = 1
        mat[4:8, -2] = 1
        mat[-4:, -1] = 1
        rhs = target_simplex.reshape(-1)
        solu = np.linalg.solve(mat, rhs)
        return solu[:9].reshape(3, 3), solu[-3:]

    else:
        raise NotImplementedError()

def invert_affine_transform(mat_a, disp_b):
    """Inverts an affine transform given by :math:`y = A x + b`.

    :param mat_A: a dim*(dim+1)-by-dim*(dim+1) :mod:`numpy` array
    :param disp_b: a dim*(dim+1) :mod:`numpy` array
    """
    iva = np.linalg.inv(mat_a)
    ivb = - iva @ disp_b
    return iva, ivb

# }}} End transform helper

# {{{ from meshmode interpolation

def interpolate_from_meshmode(actx, dof_vec, elements_to_sources_lookup,
                              order="tree"):
    """Interpolate a DoF vector from :mod:`meshmode`.

    :arg dof_vec: a DoF vector representing a field in :mod:`meshmode`
        of shape ``(..., nnodes)``.
    :arg elements_to_sources_lookup: a :class:`ElementsToSourcesLookup`.
    :arg order: order of the output potential, either "tree" or "user".

    .. note:: This function currently supports meshes with just one element
        group. Also, the element group must be simplex-based.

    .. note:: This function does some heavy-lifting computation in Python,
        which we intend to optimize in the future. In particular, we plan
        to shift the batched linear solves and basis evaluations to
        :mod:`loopy`.

    TODO: make linear solvers available as :mod:`loopy` callables.
    TODO: make :mod:`modepy` emit :mod:`loopy` callables for basis evaluation.
    """
    if not isinstance(dof_vec, cl.array.Array):
        raise TypeError("non-array passed to interpolator")

    if not isinstance(actx, PyOpenCLArrayContext):
        if isinstance(actx, cl.CommandQueue):
            from warnings import warn
            warn("Command queue passed to the interpolator. "
                 "Supply an array context to enable proper caching.")
            actx = PyOpenCLArrayContext(actx)
        else:
            raise ValueError

    assert len(elements_to_sources_lookup.discr.groups) == 1
    assert len(elements_to_sources_lookup.discr.mesh.groups) == 1
    degroup = elements_to_sources_lookup.discr.groups[0]
    megroup = elements_to_sources_lookup.discr.mesh.groups[0]

    if not degroup.is_affine:
        raise ValueError(
            "interpolation requires global-to-local map, "
            "which is only available for affinely mapped elements")

    mesh = elements_to_sources_lookup.discr.mesh
    dim = elements_to_sources_lookup.discr.dim
    template_simplex = mesh.groups[0].vertex_unit_coordinates().T

    # -------------------------------------------------------
    # Inversely map source points with a global-to-local map.
    #
    # 1. For each element, solve for the affine map.
    #
    # 2. Apply the map to corresponding source points.
    #
    # This step computes `unit_sources`, the list of inversely
    # mapped source points.

    sources_in_element_starts = \
        elements_to_sources_lookup.sources_in_element_starts.get(actx.queue)
    sources_in_element_lists = \
        elements_to_sources_lookup.sources_in_element_lists.get(actx.queue)
    tree = elements_to_sources_lookup.tree.get(actx.queue)

    unit_sources_host = make_obj_array(
            [np.zeros_like(srccrd) for srccrd in tree.sources])

    for iel in range(degroup.nelements):
        vertex_ids = megroup.vertex_indices[iel]
        vertices = mesh.vertices[:, vertex_ids]
        afa, afb = compute_affine_transform(vertices, template_simplex)

        beg = sources_in_element_starts[iel]
        end = sources_in_element_starts[iel + 1]
        source_ids_in_el = sources_in_element_lists[beg:end]
        sources_in_el = np.vstack(
            [tree.sources[iaxis][source_ids_in_el] for iaxis in range(dim)])

        ivmapped_el_sources = afa @ sources_in_el + afb.reshape([dim, 1])
        for iaxis in range(dim):
            unit_sources_host[iaxis][source_ids_in_el] = \
                ivmapped_el_sources[iaxis, :]

    unit_sources = make_obj_array(
        [cl.array.to_device(actx.queue, usc) for usc in unit_sources_host])

    # -----------------------------------------------------
    # Carry out evaluations in the local (template) frames.
    #
    # 1. Assemble a resampling matrix for each element, with
    #    the basis functions and the local source points.
    #
    # 2. For each element, perform matvec on the resampling
    #    matrix and the local DoF coefficients.
    #
    # This step assumes `unit_sources` computed on device, so
    # that the previous step can be swapped with a kernel without
    # interrupting the followed computation.

    mapped_sources = np.vstack(
        [usc.get(actx.queue) for usc in unit_sources])

    basis_funcs = degroup.basis()

    dof_vec_view = unflatten(
            actx, elements_to_sources_lookup.discr, dof_vec)[0]
    dof_vec_view = dof_vec_view.get()

    sym_shape = dof_vec.shape[:-1]
    source_vec = np.zeros(sym_shape + (tree.nsources, ))

    for iel in range(degroup.nelements):
        beg = sources_in_element_starts[iel]
        end = sources_in_element_starts[iel + 1]
        source_ids_in_el = sources_in_element_lists[beg:end]
        mapped_sources_in_el = mapped_sources[:, source_ids_in_el]
        local_dof_vec = dof_vec_view[..., iel, :]

        # resampling matrix built from Vandermonde matrices
        import modepy as mp
        rsplm = mp.resampling_matrix(
                basis=basis_funcs,
                new_nodes=mapped_sources_in_el,
                old_nodes=degroup.unit_nodes)

        if len(sym_shape) == 0:
            local_coeffs = local_dof_vec
            source_vec[source_ids_in_el] = rsplm @ local_coeffs
        else:
            from pytools import indices_in_shape
            for sym_id in indices_in_shape(sym_shape):
                source_vec[sym_id + (source_ids_in_el, )] = \
                    rsplm @ local_dof_vec[sym_id]

    source_vec = cl.array.to_device(actx.queue, source_vec)

    if order == "tree":
        pass  # no need to do anything
    elif order == "user":
        source_vec = source_vec[tree.sorted_target_ids]  # into user order
    else:
        raise ValueError(f"order must be 'tree' or 'user' (got {order}).")

    return source_vec

# }}} End from meshmode interpolation

# {{{ to meshmode interpolation

def interpolate_to_meshmode(actx, potential, leaves_to_nodes_lookup,
                            order="tree"):
    """
    :arg potential: a DoF vector representing a field in :mod:`volumential`,
        in tree order.
    :arg leaves_to_nodes_lookup: a :class:`LeavesToNodesLookup`.
    :arg order: order of the input potential, either "tree" or "user".

    :returns: a :class:`pyopencl.Array` of shape (nnodes, 1) containing the
        interpolated data.
    """
    if order == "tree":
        potential_in_tree_order = True
    elif order == "user":
        potential_in_tree_order = False
    else:
        raise ValueError(f"order must be 'tree' or 'user' (got {order}).")

    if not isinstance(actx, PyOpenCLArrayContext):
        if isinstance(actx, cl.CommandQueue):
            from warnings import warn
            warn("Command queue passed to the interpolator. "
                 "Supply an array context to enable proper caching.")
            actx = PyOpenCLArrayContext(actx)
        else:
            raise ValueError

    target_points = flatten(thaw(actx, leaves_to_nodes_lookup.discr.nodes()))

    traversal = leaves_to_nodes_lookup.trav
    tree = leaves_to_nodes_lookup.trav.tree

    dim = tree.dimensions

    # infer q_order from tree
    pts_per_box = tree.ntargets // traversal.ntarget_boxes
    assert pts_per_box * traversal.ntarget_boxes == tree.ntargets

    # allow for +/- 0.25 floating point error
    q_order = int(pts_per_box**(1 / dim) + 0.25)
    assert q_order**dim == pts_per_box

    interp_p = interpolate_volume_potential(
            target_points=target_points, traversal=traversal,
            wrangler=None, potential=potential,
            potential_in_tree_order=potential_in_tree_order,
            dim=dim, tree=tree, queue=actx.queue, q_order=q_order,
            dtype=potential.dtype, lbl_lookup=None,
            balls_near_box_starts=leaves_to_nodes_lookup.nodes_in_leaf_starts,
            balls_near_box_lists=leaves_to_nodes_lookup.nodes_in_leaf_lists)

    return interp_p

# }}} End to meshmode interpolation


#! /usr/bin/env python

"LiDAR ground algorithm"

import os, glob, argparse, time, math
import numpy as np
from laspy.file import File
import sqlite3
import ConfigParser

def getArgs():
	parser = argparse.ArgumentParser(
		description="LiDAR ground and canopy classification."
	)
	parser.add_argument(
		"-c",
		"--config",
		type=str,
		required=True,
		help="Configuration File."
	)
	parser.add_argument(
		"-g",
		"--ground",
		action = "store_true",
		help="Ground classification."

	)
	parser.add_argument(
		"-t",
		"--toc",
		action = "store_true",
		help="Canopy classification."

	)
	parser.add_argument(
		"-v",
		"--verbose",
		action = "store_true",
		help="Print status to screen."

	)

	return parser.parse_args()

def getConfigs(configFile):
	Configs={}
	try:
		config = ConfigParser.ConfigParser()
		config.read(configFile)
		Configs["paths"]=dict(config.items("paths"))
		Configs["vars"]=dict(config.items("vars"))
	except Exception as e:
		print "Problem parsing configuration file: {}.".format(configFile)
		raise e
	return Configs

class Points(object):
	def __init__(self, points, configs, filename):
		self.points = points
		self.configs = configs
		self.filename = filename

	def getParams(self):
		"""
		"""
		indirs = self.configs["paths"]["input"]
		indir = indirs.split(":")[0]
		os.chdir(indir)
		infile=glob.glob("*.las")[0]
		data = File(infile, mode="r")
		self.header = data.header
		#print self.header
		self.dtype = data.points.dtype

	def write(self):
		output = self.configs["paths"]["output"]
		points = np.array(self.points, dtype=self.dtype)
		#print self.header
		outfile = File(output + self.filename, mode="w", header=self.header)
		outfile.points = points
		outfile.close()

def setGround(pt):
	pt = list(pt)
	pt[5]=2
	return tuple(pt)

def setCanopy(pt):
	pt = list(pt)
	pt[5]=1
	return tuple(pt)

def height(dz, thresh=1.5):
	return dz < 1.5

def angle(dX, dY, dZ, thresh=5.5, eps=0.000000001):
	hyp = math.sqrt((dX)**2+(dY)**2+(dZ)**2)
	# Use epsilon equality for floats
	if hyp < eps:
		return False
	degrees = math.asin(dZ/hyp)*(180/math.pi)
	return degrees < thresh

def groundFilter(seed, pt):
	X0=seed[0]*0.001
	Y0=seed[1]*0.001
	Z0=seed[2]*0.001
	X1=pt[0]*0.001
	Y1=pt[1]*0.001
	Z1=pt[2]*0.001
	dX = X1-X0
	dY = Y1-Y0
	dZ = Z1-Z0
	# skip duplicated point
	if seed == pt:
		return False
	return height(dZ) and angle(dX, dY, dZ)

def groundClassifier(configs):
	"""
	"""
	t_i = time.time()
	points = []

	database = configs["paths"]["db"]
	conn = sqlite3.connect(database)
	c = conn.cursor()
	minZs = c.execute("""SELECT distinct(hash10), min(Z) FROM pointcloud
							WHERE num_returns=return_number
							GROUP BY hash10;"""
	).fetchall()
	minZs = list(minZs)
	print minZs
	count = 0
	for h, z in minZs:
		seed = c.execute("""
			SELECT
				X,
				Y,
				Z,
				intensity,
				flag_byte,
				raw_classification,
				scan_angle_rank,
				user_data,
				pt_src_id,
				gps_time
			FROM pointcloud
			WHERE Z=? and hash10=?;
			""", (z, h)).fetchall()[0]
		print seed
		seed = setGround(seed)
		points.append((seed,))
		putative_grounds = c.execute("""
			SELECT
				X,
				Y,
				Z,
				intensity,
				flag_byte,
				raw_classification,
				scan_angle_rank,
				user_data,
				pt_src_id,
				gps_time
			FROM pointcloud
			WHERE hash10=? AND num_returns=return_number;
			""", (h,)).fetchall()
		for putative_ground in putative_grounds:
			if groundFilter(seed, putative_ground):
				putative_ground=setGround(putative_ground)
				points.append((putative_ground,))
		count += 1
		print "Finished {}; block hash {}.".format(count, h)
	conn.close()
	t_f = time.time()
	print "Ground classifier took {} minutes.".format((t_f-t_i)/60.0)
	return points

def canopyClassifier(configs):
	"""
	"""
	t_i = time.time()
	points = []
	scale = float(configs["vars"]["scale"])
	height_max = float(configs["vars"]["height_max"])/scale # depends on scale
	database = configs["paths"]["db"]
	conn = sqlite3.connect(database)
	c = conn.cursor()

	hashmap_list = c.execute("SELECT distinct(hash1),hash10 FROM pointcloud;").fetchall()
	hashmap = {}
	for k,v in hashmap_list:
		hashmap[k]=v

	minZ_list = c.execute("""SELECT distinct(hash10), min(Z) FROM pointcloud
							GROUP BY hash10"""
	).fetchall()
	min_z = {}
	for k,v in minZ_list:
		min_z[k]=v

	maxZs = c.execute("""SELECT distinct(hash1), max(Z) FROM pointcloud
							WHERE return_number=1
							GROUP BY hash1"""
	).fetchall()
	count = 0
	for h1,h10 in hashmap.iteritems():
		z_max = min_z[h10] + height_max
		results = c.execute("""
			SELECT
				X,
				Y,
				Z,
				intensity,
				flag_byte,
				raw_classification,
				scan_angle_rank,
				user_data,
				pt_src_id,
				gps_time
			FROM pointcloud
			WHERE Z<? and hash1=?;
			""", (z_max, h1)).fetchall()
		heights={}
		for result in results:
			Z = result[2]
			heights[Z]=result
		if len(heights)>0:
			TOC = heights[max(heights)]
			TOC = setCanopy(TOC)
			points.append((TOC,))
			count += 1
			print "Finished {}; block hash {}.".format(count, h1)
	conn.close()
	t_f = time.time()
	print "Canopy classifier took {} minutes.".format((t_f-t_i)/60.0)
	return points

def main():
	t_i = time.time()
	args=getArgs()
	if args.verbose:
		print args
	#base = os.getcwd()
	configs = getConfigs(args.config)
	if args.verbose:
		print configs
	if args.ground:
		points = groundClassifier(configs)
		print points
		groundPoints = Points(points, configs, "ground.las")
		groundPoints.getParams()
		groundPoints.write()
		del points
	if args.toc:
		points = canopyClassifier(configs)
		print points
		canopyPoints = Points(points, configs, "canopy.las")
		canopyPoints.getParams()
		canopyPoints.write()
		del points
	t_f = time.time()
	if args.verbose:
		print "Total elapsed time {} minutes.".format((t_f-t_i)/60)

			#list_data.append((window_index, data_text))
		#c.executemany('INSERT INTO pointcloud VALUES (?,?)', list_data)

	#data = group(args.input)
	#os.chdir(base)
	#print data

if __name__ == "__main__":
	main()


# Copyright (c) 2019 - The Procedural Generation for Gazebo authors
# For information on the respective copyright owner see the NOTICE file
#
# 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.

from ...parsers.sdf import create_sdf_element
import collections
import numpy as np
from copy import deepcopy
from .pose import Pose

class Inertial(object):
    def __init__(self, mass=0, ixx=0, iyy=0, izz=0, ixy=0, ixz=0, iyz=0):
        self._mass = mass
        self._pose = Pose()
        self._ixx = ixx
        self._iyy = iyy
        self._izz = izz
        self._ixy = ixy
        self._iyz = iyz
        self._ixz = ixz

    def __str__(self):
        pose = self._pose.position + self._pose.rpy
        msg = 'Mass [Kg]={}\n'.format(self._mass)
        msg += 'Pose={}\n'.format(pose)
        msg += 'I =\n'
        msg += '\tIxx={}\n'.format(self._ixx)
        msg += '\tIyy={}\n'.format(self._iyy)
        msg += '\tIzz={}\n'.format(self._izz)
        msg += '\tIxy={}\n'.format(self._ixy)
        msg += '\tIxz={}\n'.format(self._ixz)
        msg += '\tIyz={}\n'.format(self._iyz)
        return msg

    @property
    def mass(self):
        return self._mass

    @mass.setter
    def mass(self, value):
        assert value > 0, 'Mass must be greater than zero'
        self._mass = value

    @property
    def pose(self):
        return self._pose

    @pose.setter
    def pose(self, vec):
        if isinstance(vec, Pose):
            self._pose = vec
        else:
            assert isinstance(vec, collections.Iterable), \
                'Input pose vector must be iterable'
            assert len(vec) == 6 or len(vec) == 7, \
                'Pose must be given as position and Euler angles (x, y, z, ' \
                'roll, pitch, yaw) or position and quaternions (x, y, z, ' \
                'qx, qy, qz, qw)'
            for item in vec:
                assert isinstance(item, float) or isinstance(item, int), \
                    'All elements in pose vector must be a float or an integer'

            self._pose = Pose(pos=vec[0:3], rot=vec[3::])

    @property
    def ixx(self):
        return self._ixx

    @ixx.setter
    def ixx(self, value):
        assert isinstance(value, float) or isinstance(value, int), \
            'Input value must be a float or an integer, provided={}'.format(type(value))
        self._ixx = value

    @property
    def iyy(self):
        return self._iyy

    @iyy.setter
    def iyy(self, value):
        assert isinstance(value, float) or isinstance(value, int), \
            'Input value must be a float or an integer, provided={}'.format(type(value))
        self._iyy= value

    @property
    def izz(self):
        return self._izz

    @izz.setter
    def izz(self, value):
        assert isinstance(value, float) or isinstance(value, int), \
            'Input value must be a float or an integer, provided={}'.format(type(value))
        self._izz = value

    @property
    def ixy(self):
        return self._ixy

    @ixy.setter
    def ixy(self, value):
        assert isinstance(value, float) or isinstance(value, int), \
            'Input value must be a float or an integer, provided={}'.format(type(value))
        self._ixy = value

    @property
    def ixz(self):
        return self._ixz

    @ixz.setter
    def ixz(self, value):
        assert isinstance(value, float) or isinstance(value, int), \
            'Input value must be a float or an integer, provided={}'.format(type(value))
        self._ixz = value

    @property
    def iyz(self):
        return self._iyz

    @iyz.setter
    def iyz(self, value):
        assert isinstance(value, float) or isinstance(value, int), \
            'Input value must be a float or an integer, provided={}'.format(type(value))
        self._iyz = value

    @property
    def moi(self):
        return np.array([
            [self.ixx, self.ixy, self.ixz],
            [-self.ixy, self.iyy, self.iyz],
            [-self.ixz, -self.iyz, self.izz]])

    @staticmethod
    def create_inertia(inertia_type, **kwargs):
        if inertia_type == 'solid_sphere':
            return Inertial.create_solid_sphere_inertia(**kwargs)
        elif inertia_type == 'hollow_sphere':
            return Inertial.create_hollow_sphere_inertia(**kwargs)
        elif inertia_type == 'ellipsoid':
            return Inertial.create_ellipsoid_inertia(**kwargs)
        elif inertia_type == 'cuboid':
            return Inertial.create_cuboid_inertia(**kwargs)
        elif inertia_type == 'centered_rod':
            return Inertial.create_centered_rod_inertia(**kwargs)
        elif inertia_type == 'solid_cylinder':
            return Inertial.create_solid_cylinder_inertia(**kwargs)
        elif inertia_type == 'custom':
            return Inertial(**kwargs)
        else:
            return None

    @staticmethod
    def create_solid_sphere_inertia(mass, radius):
        assert mass > 0, 'Mass must be greater than zero'
        assert radius > 0, 'Radius must be greater than zero'
        inertia = Inertial()

        fac = 2. / 5
        inertia.mass = mass
        inertia.ixx = fac * mass * radius**2
        inertia.iyy = fac * mass * radius**2
        inertia.izz = fac * mass * radius**2

        return inertia

    @staticmethod
    def create_hollow_sphere_inertia(mass, radius):
        assert mass > 0, 'Mass must be greater than zero'
        assert radius > 0, 'Radius must be greater than zero'
        inertia = Inertial()

        fac = 2. / 3
        inertia.mass = mass
        inertia.ixx = fac * mass * radius**2
        inertia.iyy = fac * mass * radius**2
        inertia.izz = fac * mass * radius**2

        return inertia

    @staticmethod
    def create_ellipsoid_inertia(mass, axis_length_x, axis_length_y, axis_length_z):
        assert mass > 0
        assert axis_length_x > 0
        assert axis_length_y > 0
        assert axis_length_z > 0
        inertia = Inertial()

        fac = 1. / 5

        inertia.mass = mass
        inertia.ixx = fac * mass * (axis_length_y**2 + axis_length_z**2)
        inertia.iyy = fac * mass * (axis_length_x**2 + axis_length_z**2)
        inertia.izz = fac * mass * (axis_length_x**2 + axis_length_y**2)

        return inertia

    @staticmethod
    def create_cuboid_inertia(mass, length_x, length_y, length_z):
        assert mass > 0
        assert length_x > 0
        assert length_y > 0
        assert length_z > 0
        inertia = Inertial()

        fac = 1. / 12

        inertia.mass = mass
        inertia.ixx = fac * mass * (length_y**2 + length_z**2)
        inertia.iyy = fac * mass * (length_x**2 + length_z**2)
        inertia.izz = fac * mass * (length_x**2 + length_y**2)

        return inertia

    @staticmethod
    def create_centered_rod_inertia(mass, length, axis):
        assert mass > 0
        assert length > 0
        assert isinstance(axis, list)
        assert len(axis) == 3
        assert sum(axis) == 1
        assert axis[0] == 1 or axis[1] == 1 or axis[2] == 1
        inertia = Inertial()

        fac = 1. / 12
        inertia.mass = mass
        inertia.ixx = axis[0] * fac * mass * length**2
        inertia.iyy = axis[1] * fac * mass * length**2
        inertia.izz = axis[2] * fac * mass * length**2

        return inertia

    @staticmethod
    def create_solid_cylinder_inertia(mass, radius, length, axis):
        assert mass > 0, 'Mass must be greater than zero'
        assert length > 0, 'Length must be greater than zero'
        assert radius > 0, 'Radius must be greater than zero'
        assert isinstance(axis, list), 'Axis vector must be provided as a list'
        assert len(axis) == 3, 'Axis vector must have three elements'
        assert sum(axis) == 1, 'Axis vector must be a unit vector'
        assert axis[0] == 1 or axis[1] == 1 or axis[2] == 1
        inertia = Inertial()

        fac = 1. / 12
        inertia.mass = mass
        i_axis = 0.5 * mass * radius**2
        i_side = 1. / 12 * mass * (3 * radius**2 + length**2)

        inertia.ixx = i_side
        inertia.iyy = i_side
        inertia.izz = i_side

        if axis[0] == 1:
            inertia.ixx = i_axis
        elif axis[1] == 1:
            inertia.iyy = i_axis
        else:
            inertia.izz = i_axis

        return inertia

    def to_sdf(self):
        sdf = create_sdf_element('inertial')
        sdf.mass = self._mass
        sdf.pose = self._pose.to_sdf()
        sdf.inertia.ixx = self._ixx
        sdf.inertia.iyy = self._iyy
        sdf.inertia.izz = self._izz
        sdf.inertia.ixy = self._ixy
        sdf.inertia.ixz = self._ixz
        sdf.inertia.iyz = self._iyz
        return sdf

    @staticmethod
    def from_sdf(sdf):
        assert sdf._NAME == 'inertial', 'Input SDF element must be of type inertial'
        inertial = Inertial()
        inertial.mass = sdf.mass.value
        inertial._pose = Pose.from_sdf(sdf.pose)
        inertial.ixx = sdf.inertia.ixx.value
        inertial.iyy = sdf.inertia.iyy.value
        inertial.izz = sdf.inertia.izz.value
        inertial.ixy = sdf.inertia.ixy.value
        inertial.ixz = sdf.inertia.ixz.value
        inertial.iyz = sdf.inertia.iyz.value
        return inertial


"""
Module for snow masking

Credits:
Copyright (c) 2017-2022 Matej Aleksandrov, Matej Batič, Grega Milčinski, Domagoj Korais, Matic Lubej (Sinergise)
Copyright (c) 2017-2022 Žiga Lukšič, Devis Peressutti, Nejc Vesel, Jovan Višnjić, Anže Zupanc (Sinergise)
Copyright (c) 2019-2020 Jernej Puc, Lojze Žust (Sinergise)
Copyright (c) 2017-2019 Blaž Sovdat, Andrej Burja, Eva Erzin (Sinergise)

This source code is licensed under the MIT license found in the LICENSE
file in the root directory of this source tree.
"""

import logging
import itertools

import numpy as np
from skimage.morphology import disk, binary_dilation

from eolearn.core import EOTask, FeatureType
from .utils import resize_images

LOGGER = logging.getLogger(__name__)

class BaseSnowMaskTask(EOTask):
    """Base class for snow detection and masking"""

    def __init__(self, data_feature, band_indices, dilation_size=0, undefined_value=0, mask_name="SNOW_MASK"):
        """
        :param data_feature: EOPatch feature represented by a tuple in the form of `(FeatureType, 'feature_name')`
        :type data_feature: tuple(FeatureType, str)
        :param band_indices: A list containing the indices at which the required bands can be found in the data_feature.
        :type band_indices: list(int)
        :param dilation_size: Size of the disk in pixels for performing dilation. Value 0 means do not perform
            this post-processing step.
        :type dilation_size: int
        """
        self.bands_feature = self.parse_feature(data_feature)
        self.band_indices = band_indices
        self.dilation_size = dilation_size
        self.undefined_value = undefined_value
        self.mask_feature = (FeatureType.MASK, mask_name)

    def _apply_dilation(self, snow_masks):
        """Apply binary dilation for each mask in the series"""
        if self.dilation_size:
            snow_masks = np.array([binary_dilation(mask, disk(self.dilation_size)) for mask in snow_masks])
        return snow_masks

    def execute(self, eopatch):
        raise NotImplementedError

class SnowMaskTask(BaseSnowMaskTask):
    """The task calculates the snow mask using the given thresholds.

    The default values were optimised based on the Sentinel-2 L1C processing level. Values might not be optimal for L2A
    processing level
    """

    NDVI_THRESHOLD = 0.1

    def __init__(self, data_feature, band_indices, ndsi_threshold=0.4, brightness_threshold=0.3, **kwargs):
        """
        :param data_feature: EOPatch feature represented by a tuple in the form of `(FeatureType, 'feature_name')`
            containing the bands 2, 3, 7, 11, i.e. (FeatureType.DATA, 'BANDS')
        :type data_feature: tuple(FeatureType, str)
        :param band_indices: A list containing the indices at which the required bands can be found in the data_feature.
            The required bands are B03, B04, B08 and B11 and the indices should be provided in this order. If the
            'BANDS' array contains all 13 L1C bands, then `band_indices=[2, 3, 7, 11]`. If the 'BANDS' are the 12 bands
            with L2A values, then `band_indices=[2, 3, 7, 10]`
        :type band_indices: list(int)
        :param ndsi_threshold: Minimum value of the NDSI required to classify the pixel as snow
        :type ndsi_threshold: float
        :param brightness_threshold: Minimum value of the red band for a pixel to be classified as bright
        :type brightness_threshold: float
        """
        super().__init__(data_feature, band_indices, **kwargs)
        self.ndsi_threshold = ndsi_threshold
        self.brightness_threshold = brightness_threshold

    def execute(self, eopatch):
        bands = eopatch[self.bands_feature][..., self.band_indices]
        with np.errstate(divide="ignore"):
            # (B03 - B11) / (B03 + B11)
            ndsi = (bands[..., 0] - bands[..., 3]) / (bands[..., 0] + bands[..., 3])
            # (B08 - B04) / (B08 + B04)
            ndvi = (bands[..., 2] - bands[..., 1]) / (bands[..., 2] + bands[..., 1])

        ndsi_invalid, ndvi_invalid = ~np.isfinite(ndsi), ~np.isfinite(ndvi)
        ndsi[ndsi_invalid] = self.undefined_value
        ndvi[ndvi_invalid] = self.undefined_value

        snow_mask = np.where(
            np.logical_and(
                np.logical_or(
                    ndsi >= self.ndsi_threshold, np.abs(ndvi - self.NDVI_THRESHOLD) < self.NDVI_THRESHOLD / 2
                ),
                bands[..., 0] >= self.brightness_threshold,
            ),
            1,
            0,
        )

        snow_mask = self._apply_dilation(snow_mask)

        snow_mask[np.logical_or(ndsi_invalid, ndvi_invalid)] = self.undefined_value

        eopatch[self.mask_feature] = snow_mask[..., np.newaxis].astype(bool)
        return eopatch

class TheiaSnowMaskTask(BaseSnowMaskTask):
    """Task to add a snow mask to an EOPatch. The input data is either Sentinel-2 L1C or L2A level

    Original implementation and documentation available at https://gitlab.orfeo-toolbox.org/remote_modules/let-it-snow

    ATBD https://gitlab.orfeo-toolbox.org/remote_modules/let-it-snow/blob/master/doc/atbd/ATBD_CES-Neige.pdf

    This task computes a snow mask for the input EOPatch. The `data_feature` to be used as input to the
    classifier is a mandatory argument. If all of the needed features exist already, the classifier is run.
    `linear` interpolation is used for resampling of the `data_feature` and cloud probability map, while `nearest`
    interpolation is used to upsample the binary cloud mask.
    """

    B10_THR = 0.015
    DEM_FACTOR = 0.00001

    def __init__(
        self,
        data_feature,
        band_indices,
        cloud_mask_feature,
        dem_feature,
        dem_params=(100, 0.1),
        red_params=(12, 0.3, 0.1, 0.2, 0.040),
        ndsi_params=(0.4, 0.15, 0.001),
        b10_index=None,
        **kwargs,
    ):
        """
        :param data_feature: EOPatch feature represented by a tuple in the form of `(FeatureType, 'feature_name')`
            containing the bands B3, B4, and B11

            Example: `(FeatureType.DATA, 'ALL-BANDS')`
        :type data_feature: tuple(FeatureType, str)
        :param band_indices: A list containing the indices at which the required bands can be found in the bands
            feature. If all L1C band values are provided, `band_indices=[2, 3, 11]`. If all L2A band values are
            provided, then `band_indices=[2, 3, 10]`
        :type band_indices: list(int)
        :param cloud_mask_feature: EOPatch CLM feature represented by a tuple in the form of
            `(FeatureType, 'feature_name')` containing the cloud mask
        :type cloud_mask_feature: tuple(FeatureType, str)
        :param dem_feature: EOPatch DEM feature represented by a tuple in the form of `(FeatureType, 'feature_name')`
            containing the digital elevation model
        :type dem_feature: tuple(FeatureType, str)
        :param b10_index: Array index where the B10 band is stored in the bands feature. This is used to refine the
            initial cloud mask
        :type b10_index: int
        :param dem_params: Tuple with parameters pertaining DEM processing. The first value specifies the bin size
            used to group DEM values, while the second value specifies the minimum snow fraction in an elevation band
            to define z_s. With reference to the ATBD, the tuple is (d_z, f_t)
        :type dem_params: (float, float)
        :param red_params: Tuple specifying parameters to process the B04 red band. The first parameter defines the
            scaling factor for down-sampling the red band, the second parameter is the maximum value of the
            down-sampled red band for a dark cloud pixel, the third parameter is the minimum value
            to return a non-snow pixel to the cloud mask, the fourth is the minimum reflectance value to pass the 1st
            snow test, and the fifth is the minimum reflectance value to pass the 2nd snow test. With reference to the
            ATBD, the tuple is (r_f, r_d, r_b, r_1, r_2)
        :type red_params: (float, float, float, float, float)
        :param ndsi_params: Tuple specifying parameters for the NDSI. First parameter is the minimum value to pass the
            1st snow test, the second parameter is the minimum value to pass the 2nd snow test, and the third parameter
            is the minimum snow fraction in the image to activate the pass 2 snow test. With reference to the
            ATBD, the tuple is (n_1, n_2, f_s)
        :type ndsi_params: (float, float, float)
        """
        super().__init__(data_feature, band_indices, **kwargs)
        self.dem_feature = self.parse_feature(dem_feature)
        self.clm_feature = self.parse_feature(cloud_mask_feature)
        self.dem_params = dem_params
        self.red_params = red_params
        self.ndsi_params = ndsi_params
        self.b10_index = b10_index
        self._validate_params()

    def _validate_params(self):
        """Check length of parameters defining threshold values"""
        for params, n_params in [(self.dem_params, 2), (self.red_params, 5), (self.ndsi_params, 3)]:
            if not isinstance(params, (tuple, list)) or len(params) != n_params:
                raise ValueError(
                    f"Incorrect format or number of parameters for {params}. Has to be a tuple of length {n_params}"
                )

    def _resample_red(self, input_array):
        """Method to resample the values of the red band

        The input array is first down-scaled using bicubic interpolation and up-scaled back using nearest neighbour
        interpolation

        :param input_array: input values
        :return: resampled values
        """
        height, width = input_array.shape[1:]
        size = (height // self.red_params[0], width // self.red_params[0])
        return resize_images(
            resize_images(input_array[..., np.newaxis], new_size=size), new_size=(height, width)
        ).squeeze()

    def _adjust_cloud_mask(self, bands, cloud_mask, dem, b10):
        """Adjust existing cloud mask using cirrus band if L1C data and resampled red band

        Add to the existing cloud mask pixels found thresholding down-sampled red band and cirrus band/DEM
        """
        clm_b10 = (
            np.where(b10 > self.B10_THR + self.DEM_FACTOR * dem, 1, 0)
            if b10 is not None
            else np.ones(shape=cloud_mask.shape, dtype=np.uint8)
        )
        return np.logical_or(
            np.where(np.logical_and(cloud_mask == 1, self._resample_red(bands[..., 1]) > self.red_params[1]), 1, 0),
            clm_b10,
        ).astype(np.uint8)

    def _apply_first_pass(self, bands, ndsi, clm, dem, clm_temp):
        """Apply first pass of snow detection"""
        snow_mask_pass1 = np.where(
            np.logical_and(
                np.logical_not(clm_temp), np.logical_and(ndsi > self.ndsi_params[0], bands[..., 1] > self.red_params[3])
            ),
            1,
            0,
        )

        clm_pass1 = np.where(
            np.logical_or(clm_temp, (bands[..., 1] > self.red_params[2]) & np.logical_not(snow_mask_pass1) & clm), 1, 0
        )

        dem_edges = np.linspace(
            np.min(dem), np.max(dem), int(np.ceil((np.max(dem) - np.min(dem)) / self.dem_params[0]))
        )
        nbins = len(dem_edges) - 1
        dem_hist_clear_pixels, snow_frac = None, None
        if nbins > 0:
            snow_frac = np.zeros(shape=(bands.shape[0], nbins))
            dem_hist_clear_pixels = np.array(
                [np.histogram(dem[np.logical_not(mask)], bins=dem_edges)[0] for mask in clm_pass1]
            )

            for date, nbin in itertools.product(range(bands.shape[0]), range(nbins)):
                if dem_hist_clear_pixels[date, nbin] > 0:
                    dem_mask = np.logical_and(dem_edges[nbin] <= dem, dem < dem_edges[nbin + 1])
                    in_dem_range_clear = np.where(np.logical_and(dem_mask, np.logical_not(clm_pass1[date])))
                    snow_frac[date, nbin] = (
                        np.sum(snow_mask_pass1[date][in_dem_range_clear]) / dem_hist_clear_pixels[date, nbin]
                    )
        return snow_mask_pass1, snow_frac, dem_edges

    def _apply_second_pass(self, bands, ndsi, dem, clm_temp, snow_mask_pass1, snow_frac, dem_edges):
        """Second pass of snow detection"""
        _, height, width, _ = bands.shape
        total_snow_frac = np.sum(snow_mask_pass1, axis=(1, 2)) / (height * width)
        snow_mask_pass2 = np.zeros(snow_mask_pass1.shape)
        for date in range(bands.shape[0]):
            if (total_snow_frac[date] > self.ndsi_params[2]) and (
                snow_frac is not None and np.any(snow_frac[date] > self.dem_params[1])
            ):
                z_s = dem_edges[max(np.argmax(snow_frac[date] > self.dem_params[1]) - 2, 0)]
                snow_mask_pass2[date, :, :] = np.where(
                    np.logical_and(
                        dem > z_s,
                        np.logical_and(
                            np.logical_not(clm_temp[date]),
                            np.logical_and(ndsi[date] > self.ndsi_params[1], bands[date, ..., 1] > self.red_params[-1]),
                        ),
                    ),
                    1,
                    0,
                )
        return snow_mask_pass2

    def execute(self, eopatch):
        """Run multi-pass snow detection"""
        bands = eopatch[self.bands_feature][..., self.band_indices]
        b10 = eopatch[self.bands_feature][..., self.b10_index] if self.b10_index is not None else None
        dem = eopatch[self.dem_feature][..., 0]
        clm = eopatch[self.clm_feature][..., 0]

        with np.errstate(divide="ignore"):
            # (B03 - B11) / (B03 + B11)
            ndsi = (bands[..., 0] - bands[..., 2]) / (bands[..., 0] + bands[..., 2])

        ndsi_invalid = ~np.isfinite(ndsi)
        ndsi[ndsi_invalid] = self.undefined_value

        clm_temp = self._adjust_cloud_mask(bands, clm, dem, b10)

        snow_mask_pass1, snow_frac, dem_edges = self._apply_first_pass(bands, ndsi, clm, dem, clm_temp)

        snow_mask_pass2 = self._apply_second_pass(bands, ndsi, dem, clm_temp, snow_mask_pass1, snow_frac, dem_edges)

        snow_mask = self._apply_dilation(np.logical_or(snow_mask_pass1, snow_mask_pass2))

        eopatch[self.mask_feature] = snow_mask[..., np.newaxis].astype(bool)

        return eopatch


import hashlib
import logging
import subprocess
from abc import abstractmethod
from collections import namedtuple
from os.path import exists, getsize
from typing import Any, Dict, NamedTuple, Optional, Tuple

import casadi as cas
import numpy as np
from jax import jit, vmap

import lumos.numpy as lnp
from lumos.optimal_control.nlp import MappedConstraints, JaxConstraints, CasConstraints
from lumos.models.composition import CompositeModel
from lumos.optimal_control.nlp import CasConstraints, JaxConstraints

logger = logging.getLogger(__name__)

# Convert from names and arrays to dictionary
def _array_to_dict(names, values):
    # We could do dict(zip(names, values)), but unfortunately this does NOT work
    # for casadi as casadi matrices are designed to be non-iterable
    # see: https://github.com/casadi/casadi/issues/2278
    return {name: values[idx] for idx, name in enumerate(names)}

# We use namedtuple to make the io names immutable after generation. They will be used
# for two types of names:
# the direct names: which are fully defined by the model itself
# the 'names': which is the eventual IO taking into consideration submodels and etc.
ModelIO = namedtuple("BaseIO", ("inputs", "outputs", "residuals"))
StateSpaceIO = namedtuple(
    "StateSpaceIO",
    ("states", "states_dot", "inputs", "outputs", "con_outputs", "residuals"),
)

def model_io(
    inputs: Tuple[str] = (), outputs: Tuple[str] = (), residuals: Tuple[str] = (),
):
    """Decorator to set the input and output names of a stateless model."""

    def wrapper(cls):
        cls._direct_names = ModelIO(inputs=inputs, outputs=outputs, residuals=residuals)
        return cls

    return wrapper

def state_space_io(
    states: Tuple[str] = (),
    inputs: Tuple[str] = (),
    outputs: Tuple[str] = (),
    con_outputs: Tuple[str] = (),
    residuals: Tuple[str] = (),
):
    """Decorator to set the input and output names of a StateSpaceModel"""

    def wrapper(cls):

        # TODO: an issue with using a decorator instead of directly in the class is that
        # the decorator is executed during import time, which might be harder to debug.
        if not issubclass(cls, StateSpaceModel):
            raise RuntimeError(
                "state_space_io can only be used for setting the io for "
                "StateSpaceModel. Use model_io for other models"
            )

        cls._direct_names = StateSpaceIO(
            states=states,
            states_dot=states,  # use the same name as states for the states_dot
            inputs=inputs,
            outputs=outputs,
            con_outputs=con_outputs,
            residuals=residuals,
        )

        return cls

    return wrapper

class ModelReturn(NamedTuple):
    """Return data structure for a static model."""

    outputs: Dict = {}
    residuals: Dict = {}

class StateSpaceModelReturn(NamedTuple):
    """Return data structure for state space model."""

    states_dot: Dict = {}
    outputs: Dict = {}
    con_outputs: Dict = {}
    residuals: Dict = {}

class ArrayModelReturn(NamedTuple):
    """Return data structure for a static model in array forms."""

    outputs: lnp.ndarray = np.array([])
    residuals: lnp.ndarray = np.array([])

class ArrayStateSpaceModelReturn(NamedTuple):
    """Return data structure for state space model in array forms."""

    # NOTE: we use np empty array instead of lnp.array([]) here because the type is
    # not really affected by user backend choice as it's determined already during the
    # import time as the default values.
    states_dot: lnp.ndarray = np.array([])
    outputs: lnp.ndarray = np.array([])
    con_outputs: lnp.ndarray = np.array([])
    residuals: lnp.ndarray = np.array([])

class Model(CompositeModel):
    """Abstract class for mathematical models of the simplest form."""

    # parameters of the model.
    # TODO: need to constrain this more.
    _params: Dict[str, Any]

    # TODO: not designed yet
    model_config: Dict[str, Any]

    def __init__(
        self, model_config: Dict[str, Any] = {}, params: Dict[str, Any] = {},
    ):
        super().__init__(model_config=model_config, params=params)

        # Build instance specific name attributes
        self._construct_io_names()

    @abstractmethod
    def forward(self, *args, **kwargs):
        """Abstract method for calling the model.

        We name the method to 'forward' to emphasis it's the forward execution of a
        model.h
        """
        pass

    def apply_and_forward(self, inputs, params):
        self.set_recursive_params(params)
        return self.forward(inputs)

    def apply_and_forward_with_arrays(self, inputs, params):
        self.set_recursive_params(params)
        return self.forward_with_array(inputs)

    def _construct_io_names(self):
        """Create model io names while also taking into account submodels compositoin"""

        self.names = ModelIO(
            inputs=self._direct_names.inputs,
            residuals=self._direct_names.residuals,
            outputs=self._direct_names.outputs
            + tuple(self._collect_children_outputs()),
        )

    def _collect_children_outputs(self):
        """Collect all children outputs and prefix them with submodel_name

        eg: the 'power' output of the 'engine' submodel becomes 'engine.power'
        """
        children_outputs = []
        if not self.is_leaf():
            for submodel_name, model in self._submodels.items():
                children_outputs += [
                    submodel_name + "." + n for n in model.names.outputs
                ]
        return children_outputs

    def combine_submodel_outputs(self, **kwargs):
        """combine the outputs from submodels into large dictionary

        kwargs: {name_of_submodel: vector_of_outputs}
        """
        combined_dict = {}
        for submodel_name, submodel_outputs in kwargs.items():
            # TODO: maybe we could convert the following to a helper method?
            combined_dict.update(
                {submodel_name + "." + n: v for n, v in submodel_outputs.items()}
            )

        return combined_dict

    @classmethod
    def get_direct_group_names(cls, group: str) -> Tuple[str, ...]:
        """Return the direct names of variables inside an IO group."""

        return getattr(cls._direct_names, group)

    def get_group_names(self, group: str) -> Tuple[str, ...]:
        """Return the names of variables inside an IO group."""

        return getattr(self.names, group)

    def get_var_index(self, group: str, name: str) -> int:
        if name not in self.get_group_names(group):
            raise KeyError(
                f"{name} doesnot exist in group: {group}. Valid names are {self.get_group_names(group)} "
            )
        return self.get_group_names(group).index(name)

    def get_var_index_in_flat(self, group: str, name: str) -> int:
        """Return variable index in the flat input"""

        # offset before the group
        offsets = np.cumsum([self.get_num_vars(g) for g in self._implicit_inputs])[:-1]
        offsets = dict(zip(self._implicit_inputs, np.insert(offsets, 0, 0)))

        return self.get_var_index(group=group, name=name) + offsets[group]

    def get_group_indices_in_flat(self, group: str) -> np.ndarray:
        """Return the indices for a group of varialbes in the flat input."""
        return np.array(
            [self.get_var_index_in_flat(group, n) for n in self.get_group_names(group)],
            dtype=np.int32,
        )

    def get_num_vars(self, group: str) -> int:
        return len(self.get_group_names(group))

    # TODO: these are convenience methods. Can we dynamically generate them?
    def get_input(self, inputs: lnp.ndarray, name: str) -> float:
        return inputs[self.get_var_index(group="inputs", name=name)]

    def get_output(self, outputs: lnp.ndarray, name: str) -> float:
        return outputs[self.get_var_index(group="outputs", name=name)]

    def forward_with_arrays(self, inputs):
        inputs = _array_to_dict(self.names.inputs, inputs)
        model_return = self.forward(inputs)

        # Convert from dictionary to arrays for the outputs
        kwargs = {
            g: self.make_vector(g, **getattr(model_return, g))
            for g in model_return._fields
        }

        return ArrayModelReturn(**kwargs)

    @property
    def num_inputs(self):
        return self.get_num_vars(group="inputs")

    @property
    def num_outputs(self):
        return self.get_num_vars(group="outputs")

    @property
    def num_residuals(self):
        return self.get_num_vars(group="residuals")

    @property
    def num_con_outputs(self):
        return self.get_num_vars(group="con_outputs")

    def make_const_vector(self, group: str, val: float = 0.0):
        """Create an array of constant value representing a group."""

        return np.ones(self.get_num_vars(group=group)) * val

    def make_random_vector(self, group: str):
        """Create a random vector representing a group

        We make all the values positive for now.
        """
        return np.abs(np.random.randn(self.get_num_vars(group=group)))

    def _check_keys(self, group, kwargs):
        # check missing names
        input_set = set(kwargs.keys())
        expected_set = set(self.get_group_names(group))
        if not input_set == expected_set:
            if expected_set.issubset(input_set):
                logger.warning(
                    f"{input_set - expected_set} are not valid {group} names."
                    " These values are ignored."
                )
            else:
                raise ValueError(
                    f"Missing {group} values for {expected_set - input_set}"
                )

    def make_vector(self, group: str, **kwargs) -> lnp.ndarray:
        """Create a state vector from kwargs. All values must be provided."""
        self._check_keys(group, kwargs)
        return lnp.array(list(kwargs[name] for name in self.get_group_names(group)))

    def make_dict(self, group: str, **kwargs) -> Dict[str, Any]:
        """Create a dictionary from kwargs. All values must be provided.

        This is actually just a thing wrapper on the standard dictionary construction,
        but it additionally checks if all the necessary keys exist
        """
        self._check_keys(group, kwargs)
        return kwargs

    def make_const_dict(self, group, value: float) -> Dict[str, Any]:
        """Create a dictionary for a group filled with constant values."""
        return {n: value for n in self.get_group_names(group)}

    def plot(self, *args, **kwargs):
        raise NotImplementedError

class StateSpaceModel(Model):
    """Compare to a standard model, now we have states."""

    # We name the inputs to the system simply "inputs" instead of "controls" in control
    # literature. This is so that StateSpaceModel can be naturally seen as a child of
    # the standard Model, where the only difference is the addition of states.

    # TODO: maybe we need a better name for this as this is now only used for the
    # flat input calls which is then used in ocp
    # Maybe explicit_inputs, and implicit_inputs? But for explicit formulation, we also
    # need the outputs defined, so it's not just 'inputs'
    _implicit_inputs: Tuple[str, str, str, str] = (
        "states",
        "inputs",
        "states_dot",
        "con_outputs",
    )

    def __init__(
        self, params: Dict[str, Any] = {}, model_config: Dict[str, Any] = {},
    ):
        super().__init__(model_config=model_config, params=params)
        self._check_names()

    @abstractmethod
    def forward(
        self, states: lnp.ndarray, inputs: lnp.ndarray, mesh: float
    ) -> StateSpaceModelReturn:
        """The canonical form.

        x_dot = f(x, u, t, p)
        y     = g(x, u, t, p)
        """
        pass

    def apply_and_forward(self, states, inputs, mesh, params):
        self.set_recursive_params(params)
        return self.forward(states, inputs, mesh)

    def apply_and_forward_with_arrays(self, states, inputs, mesh, params):
        self.set_recursive_params(params)
        return self.forward_with_arrays(states, inputs, mesh)

    def make_state_space_model_return(
        self,
        states_dot: lnp.ndarray,
        outputs: lnp.ndarray = None,
        residuals: lnp.ndarray = None,
    ) -> StateSpaceModelReturn:
        """Thin wrapper for StateSpaceModelReturn to handle con_outputs automatically."""
        kwargs = {"states_dot": states_dot}
        if outputs is not None:
            kwargs["outputs"] = outputs

        if residuals is not None:
            kwargs["residuals"] = residuals

        return StateSpaceModelReturn(**kwargs)

    def _construct_io_names(self):
        """Create model io names while also taking into account submodels compositoin.

        Similar to Model._construct_io_names, but now needs to operate on more groups
        for state space model.
        """

        self.names = StateSpaceIO(
            inputs=self._direct_names.inputs,
            states=self._direct_names.states,
            states_dot=self._direct_names.states_dot,
            con_outputs=self._direct_names.con_outputs,
            residuals=self._direct_names.residuals,
            outputs=self._direct_names.outputs
            + tuple(self._collect_children_outputs()),
        )

    def _check_names(self):
        # Ensure con_outputs all exist
        for c in self.get_group_names("con_outputs"):
            if c not in self.get_group_names("outputs"):
                raise ValueError(f"constraint outputs {c} not found in outputs")

    def _extract_con_outputs(self, outputs: Dict[str, Any]) -> Dict[str, Any]:
        return {c: outputs[c] for c in self.get_group_names("con_outputs")}

    @property
    def num_implicit_res(self):
        return self.num_states + self.num_con_outputs + self.num_residuals

    @property
    def num_implicit_var(self):
        return sum([self.get_num_vars(g) for g in self._implicit_inputs])

    @property
    def num_states(self):
        return self.get_num_vars(group="states")

    def get_state(self, states: lnp.ndarray, name: str) -> float:
        return states[self.get_var_index(group="states", name=name)]

    def forward_with_arrays(self, states, inputs, mesh):

        states = _array_to_dict(self.names.states, states)
        inputs = _array_to_dict(self.names.inputs, inputs)
        model_return = self.forward(states, inputs, mesh)

        # Convert from dictionary to arrays for the outputs
        kwargs = {
            g: self.make_vector(g, **getattr(model_return, g))
            for g in model_return._fields
            if g != "con_outputs"
        }

        kwargs["con_outputs"] = self.make_vector(
            "con_outputs", **self._extract_con_outputs(model_return.outputs)
        )

        return ArrayStateSpaceModelReturn(**kwargs)

    def implicit(
        self,
        states: lnp.ndarray,
        inputs: lnp.ndarray,
        states_dot: lnp.ndarray,
        con_outputs: lnp.ndarray,
        mesh: float,
    ) -> lnp.ndarray:
        """Implicit form: f(x, u, t, x_dot, y, p) = 0

        This is the base class method for models with a forward method to be turned into
        an implicit model.

        For models that are really implicit, the user should implement the implicit form
        directly.

        the order of the outputs should be: [states_dot, con_outputs, residuals]. This
        is a hard-coded limitation of the current design.
        """

        model_return = self.forward_with_arrays(states, inputs, mesh)

        # TODO: here we return an array, but maybe we should at least always check
        # (especially for user-defined ones) that the residual size is correct.
        #
        # TODO: maybe we should also refer to constraints by name, just like what we did
        # for the decision variables.

        # Assemble scaled residuals. We could move scales outside of the residual
        # functions so the scales are no longer compiled into the residual calls. But
        # that would require manually scaling the jacobian and hessian as well -- Doable
        # but not necessary for now.

        res = lnp.vector_concat(
            [
                model_return.states_dot - states_dot,
                model_return.con_outputs - con_outputs,
                model_return.residuals,
            ]
        )
        return res

    def _split_flat_vars(self, flat_vars: lnp.ndarray):
        split_indices = list(
            np.cumsum([self.get_num_vars(g) for g in self._implicit_inputs])[:-1]
        )
        list_vars = lnp.vector_split(flat_vars, split_indices)

        return dict(zip(self._implicit_inputs, list_vars))

    def _apply_and_flat_implicit(
        self, flat_vars: lnp.ndarray, mesh: float, params,
    ) -> lnp.ndarray:
        self.set_recursive_params(params)
        dict_vars = self._split_flat_vars(flat_vars)

        return self.implicit(**dict_vars, mesh=mesh)

    def _implicit_jac(self, flat_vars, mesh, params):
        raise NotImplementedError(
            "_implicit_jac and _implicit_jacobianstructure needs to be implemented to "
            "use custom jacobians"
        )

    def _implicit_jacobianstructure(self):
        raise NotImplementedError(
            "_implicit_jac and _implicit_jacobianstructure needs to be implemented to "
            "use custom jacobians"
        )

    def _implicit_hess(self, flat_vars, mesh, params, mult):
        raise NotImplementedError(
            "_implicit_hess and _implicit_hessianstructure needs to be implemented to "
            "use custom hessian"
        )

    def _implicit_hessianstructure(self):
        raise NotImplementedError(
            "_implicit_hess and _implicit_hessianstructure needs to be implemented to "
            "use custom hessian"
        )

    def batched_forward(
        self,
        states: lnp.ndarray,
        inputs: lnp.ndarray,
        mesh: float,
        params: Optional[Dict[str, Any]] = None,
    ) -> StateSpaceModelReturn:
        return self._batched_forward(states, inputs, mesh, params)

    def make_model_algebra_cons(self, backend: str):
        """Create the model_algebra constraints that are needed for the OCP.

        FIXME: currently we also require this function to formulate the _batched_forward
        function, which we should probaly consider moving somewhere else.
        """
        if backend == "casadi":
            self._make_casadi_model_algebra_cons()
        elif backend == "jax":
            self._make_jax_model_algebra_cons()
        elif backend == "custom":
            self._make_custom_model_algebra_cons()
        else:
            raise ValueError(
                f"{backend} is not supported. Only 'jax', 'casadi' and 'custom' backends are supported."
            )

    def _model_algebra_jacobianstructure(self):
        """A pessimistic estimate for the jacobian structure of model algebra.

        NOTE: Here we rely on a few dangerous assumptions:
        1) the implicit form is a explicit turned into implicit (see model.implicit),
        with states_dot and con_outputs only acting as linear variables.
        2) the ordering of constraints are assumed to be defined by self.implicit

        TODO: if we really want to, we could use casadi to get the symbolic jac struct
        and then apply it to jax functions (and potentially other backend).
        """
        # states and inputs are involved in all constraints
        # FIXME: this relies on the ordering of the constraints
        rows = np.stack(
            [np.arange(self.num_implicit_res)] * (self.num_states + self.num_inputs)
        ).T.ravel()
        cols = np.stack(
            [
                np.concatenate(
                    [
                        self.get_group_indices_in_flat("states"),
                        self.get_group_indices_in_flat("inputs"),
                    ]
                )
            ]
            * self.num_implicit_res
        ).ravel()

        # states_dot part
        # FIXME: this relies on both the ordering of the constraints and the variables
        states_dot_rows = np.arange(self.num_states)
        states_dot_cols = self.get_group_indices_in_flat("states_dot")

        # con_outputs part
        # FIXME: this relies on both the ordering of the constraints and the variables
        con_outputs_rows = self.num_states + np.arange(self.num_con_outputs)
        con_outputs_cols = self.get_group_indices_in_flat("con_outputs")

        rows = np.concatenate([rows, states_dot_rows, con_outputs_rows])
        cols = np.concatenate([cols, states_dot_cols, con_outputs_cols])

        return rows, cols

    def _model_algebra_hessianstructure(self):
        """A pessimistic estimate for the hessian structure of model algebra.

        NOTE: states_dot and con_outputs are only linear, so no hessian entries for
        them, but for truely implicit equations where states_dot or con_outputs are
        algebraic variables they could become nonlinear! But in those case, the
        condensed approach also won't work. (because it has no explicit ODE to work on)

        TODO: if we really want to, we could use casadi to get the symbolic jac struct
        and then apply it to jax functions (and potentially other backend).
        """

        rows, cols = np.nonzero(np.ones((self.num_implicit_var, self.num_implicit_var)))

        # remove those related to states_dot and con_outputs
        idx_remove = np.hstack(
            [
                self.get_group_indices_in_flat("states_dot"),
                self.get_group_indices_in_flat("con_outputs"),
            ]
        )

        keep_rows = np.array([r not in idx_remove for r in rows])
        keep_cols = np.array([c not in idx_remove for c in cols])
        keep = keep_rows & keep_cols

        return rows[keep], cols[keep]

    def _make_jax_model_algebra_cons(self):
        # For jax, we rely on jac and hessian to compute the constraints at a later
        # stage. See JaxConstraints
        implicit_functions = {
            "constraints": self._apply_and_flat_implicit,
        }

        # TODO: _stage_hessianstructure not yet implemented
        self.model_algebra = JaxConstraints(
            num_in=self.num_implicit_var,
            num_con=self.num_implicit_res,
            **implicit_functions,
            jacobian_structure=self._model_algebra_jacobianstructure(),
            hessian_structure=self._model_algebra_hessianstructure(),
        )

        # FIXME: handle batched_forward better
        self._batched_forward = lnp.use_backend("jax")(
            jit(vmap(self.apply_and_forward_with_arrays, in_axes=[0, 0, 0, None]))
        )

    def _make_custom_model_algebra_cons(self):
        implicit_functions = {
            "constraints": self._apply_and_flat_implicit,
            "jacobian": self._implicit_jac,
            "jacobian_structure": self._implicit_jacobianstructure(),
            "hessian": self._implicit_hess,
            "hessian_structure": self._implicit_hessianstructure(),
        }

        self.model_algebra = MappedConstraints(
            num_in=self.num_implicit_var,
            num_con=self.num_implicit_res,
            **implicit_functions,
        )

        def custom_batched_forward(_states, _inputs, _mesh, _params):
            self.set_recursive_params(_params)
            out = lnp.lmap(self.forward)(_states, _inputs, _mesh)
            return StateSpaceModelReturn(*out)

        self._batched_forward = custom_batched_forward

    def _make_casadi_model_algebra_cons(self, CasType: type = cas.MX):
        # NOTE: for large linear operations, like those in FC layers, MX is much faster
        # than SX both for compilation and executtion
        params = self.get_recursive_params()
        flat_params, unravel = params.tree_ravel()

        cas_flat_params = CasType.sym("params", len(flat_params))
        cas_dict_params = unravel(cas_flat_params)

        mesh = CasType.sym("distance")
        states = CasType.sym("states", self.num_states)
        inputs = CasType.sym("inputs", self.num_inputs)
        stage_vars = CasType.sym("stage_vars", self.num_implicit_var)
        mult = CasType.sym("mult", self.num_implicit_res)

        # Only the model calls need to go inside the context manager.
        with lnp.use_backend("casadi"):
            model_return = self.apply_and_forward_with_arrays(
                states, inputs, mesh, cas_dict_params
            )
            res = self._apply_and_flat_implicit(stage_vars, mesh, cas_dict_params)
            lagrange = lnp.dot(mult, res)

        jac = cas.jacobian(res, stage_vars)
        lagrange_hessian, lagrange_gradient = cas.hessian(lagrange, stage_vars)

        # Generate code
        filename = "nlpfunctions"
        cfile = filename + ".c"
        # FIXME: path management, currently local directory only
        codegen = cas.CodeGenerator(cfile)

        codegen.add(
            cas.Function(
                "forward", [states, inputs, mesh, cas_flat_params], [*model_return]
            )
        )
        codegen.add(
            cas.Function("implicit_con", [stage_vars, mesh, cas_flat_params], [res])
        )
        codegen.add(
            cas.Function("implicit_jac", [stage_vars, mesh, cas_flat_params], [jac])
        )
        codegen.add(
            cas.Function(
                "implicit_hess",
                [stage_vars, mesh, cas_flat_params, mult],
                [lagrange_hessian],
            )
        )
        codegen.generate()

        logger.info(f"Generated c-code with {getsize(cfile)} lines: ")
        # Call compiler if no library exists already
        with open(cfile, "rb") as f:
            file_hash = hashlib.md5()
            for chunk in iter(lambda: f.read(8192), b""):
                file_hash.update(chunk)

        library_name = file_hash.hexdigest()
        library_path = "./" + library_name + ".so"

        if exists(library_path):
            logger.info(f"{library_path} already exists, no compliation is needed.")
        else:
            logger.info(f"Compiling casadi library {library_path}")
            cmd = ["gcc", "-fPIC", "-o2", "-shared", cfile, "-o", library_path]
            p = subprocess.Popen(cmd)
            p.wait()

        # Wrap the functions to take in dict params
        # FIXME: fixed thread number
        _mapped_forward = lnp.cmap(
            cas.external("forward", library_path),
            num_workers=32,
            in_axes=[0, 0, 0, None],
        )

        def cas_batched_forward(_states, _inputs, _mesh, _params):
            _flat_params, _ = _params.tree_ravel()
            out = _mapped_forward(_states, _inputs, _mesh, _flat_params)
            return StateSpaceModelReturn(*out)

        self._batched_forward = cas_batched_forward

        # add sparsity structure
        # Can we not get this from the symbolic SX? eg, jac.sparsity(), but then what?
        # This get_triplet interface almost looks like a bug...
        # [*row_indices, [*col_indices]] is what is returned
        # see: http://casadi.sourceforge.net/api/html/d5/da8/classcasadi_1_1Sparsity.html#a1f3eed93488c3cf121f47bec4956927f
        *jac_rows, jac_cols = jac.sparsity().get_triplet()
        *hess_rows, hess_cols = lagrange_hessian.sparsity().get_triplet()

        implicit_functions = {
            "constraints": cas.external("implicit_con", library_path),
            "jacobian": cas.external("implicit_jac", library_path),
            "hessian": cas.external("implicit_hess", library_path),
            "jacobian_structure": (np.array(jac_rows), np.array(jac_cols)),
            "hessian_structure": (np.array(hess_rows), np.array(hess_cols)),
        }

        self.model_algebra = CasConstraints(
            num_in=self.num_implicit_var,
            num_con=self.num_implicit_res,
            **implicit_functions,
        )

    def export_c_mex(self, cfile: str, CasType: type = cas.MX):
        """Export a state space model into c-code that is ready for mex.

        The exported function has the following API:

        states_dot, outputs, con_outputs, residuals = function(states, inputs, mesh, params)
        where each I/O is an array of the corresponding size.

        Args:
            cfile (str): path of the c-file for export, includg .c extension. Limited to
                current working directory only.
            CasType (type, optional): casadi type to use, SX or MX. Defaults to MX.

        """
        # FIXME: path management, currently local directory only, which is a limitation
        # that comes from casadi
        params = self.get_recursive_params()
        flat_params, unravel = params.tree_ravel()

        cas_flat_params = CasType.sym("params", len(flat_params))
        cas_dict_params = unravel(cas_flat_params)

        mesh = CasType.sym("distance")
        states = CasType.sym("states", self.num_states)
        inputs = CasType.sym("inputs", self.num_inputs)

        # Only the model calls need to go inside the context manager.
        with lnp.use_backend("casadi"):
            model_return = self.apply_and_forward_with_arrays(
                states, inputs, mesh, cas_dict_params
            )

        # Generate code
        codegen = cas.CodeGenerator(cfile, dict(mex=True, main=True))

        codegen.add(
            cas.Function(
                "forward", [states, inputs, mesh, cas_flat_params], [*model_return]
            )
        )

        codegen.generate()

        # Need to set the parameters back to the original params, as they were replaced
        # with casadi variables during the apply_and_forward_with_arrays
        # TODO: the same also happens to other casadi export and also jax tracing right?
        # Can we generalize this problem?
        self.set_recursive_params(params)


# -*- coding: utf-8 -*-
"""
sRGB Colourspace
================

Defines the *sRGB* colourspace:

-   :attr:`colour.models.RGB_COLOURSPACE_sRGB`.

References
----------
-   :cite:`InternationalElectrotechnicalCommission1999a` : International
    Electrotechnical Commission. (1999). IEC 61966-2-1:1999 - Multimedia
    systems and equipment - Colour measurement and management - Part 2-1:
    Colour management - Default RGB colour space - sRGB (p. 51).
    https://webstore.iec.ch/publication/6169
-   :cite:`InternationalTelecommunicationUnion2015i` : International
    Telecommunication Union. (2015). Recommendation ITU-R BT.709-6 - Parameter
    values for the HDTV standards for production and international programme
    exchange BT Series Broadcasting service (pp. 1-32).
    https://www.itu.int/dms_pubrec/itu-r/rec/bt/\
R-REC-BT.709-6-201506-I!!PDF-E.pdf
"""

from __future__ import division, unicode_literals

import numpy as np

from colour.colorimetry import CCS_ILLUMINANTS
from colour.models.rgb import RGB_Colourspace, eotf_inverse_sRGB, eotf_sRGB

__author__ = 'Colour Developers'
__copyright__ = 'Copyright (C) 2013-2020 - Colour Developers'
__license__ = 'New BSD License - https://opensource.org/licenses/BSD-3-Clause'
__maintainer__ = 'Colour Developers'
__email__ = 'colour-developers@colour-science.org'
__status__ = 'Production'

__all__ = [
    'PRIMARIES_sRGB', 'WHITEPOINT_NAME_sRGB', 'CCS_WHITEPOINT_sRGB',
    'MATRIX_sRGB_TO_XYZ', 'MATRIX_XYZ_TO_sRGB', 'RGB_COLOURSPACE_sRGB'
]

PRIMARIES_sRGB = np.array([
    [0.6400, 0.3300],
    [0.3000, 0.6000],
    [0.1500, 0.0600],
])
"""
*sRGB* colourspace primaries.

PRIMARIES_sRGB : ndarray, (3, 2)
"""

WHITEPOINT_NAME_sRGB = 'D65'
"""
*sRGB* colourspace whitepoint name.

CCS_WHITEPOINT_sRGB : unicode
"""

CCS_WHITEPOINT_sRGB = (CCS_ILLUMINANTS['CIE 1931 2 Degree Standard Observer'][
    WHITEPOINT_NAME_sRGB])
"""
*sRGB* colourspace whitepoint chromaticity coordinates.

CCS_WHITEPOINT_sRGB : ndarray
"""

MATRIX_sRGB_TO_XYZ = np.array([
    [0.4124, 0.3576, 0.1805],
    [0.2126, 0.7152, 0.0722],
    [0.0193, 0.1192, 0.9505],
])
"""
*sRGB* colourspace to *CIE XYZ* tristimulus values matrix.

MATRIX_sRGB_TO_XYZ : array_like, (3, 3)
"""

MATRIX_XYZ_TO_sRGB = np.array([
    [3.2406, -1.5372, -0.4986],
    [-0.9689, 1.8758, 0.0415],
    [0.0557, -0.2040, 1.0570],
])
"""
*CIE XYZ* tristimulus values to *sRGB* colourspace matrix.

MATRIX_XYZ_TO_sRGB : array_like, (3, 3)
"""

RGB_COLOURSPACE_sRGB = RGB_Colourspace(
    'sRGB',
    PRIMARIES_sRGB,
    CCS_WHITEPOINT_sRGB,
    WHITEPOINT_NAME_sRGB,
    MATRIX_sRGB_TO_XYZ,
    MATRIX_XYZ_TO_sRGB,
    eotf_inverse_sRGB,
    eotf_sRGB,
)
RGB_COLOURSPACE_sRGB.__doc__ = """
*sRGB* colourspace.

References
----------
:cite:`InternationalElectrotechnicalCommission1999a`,
:cite:`InternationalTelecommunicationUnion2015i`

RGB_COLOURSPACE_sRGB : RGB_Colourspace
"""