from __future__ import absolute_import

from dune.perftool.options import get_option
from dune.perftool.generation import (base_class,
                                      class_basename,
                                      class_member,
                                      constructor_parameter,
                                      global_context,
                                      include_file,
                                      initializer_list,
                                      symbol,
                                      template_parameter,
                                      )
from dune.perftool.cgen.clazz import (AccessModifier,
                                      BaseClass,
                                      ClassMember,
                                      )

from pytools import memoize


@template_parameter("operator")
def lop_template_ansatz_gfs():
    return "GFSU"


@template_parameter("operator")
def lop_template_test_gfs():
    return "GFSV"


@symbol
def name_initree_constructor_param():
    return "iniParams"


@class_member("operator")
def define_initree(name):
    param_name = name_initree_constructor_param()
    include_file('dune/common/parametertree.hh', filetag="operatorfile")
    constructor_parameter("const Dune::ParameterTree&", param_name, classtag="operator")
    initializer_list(name, [param_name], classtag="operator")

    return "const Dune::ParameterTree& {};".format(name)


def ufl_measure_to_pdelab_measure(which):
    return {'cell': 'Volume',
            'exterior_facet': 'Boundary',
            'interior_facet': 'Skeleton',
            }.get(which)


@class_member(classtag="operator", access=AccessModifier.PUBLIC)
def _enum_pattern(which):
    return "enum {{ doPattern{} = true }};".format(which)


def enum_pattern():
    from dune.perftool.generation import get_global_context_value
    integral_type = get_global_context_value("integral_type")
    return _enum_pattern(ufl_measure_to_pdelab_measure(integral_type))


def _pattern_baseclass(measure):
    return base_class('Dune::PDELab::Full{}Pattern'.format(measure), classtag="operator")


def pattern_baseclass():
    from dune.perftool.generation import get_global_context_value
    integral_type = get_global_context_value("integral_type")
    return _pattern_baseclass(ufl_measure_to_pdelab_measure(integral_type))


@class_member(classtag="operator", access=AccessModifier.PUBLIC)
def _enum_alpha(which):
    return "enum {{ doAlpha{} = true }};".format(which)


def enum_alpha():
    from dune.perftool.generation import get_global_context_value
    integral_type = get_global_context_value("integral_type")
    return _enum_alpha(ufl_measure_to_pdelab_measure(integral_type))


@symbol
def name_initree_member():
    define_initree("_iniParams")
    return "_iniParams"


@class_basename("operator")
def localoperator_basename():
    return "LocalOperator"


def class_type_from_cache(classtag):
    from dune.perftool.generation import retrieve_cache_items

    # get the basename
    basename = [i for i in retrieve_cache_items(condition="{} and basename".format(classtag))]
    assert len(basename) == 1
    basename = basename[0]

    # get the template parameters
    tparams = [i for i in retrieve_cache_items(condition="{} and template_param".format(classtag))]
    tparam_str = ''
    if len(tparams) > 0:
        tparam_str = '<{}>'.format(', '.join(t for t in tparams))

    return basename, basename + tparam_str


def assembly_routine_signature():
    from dune.perftool.generation import get_global_context_value
    integral_type = get_global_context_value("integral_type")
    form_type = get_global_context_value("form_type")

    aj = "alpha" if form_type == 'residual' else "jacobian"
    which = ufl_measure_to_pdelab_measure(integral_type).lower()

    if integral_type == 'interior_facet':
        raise NotImplementedError

    from dune.perftool.pdelab.geometry import (name_geometry_wrapper,
                                               type_geometry_wrapper,
                                               )
    geot = type_geometry_wrapper()
    geo = name_geometry_wrapper()

    from dune.perftool.pdelab.argument import (name_accumulation_variable,
                                               type_accumulation_variable,
                                               name_coefficientcontainer,
                                               type_coefficientcontainer,
                                               name_testfunctionspace,
                                               type_testfunctionspace,
                                               name_trialfunctionspace,
                                               type_trialfunctionspace,
                                               )
    lfsut = type_trialfunctionspace()
    lfsu = name_trialfunctionspace()
    lfsvt = type_testfunctionspace()
    lfsv = name_testfunctionspace()
    cct = type_coefficientcontainer()
    cc = name_coefficientcontainer()
    avt = type_accumulation_variable()
    av = name_accumulation_variable()

    return ['template<typename {}, typename {}, typename {}, typename {}, typename {}>'.format(geot,
                                                                                               lfsut,
                                                                                               cct,
                                                                                               lfsvt,
                                                                                               avt,
                                                                                               ),
            'void {}_{}(const {}& {}, const {}& {}, const {}& {}, const {}& {}, {}& {}) const'.format(aj,
                                                                                                      which,
                                                                                                      geot,
                                                                                                      geo,
                                                                                                      lfsut,
                                                                                                      lfsu,
                                                                                                      cct,
                                                                                                      cc,
                                                                                                      lfsvt,
                                                                                                      lfsv,
                                                                                                      avt,
                                                                                                      av,
                                                                                                      )
            ]


def generate_kernel(integral):
    integrand = integral.integrand()
    measure = integral.integral_type()
    subdomain_id = integral.subdomain_id()
    subdomain_data = integral.subdomain_data()

    # 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, measure, subdomain_id)

    # 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, p) for i, p in enumerate(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(), preambles=preambles)
    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("(not file) and (not clazz)")

    # 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, Block
        content = signature
        content.append('{')
        if kernel is not None:
            content.extend('  ' + l for l in generate_code(kernel)[0].split('\n'))
        content.append('}')
        ClassMember.__init__(self, content)


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

    # Generate the name by concatenating basename and template parameters
    basename, fullname = class_type_from_cache(tag)

    base_classes = [bc for bc in retrieve_cache_items('{} and baseclass'.format(tag))]
    constructor_params = [bc for bc in retrieve_cache_items('{} and constructor_param'.format(tag))]
    il = [i for i in retrieve_cache_items('{} and initializer'.format(tag))]
    pm = [m for m in retrieve_cache_items('{} and member'.format(tag))]
    tparams = [i for i in retrieve_cache_items('{} and template_param'.format(tag))]

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

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


def generate_localoperator_kernels(formdata, namedata):
    # Extract the relevant attributes of the form data
    form = formdata.preprocessed_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_items
    delete_cache_items()

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

    # Trigger this one once early on to avoid wrong stuff happening
    localoperator_basename()
    lop_template_ansatz_gfs()
    lop_template_test_gfs()
    from dune.perftool.pdelab.parameter import parameterclass_basename
    parameterclass_basename()

    # Make sure there is always the same constructor arguments (even if parameter class is empty)
    from dune.perftool.pdelab.localoperator import name_initree_member
    name_initree_member()
    from dune.perftool.pdelab.parameter import name_paramclass
    name_paramclass()

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

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

    from dune.perftool.generation import global_context
    with global_context(formdata=formdata, namedata=namedata):
        with global_context(form_type='residual'):
            # Generate the necessary residual methods
            for integral in form.integrals():
                with global_context(integral_type=integral.integral_type()):
                    enum_pattern()
                    pattern_baseclass()
                    enum_alpha()
                    kernel = generate_kernel(integral)

                    # Maybe add numerical differentiation
                    if get_option("numerical_jacobian"):
                        include_file("dune/pdelab/localoperator/defaultimp.hh", filetag="operatorfile")

                        _, loptype = class_type_from_cache("operator")
                        which = ufl_measure_to_pdelab_measure(integral.integral_type())
                        base_class("Dune::PDELab::NumericalJacobian{}<{}>".format(which, loptype), classtag="operator")

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

                operator_kernels[(integral.integral_type(), 'residual')] = kernel

        # Generate the necessary jacobian methods
        if not get_option("numerical_jacobian"):
            from ufl import derivative
            from ufl.algorithms import expand_derivatives
            jacform = expand_derivatives(derivative(form, form.coefficients()[0]))

            with global_context(form_type="jacobian"):
                for integral in jacform.integrals():
                    with global_context(integral_type=integral.integral_type()):
                        kernel = generate_kernel(integral)
                    operator_kernels[(integral.integral_type(), 'jacobian')] = kernel

                # Generate dummy functions for those kernels, that vanished in the differentiation process
                # We *could* solve this problem by using lambda_* terms but we do not really want that, so
                # we use empty jacobian assembly methods instead
                alpha_measures = set(i.integral_type() for i in form.integrals())
                jacobian_measures = set(i.integral_type() for i in jacform.integrals())
                for it in alpha_measures - jacobian_measures:
                    operator_kernels[(it, 'jacobian')] = None

        # 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():
        it, ft = method
        with global_context(integral_type=it, form_type=ft):
            signature = assembly_routine_signature()
            operator_methods.append(AssemblyMethod(signature, kernel))

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

    param = cgen_class_from_cache("parameterclass")
    # 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"), "operatorfile", [param, lop])