from __future__ import absolute_import

from dune.perftool.options import get_option
from dune.perftool.generation import include_file, base_class, symbol, generator_factory
from dune.perftool.cgen.clazz import BaseClass, ClassMember

from pytools import memoize


@generator_factory(item_tags=("initializer", "operator"), counted=True, cache_key_generator=lambda *a: a[0])
def initializer_list(obj, params):
    return "{}({})".format(obj, ", ".join(params))


@generator_factory(item_tags=("operator", "member"), counted=True, cache_key_generator=lambda t, n: n)
def define_private_member(_type, name):
    from cgen import Value
    from dune.perftool.cgen.clazz import ClassMember, AccessModifier
    return ClassMember(Value(_type, name), access=AccessModifier.PRIVATE)


@generator_factory(item_tags=("operator", "constructor_param"), counted=True)
def constructor_parameter(_type, name):
    from cgen import Value
    return Value(_type, name)


@symbol
def name_initree_constructor():
    include_file('dune/common/parametertree.hh', filetag="operator")
    constructor_parameter("const Dune::ParameterTree&", "iniParams")
    return "iniParams"


@symbol
def name_initree_member():
    include_file('dune/common/parametertree.hh', filetag="operator")
    define_private_member("const Dune::ParameterTree&", "_iniParams")
    in_constructor = name_initree_constructor()
    initializer_list("_iniParams", [in_constructor])
    return "_iniParams"


@symbol
def localoperator_type():
    # TODO use something from the form here to make it unique
    return "LocalOperator"


@memoize
def measure_specific_details(measure):
    # The return dictionary that this memoized method will grant direct access to.
    ret = {}

    def numerical_jacobian(which):
        if get_option("numerical_jacobian"):
            # Add a base class
            from dune.perftool.pdelab.driver import type_localoperator
            loptype = type_localoperator()
            base_class("Dune::PDELab::NumericalJacobian{}<{}>".format(which, loptype), classtag="operator")

            # Add the initializer list for that base class
            ini = name_initree_member()
            initializer_list("Dune::PDELab::NumericalJacobian{}<{}>".format(which, loptype),
                             ["{}.get(\"numerical_epsilon.{}\", 1e-9)".format(ini, which.lower())])

    if measure == "cell":
        base_class('Dune::PDELab::FullVolumePattern', classtag="operator")
        numerical_jacobian("Volume")

        ret["residual_signature"] = ['template<typename EG, typename LFSV0, typename X, typename LFSV1, typename R>',
                                     'void alpha_volume(const EG& eg, const LFSV0& lfsv0, const X& x, const LFSV1& lfsv1, R& r) const']
        ret["jacobian_signature"] = ['template<typename EG, typename LFSV0, typename X, typename LFSV1, typename J>',
                                     'void jacobian_volume(const EG& eg, const LFSV0& lfsv0, const X& x, const LFSV1& lfsv1, J& jac) const']

    if measure == "exterior_facet":
        base_class('Dune::PDELab::FullBoundaryPattern', classtag="operator")
        numerical_jacobian("Boundary")

        ret["residual_signature"] = ['template<typename IG, typename LFSV0, typename X, typename LFSV1, typename R>',
                                     'void alpha_boundary(const IG& ig, const LFSV0& lfsv0, const X& x, const LFSV1& lfsv1, R& r) const']
        ret["jacobian_signature"] = ['template<typename IG, typename LFSV0, typename X, typename LFSV1, typename J>',
                                     'void jacobian_boundary(const IG& ig, const LFSV0& lfsv0, const X& x, const LFSV1& lfsv1, J& jac) const']

    if measure == "interior_facet":
        base_class('Dune::PDELab::FullSkeletonPattern', classtag="operator")
        numerical_jacobian("Skeleton")

        ret["residual_signature"] = ['template<typename IG, typename LFSV0_S, typename X, typename LFSV1_S, typename LFSV0_N, typename R, typename LFSV1_N>',
                                     'void alpha_skeleton(const IG& ig, const LFSV0_S& lfsv0_s, const X& x_s, const LFSV1_S& lfsv1_s, const LFSV0_N& lfsv0_n, const X& x_n, const LFSV1_N& lfsv1_n, R& r_s, R& r_n) const']
        ret["jacobian_signature"] = ['template<typename IG, typename LFSV0_S, typename X, typename LFSV1_S, typename LFSV0_N, typename LFSV1_N, typename Jac>',
                                     'void jacobian_skeleton(const IG& ig, const LFSV0_S& lfsv0_s, const X& x_s, const LFSV1_S& lfsv1_s, const LFSV0_N& lfsv0_n, const X& x_n, const LFSV1_N& lfsv1_n, Jac& jac_ss, Jac& jac_sn, Jac& jac_ns, Jac& jac_nn) const']

    return ret


def generate_kernel(integrand=None, measure=None):
    assert integrand and measure

    # Get the measure specifics
    specifics = measure_specific_details(measure)

    # Now split the given integrand into accumulation expressions
    from dune.perftool.ufl.transformations.extract_accumulation_terms import split_into_accumulation_terms
    accterms = split_into_accumulation_terms(integrand)

    # Iterate over the terms and generate a kernel
    for term in accterms:
        from dune.perftool.loopy.transformer import transform_accumulation_term
        transform_accumulation_term(term)

    # Extract the information, which is needed to create a loopy kernel.
    # First extracting it, might be useful to alter it before kernel generation.
    from dune.perftool.generation import retrieve_cache_items
    from dune.perftool.loopy.target import DuneTarget
    domains = [i for i in retrieve_cache_items("domain")]
    instructions = [i for i in retrieve_cache_items("instruction")]
    temporaries = {i.name: i for i in retrieve_cache_items("temporary")}
    # preambles = [i for i in retrieve_cache_items("preamble")]
    arguments = [i for i in retrieve_cache_items("argument")]

    # Create the kernel
    from loopy import make_kernel, preprocess_kernel
    # kernel = make_kernel(domains, instructions, arguments, temporary_variables=temporaries, preambles=preambles, target=DuneTarget())
    kernel = make_kernel(domains, instructions, arguments, temporary_variables=temporaries, target=DuneTarget())
    kernel = preprocess_kernel(kernel)

    # All items with the kernel tags can be destroyed once a kernel has been generated
    from dune.perftool.generation import delete_cache_items
    delete_cache_items("kernel")

    # Return the actual code (might instead return kernels...)
    return kernel


class AssemblyMethod(ClassMember):
    def __init__(self, signature, kernel):
        from loopy import generate_code
        from cgen import LiteralLines
        content = LiteralLines('\n' + '\n'.join(signature) + '\n' + generate_code(kernel)[0])
        ClassMember.__init__(self, content)


def cgen_class_from_cache(tag, members=[]):
    from dune.perftool.generation import retrieve_cache_items

    base_classes = [bc for bc in retrieve_cache_items(tags=(tag, "baseclass"), union=False)]
    constructor_params = [bc for bc in retrieve_cache_items(tags=(tag, "constructor_param"), union=False)]
    il = [i for i in retrieve_cache_items(tags=(tag, "initializer"), union=False)]
    pm = [m for m in retrieve_cache_items(tags=(tag, "member"), union=False)]

    from dune.perftool.cgen.clazz import Constructor
    constructor = Constructor(arg_decls=constructor_params, clsname=localoperator_type(), initializer_list=il)

    from dune.perftool.cgen import Class
    return Class(localoperator_type(), base_classes=base_classes, members=members + pm, constructors=[constructor])


def generate_localoperator_kernels(form):
    # For the moment, I do assume that there is but one integral of each type. This might differ
    # if you use different quadrature orders for different terms.
    assert len(form.integrals()) == len(set(i.integral_type() for i in form.integrals()))

    # Reset the generation cache
    from dune.perftool.generation import delete_cache
    delete_cache()

    # Manage includes and base classes that we always need
    include_file('dune/pdelab/gridfunctionspace/gridfunctionspaceutilities.hh', filetag="operator")
    include_file('dune/pdelab/localoperator/idefault.hh', filetag="operator")
    include_file('dune/pdelab/localoperator/flags.hh', filetag="operator")
    include_file('dune/pdelab/localoperator/pattern.hh', filetag="operator")
    include_file('dune/geometry/quadraturerules.hh', filetag="operator")

    base_class('Dune::PDELab::LocalOperatorDefaultFlags', classtag="operator")

    # Have a data structure collect the generated kernels
    operator_kernels = {}

    # Generate the necessary residual methods
    for integral in form.integrals():
        kernel = generate_kernel(integrand=integral.integrand(), measure=integral.integral_type())
        operator_kernels[(integral.integral_type(), 'residual')] = kernel

    # Generate the necessary jacobian methods
    from dune.perftool.options import get_option
    if get_option("numerical_jacobian"):
        include_file("dune/pdelab/localoperator/defaultimp.hh", filetag="operator")
    else:
        from ufl import derivative
        from ufl.algorithms import expand_derivatives
        jacform = expand_derivatives(derivative(form, form.coefficients()[0]))

        for integral in jacform.integrals():
            kernel = generate_kernel(integrand=integral.integrand(), measure=integral.integral_type())
            operator_kernels[(integral.integral_type(), 'jacobian')] = kernel

    # TODO: JacobianApply for matrix-free computations.

    # Return the set of generated kernels
    return operator_kernels


def generate_localoperator_file(kernels):
    operator_methods = []

    # Make generables from the given kernels
    for method, kernel in kernels.items():
        signature = measure_specific_details(method[0])["{}_signature".format(method[1])]
        operator_methods.append(AssemblyMethod(signature, kernel))

    # Write the file!
    from dune.perftool.file import generate_file

    # TODO take the name of this thing from the UFL file
    lop = cgen_class_from_cache("operator", members=operator_methods)
    generate_file(get_option("operator_file"), "operator", [lop])