from pytools import memoize
from dune.perftool.options import get_option
from dune.perftool.generation import generator_factory

# Define the generators used in-here
operator_include = generator_factory(item_tags=("pdelab", "include", "operator"), on_store=lambda i: "#include<{}>".format(i), no_deco=True)
base_class = generator_factory(item_tags=("pdelab", "baseclass", "operator"), counted=True, no_deco=True)
initializer_list = generator_factory(item_tags=("pdelab", "initializer", "operator"), counted=True)

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

    if measure == "cell":
        base_class('Dune::PDELab::FullVolumePattern')
        if get_option("numerical_jacobian"):
            base_class("Dune::PDELab::NumericalJacobianVolume<{}>".format())

        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')

        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')

        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_term(integrand=None, measure=None):
    assert integrand and measure

    # Delete all non-include parts of the cache.
    # TODO: add things such as base classes as cache items.
    from dune.perftool.generation import delete_cache
    delete_cache()

    # Get the measure specifics
    specifics = measure_specific_details(measure)

    # Transform the expression. All relevant results are then in the cache
    from dune.perftool.transformer import transform_expression
    transform_expression(integrand)

    # 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.target import DuneTarget
    domains = [i for i in retrieve_cache_items("domain")]
    instructions = [i for i in retrieve_cache_items("cinstruction")] \
                 + [i for i in retrieve_cache_items("exprinstruction")]
    temporaries = {i.name:i for i in retrieve_cache_items("temporary")}
    preambles = [i[1] for i in retrieve_cache_items("preamble")]

    print "Printing the information that we found:\n\nDomains:"
    for d in domains:
        print d
    print "\nInstructions:"
    for i in instructions:
        print i
    print "\nTemporaries:"
    for t, v in temporaries.items():
        print "{}: {}".format(t, v)
    print "\nPreambles:"
    for p in preambles:
        print p

#     import loopy
#     import numpy
#     from loopy import make_kernel, preprocess_kernel
#     from dune.perftool.target import DuneTarget
#     k = make_kernel(domains, instructions,
#                 [loopy.GlobalArg("grad_a0", numpy.float64, shape=("argi_n", "dim")),
#                  loopy.GlobalArg("grad_c0", numpy.float64, shape=("dim",)),
#                  loopy.ValueArg("dim", numpy.int32),
#                  loopy.ValueArg("q_n", numpy.int32),
#                  loopy.ValueArg("argi_n", numpy.int32)],
#                     temporary_variables=temporaries, target=DuneTarget())
#     k = preprocess_kernel(k)
    # Create the kernel
    from loopy import make_kernel, preprocess_kernel, add_dtypes
    kernel = make_kernel(domains, instructions, temporary_variables=temporaries, preambles=preambles, target=DuneTarget())
    import numpy
    kernel = add_dtypes(kernel, dict(grad_a0=numpy.float64, grad_c0=numpy.float64))
    kernel = preprocess_kernel(kernel)

    # Return the actual code (might instead return kernels...)
    from loopy import generate_code
    return str(generate_code(kernel)[0])


def generate_localoperator(ufldata, operatorfile):
    form = ufldata.preprocessed_form

    operator_methods = []

    # 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()

    # Generate the necessary residual methods
    for integral in form.integrals():
        body = generate_term(integrand=integral.integrand(), measure=integral.integral_type())
        signature = measure_specific_details(integral.integral_type())["residual_signature"]
        operator_methods.append((signature, body))

    # Generate the necessary jacobian methods
    from dune.perftool.options import get_option
    if get_option("numerical_jacobian"):
        operator_include("dune/pdelab/localoperator/defaultimp.hh")
    else:
        pass

    # TODO: JacobianApply for matrix-free computations.

    # Manage includes and base classes that we always need
    operator_include('dune/pdelab/gridfunctionspace/gridfunctionspaceutilities.hh')
    operator_include('dune/pdelab/localoperator/idefault.hh')
    operator_include('dune/pdelab/localoperator/flags.hh')
    operator_include('dune/pdelab/localoperator/pattern.hh')
    operator_include('dune/common/parametertree.hh')
    operator_include('dune/geometry/quadraturerules.hh')

    base_class('Dune::PDELab::LocalOperatorDefaultFlags')

    print operator_methods[0]