# -*- coding: utf-8 -*-
"""Algorithm for replacing gradients in an expression with reference gradients and coordinate mappings."""

# Copyright (C) 2008-2016 Martin Sandve Alnæs
#
# This file is part of UFL (https://www.fenicsproject.org)
#
# SPDX-License-Identifier:    LGPL-3.0-or-later
#
# Modified by Lizao Li <lzlarryli@gmail.com>, 2016

from itertools import chain, accumulate, repeat

from ufl_legacy.log import error

from ufl_legacy.core.multiindex import indices
from ufl_legacy.corealg.multifunction import MultiFunction, memoized_handler
from ufl_legacy.algorithms.map_integrands import map_integrand_dags

from ufl_legacy.classes import (ReferenceValue,
                                Jacobian, JacobianInverse, JacobianDeterminant)

from ufl_legacy.tensors import as_tensor, as_vector
from ufl_legacy.utils.sequences import product
import numpy


def sub_elements_with_mappings(element):
    "Return an ordered list of the largest subelements that have a defined mapping."
    if element.mapping() != "undefined":
        return [element]
    elements = []
    for subelm in element.sub_elements():
        if subelm.mapping() != "undefined":
            elements.append(subelm)
        else:
            elements.extend(sub_elements_with_mappings(subelm))
    return elements


def create_nested_lists(shape):
    if len(shape) == 0:
        return [None]
    elif len(shape) == 1:
        return [None] * shape[0]
    else:
        return [create_nested_lists(shape[1:]) for i in range(shape[0])]


def reshape_to_nested_list(components, shape):
    if len(shape) == 0:
        assert len(components) == 1
        return [components[0]]
    elif len(shape) == 1:
        assert len(components) == shape[0]
        return components
    else:
        n = product(shape[1:])
        return [reshape_to_nested_list(components[n * i:n * (i + 1)], shape[1:]) for i in range(shape[0])]


def apply_known_single_pullback(r, element):
    """Apply pullback with given mapping.

    :arg r: Expression wrapped in ReferenceValue
    :arg element: The element defining the mapping
    """
    # Need to pass in r rather than the physical space thing, because
    # the latter may be a ListTensor or similar, rather than a
    # Coefficient/Argument (in the case of mixed elements, see below
    # in apply_single_function_pullbacks), to which we cannot apply ReferenceValue
    mapping = element.mapping()
    domain = r.ufl_domain()
    if mapping == "physical":
        return r
    elif mapping == "identity":
        return r
    elif mapping == "contravariant Piola":
        J = Jacobian(domain)
        detJ = JacobianDeterminant(J)
        transform = (1.0 / detJ) * J
        # Apply transform "row-wise" to TensorElement(PiolaMapped, ...)
        *k, i, j = indices(len(r.ufl_shape) + 1)
        kj = (*k, j)
        f = as_tensor(transform[i, j] * r[kj], (*k, i))
        return f
    elif mapping == "covariant Piola":
        K = JacobianInverse(domain)
        # Apply transform "row-wise" to TensorElement(PiolaMapped, ...)
        *k, i, j = indices(len(r.ufl_shape) + 1)
        kj = (*k, j)
        f = as_tensor(K[j, i] * r[kj], (*k, i))
        return f
    elif mapping == "L2 Piola":
        detJ = JacobianDeterminant(domain)
        return r / detJ
    elif mapping == "double contravariant Piola":
        J = Jacobian(domain)
        detJ = JacobianDeterminant(J)
        transform = (1.0 / detJ) * J
        # Apply transform "row-wise" to TensorElement(PiolaMapped, ...)
        *k, i, j, m, n = indices(len(r.ufl_shape) + 2)
        kmn = (*k, m, n)
        f = as_tensor((1.0 / detJ)**2 * J[i, m] * r[kmn] * J[j, n], (*k, i, j))
        return f
    elif mapping == "double covariant Piola":
        K = JacobianInverse(domain)
        # Apply transform "row-wise" to TensorElement(PiolaMapped, ...)
        *k, i, j, m, n = indices(len(r.ufl_shape) + 2)
        kmn = (*k, m, n)
        f = as_tensor(K[m, i] * r[kmn] * K[n, j], (*k, i, j))
        return f
    else:
        error("Should never be reached!")


def apply_single_function_pullbacks(r, element):
    """Apply an appropriate pullback to something in physical space

    :arg r: An expression wrapped in ReferenceValue.
    :arg element: The element this expression lives in.
    :returns: a pulled back expression."""
    mapping = element.mapping()
    if r.ufl_shape != element.reference_value_shape():
        error("Expecting reference space expression with shape '%s', got '%s'" % (element.reference_value_shape(), r.ufl_shape))
    if mapping in {"physical", "identity",
                   "contravariant Piola", "covariant Piola",
                   "double contravariant Piola", "double covariant Piola",
                   "L2 Piola"}:
        # Base case in recursion through elements. If the element
        # advertises a mapping we know how to handle, do that
        # directly.
        f = apply_known_single_pullback(r, element)
        if f.ufl_shape != element.value_shape():
            error("Expecting pulled back expression with shape '%s', got '%s'" % (element.value_shape(), f.ufl_shape))
        return f
    elif mapping in {"symmetries", "undefined"}:
        # Need to pull back each unique piece of the reference space thing
        gsh = element.value_shape()
        rsh = r.ufl_shape
        if mapping == "symmetries":
            subelem = element.sub_elements()[0]
            fcm = element.flattened_sub_element_mapping()
            offsets = (product(subelem.reference_value_shape()) * i for i in fcm)
            elements = repeat(subelem)
        else:
            elements = sub_elements_with_mappings(element)
            # Python >= 3.8 has an initial keyword argument to
            # accumulate, but 3.7 does not.
            offsets = chain([0],
                            accumulate(product(e.reference_value_shape())
                                       for e in elements))
        rflat = as_vector([r[idx] for idx in numpy.ndindex(rsh)])
        g_components = []
        # For each unique piece in reference space, apply the appropriate pullback
        for offset, subelem in zip(offsets, elements):
            sub_rsh = subelem.reference_value_shape()
            rm = product(sub_rsh)
            rsub = [rflat[offset + i] for i in range(rm)]
            rsub = as_tensor(numpy.asarray(rsub).reshape(sub_rsh))
            rmapped = apply_single_function_pullbacks(rsub, subelem)
            # Flatten into the pulled back expression for the whole thing
            g_components.extend([rmapped[idx]
                                 for idx in numpy.ndindex(rmapped.ufl_shape)])
        # And reshape appropriately
        f = as_tensor(numpy.asarray(g_components).reshape(gsh))
        if f.ufl_shape != element.value_shape():
            error("Expecting pulled back expression with shape '%s', got '%s'" % (element.value_shape(), f.ufl_shape))
        return f
    else:
        error("Unhandled mapping type '%s'" % mapping)


class FunctionPullbackApplier(MultiFunction):
    def __init__(self):
        MultiFunction.__init__(self)

    expr = MultiFunction.reuse_if_untouched

    def terminal(self, t):
        return t

    @memoized_handler
    def form_argument(self, o):
        # Represent 0-derivatives of form arguments on reference
        # element
        f = apply_single_function_pullbacks(ReferenceValue(o), o.ufl_element())
        assert f.ufl_shape == o.ufl_shape
        return f


def apply_function_pullbacks(expr):
    """Change representation of coefficients and arguments in expression
    by applying Piola mappings where applicable and representing all
    form arguments in reference value.

    @param expr:
        An Expr.
    """
    return map_integrand_dags(FunctionPullbackApplier(), expr)
