from __future__ import absolute_import
from os.path import splitext

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


def name_form(formdata, data):
    # Check wether the formdata has a name in UFL
    try:
        name = data.object_names[id(formdata.original_form)]
        return name
    except:
        for index, form in enumerate(data.forms):
            if formdata.preprocessed_form.equals(form):
                name = str(index)
                return name
    # If the form has no name and can not be found in data.forms something went wrong
    assert False


def name_localoperator_file(formdata, data):
    from dune.perftool.options import get_option
    if len(data.forms) == 1:
        filename = get_option("operator_file")
    else:
        suffix = '_' + name_form(formdata, data)
        basename, extension = splitext(get_option("operator_file"))
        filename = basename + suffix + extension
    return filename


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


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


@template_parameter("operator")
def lop_template_range_field():
    return "RF"


@class_member("operator")
def lop_domain_field(name):
    # TODO: Rethink for not Galerkin Method
    gfs = lop_template_ansatz_gfs()
    return "using {} = typename {}::Traits::GridView::ctype;".format(name, gfs)


def name_domain_field():
    name = "DF"
    lop_domain_field(name)
    return name


def lop_template_gfs(ma):
    from ufl.classes import Argument, Coefficient
    if isinstance(ma.argexpr, Argument):
        if ma.argexpr.number() == 0:
            return lop_template_test_gfs()
        if ma.argexpr.number() == 1:
            return lop_template_ansatz_gfs()
    if isinstance(ma.argexpr, Coefficient):
        # Index 0 is reserved for trialfunction, index 1 is reserved for jacobian apply function
        assert ma.argexpr.count() < 2
        return lop_template_ansatz_gfs()
    assert False


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


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


@class_basename("operator")
def localoperator_basename(formdata, data):
    form_name = name_form(formdata, data)
    return "LocalOperator" + form_name.capitalize()


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 assembler_routine_name():
    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")

    part1 = {"residual": "alpha"}.get(form_type, form_type)
    part2 = ufl_measure_to_pdelab_measure(integral_type).lower()

    return "{}_{}".format(part1, part2)


def assembly_routine_signature(formdata):
    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")

    # Check if form is linear
    from dune.perftool.pdelab.driver import is_linear
    linear = is_linear(formdata.original_form)

    if form_type == 'residual':
        if integral_type == 'cell':
            from dune.perftool.pdelab.signatures import alpha_volume_signature
            return alpha_volume_signature()
        if integral_type == 'exterior_facet':
            from dune.perftool.pdelab.signatures import alpha_boundary_signature
            return alpha_boundary_signature()
        if integral_type == 'interior_facet':
            from dune.perftool.pdelab.signatures import alpha_skeleton_signature
            return alpha_skeleton_signature()

    if form_type == 'jacobian':
        if integral_type == 'cell':
            from dune.perftool.pdelab.signatures import jacobian_volume_signature
            return jacobian_volume_signature()
        if integral_type == 'exterior_facet':
            from dune.perftool.pdelab.signatures import jacobian_boundary_signature
            return jacobian_boundary_signature()
        if integral_type == 'interior_facet':
            from dune.perftool.pdelab.signatures import jacobian_skeleton_signature
            return jacobian_skeleton_signature()

    if form_type == 'jacobian_apply':
        if linear:
            if integral_type == 'cell':
                from dune.perftool.pdelab.signatures import jacobian_apply_volume_signature
                return jacobian_apply_volume_signature()
            if integral_type == 'exterior_facet':
                from dune.perftool.pdelab.signatures import jacobian_apply_boundary_signature
                return jacobian_apply_boundary_signature()
            if integral_type == 'interior_facet':
                from dune.perftool.pdelab.signatures import jacobian_apply_skeleton_signature
                return jacobian_apply_skeleton_signature()
        else:
            if integral_type == 'cell':
                from dune.perftool.pdelab.signatures import nonlinear_jacobian_apply_volume_signature
                return nonlinear_jacobian_apply_volume_signature()
            if integral_type == 'exterior_facet':
                from dune.perftool.pdelab.signatures import nonlinear_jacobian_apply_boundary_signature
                return nonlinear_jacobian_apply_boundary_signature()
            if integral_type == 'interior_facet':
                from dune.perftool.pdelab.signatures import nonlinear_jacobian_apply_skeleton_signature
                return nonlinear_jacobian_apply_skeleton_signature()

    assert False


def generate_kernel(integrals):
    subst_rules = []
    for integral in integrals:
        integrand = integral.integrand()
        measure = integral.integral_type()
        subdomain_id = integral.subdomain_id()
        subdomain_data = integral.subdomain_data()

        from dune.perftool.ufl.dimensionindex import dimension_index_mapping
        dimension_indices = dimension_index_mapping(integrand)

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

        # Get a transformer instance for this kernel
        from dune.perftool.loopy.transformer import UFL2LoopyVisitor
        visitor = UFL2LoopyVisitor(measure, subdomain_id, dimension_indices)

        # Iterate over the terms and generate a kernel
        for term in accterms:
            visitor(term)
            subst_rules.extend(visitor.substitution_rules)

    # 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_functions, 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")}
    arguments = [i for i in retrieve_cache_items("argument")]
    manglers = retrieve_cache_functions("mangler")

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

    from loopy import make_reduction_inames_unique
    kernel = make_reduction_inames_unique(kernel)

    kernel = preprocess_kernel(kernel)

    # see whether we need iname duplication to make this thing schedulable!
    # TODO: Use a clever strategy here, instead of random transformation until the problem is resolved
    from loopy import has_schedulable_iname_nesting, get_iname_duplication_options, duplicate_inames
    while not has_schedulable_iname_nesting(kernel):
        # If there is a duplication that solves the problem with just one duplication, we pick that one
        iname, within = (None, None)
        for i, w in get_iname_duplication_options(kernel):
            if has_schedulable_iname_nesting(duplicate_inames(kernel, i, w)):
                iname, within = (i, w)

        # Otherwise pick a random one.
        if iname is None:
            iname, within = next(get_iname_duplication_options(kernel))

        # Do the transformation
        print("Applying iname duplication to measure {}: iname {}; within {}".format(measure, iname, within))
        kernel = duplicate_inames(kernel, iname, within)

    # Loopy might have introduced some temporary variables during preprocessing. As I want to have my own
    # temporary declaration code right now, I call the declaration preamble manually.
    for added_tv in set(kernel.temporary_variables.keys()) - set(temporaries.keys()):
        from dune.perftool.generation.loopy import default_declaration
        default_declaration(added_tv)

    # Now add the preambles to the kernel
    preambles = [(i, p) for i, p in enumerate(retrieve_cache_items("preamble"))]
    kernel = kernel.copy(preambles=preambles)

    # 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 kernel


def name_time_dumper_os():
    return "os"


def name_time_dumper_reset():
    return "reset"


def name_time_dumper_t():
    return "t"


def name_time_dumper_counter():
    return "counter"


def name_time_dumper_exec():
    return "exec"


class TimerMethod(ClassMember):
    def __init__(self):
        os = name_time_dumper_os()
        reset = name_time_dumper_reset()
        t = name_time_dumper_t()
        ex = name_time_dumper_exec()
        content = ["template <typename Stream>",
                   "void dump_timers(Stream& {}, char* {}, bool {})".format(os, ex, reset),
                   "{",
                   "  double {} = 0.0;".format(t),
                   "#ifdef ENABLE_COUNTERS",
                   "  auto counter = HP_TIMER_OPCOUNTERS({})",
                   "  counter.reset();",
                   "#endif",
                   ""]
        dump_timers = [i for i in retrieve_cache_items(condition='dump_timers')]
        content.extend(map(lambda x: '  ' + x, dump_timers))
        content.extend(["#ifdef ENABLE_COUNTERS",
                        "  counter.reportOperations({});".format(os),
                        "#endif"])
        content.append("}")
        ClassMember.__init__(self, content)


class AssemblyMethod(ClassMember):
    def __init__(self, signature, kernel, filename):
        from loopy import generate_body
        from cgen import LiteralLines, Block
        content = signature
        content.append('{')
        if kernel is not None:
            # Add kernel preamble
            for i, p in kernel.preambles:
                content.append('  ' + p)

            # Start timer
            if get_option('timer'):
                timer_name = assembler_routine_name() + '_kernel'
                post_include('HP_DECLARE_TIMER({});'.format(timer_name), filetag='operatorfile')
                content.append('  ' + 'HP_TIMER_START({});'.format(timer_name))
                dump_accumulate_timer(timer_name)

            # Add kernel body
            content.extend(l for l in generate_body(kernel).split('\n')[1:-1])

            # Stop timer
            if get_option('timer'):
                content.append('  ' + 'HP_TIMER_STOP({});'.format(timer_name))

        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))]
    from cgen import Block
    constructor_block = Block(contents=[i for i in retrieve_cache_items("{} and constructor_block".format(tag), make_generable=True)])
    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(block=constructor_block, 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, data):
    # Extract the relevant attributes of the form data
    form = formdata.preprocessed_form

    # 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(formdata, data)
    lop_template_ansatz_gfs()
    lop_template_test_gfs()
    lop_template_range_field()
    from dune.perftool.pdelab.parameter import parameterclass_basename
    parameterclass_basename(formdata, data)

    # 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")
    from dune.perftool.pdelab.driver import is_stationary
    if not is_stationary():
        rf = lop_template_range_field()
        base_class('Dune::PDELab::InstationaryLocalOperatorDefaultMethods<{}>'.format(rf), classtag="operator")

        # Create set time method in parameter class
        from dune.perftool.pdelab.parameter import define_set_time_method
        define_set_time_method()

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

    with global_context(form_type='residual'):
        # Generate the necessary residual methods
        for measure in set(i.integral_type() for i in form.integrals()):
            with global_context(integral_type=measure):
                enum_pattern()
                pattern_baseclass()
                enum_alpha()
                kernel = generate_kernel(form.integrals_by_type(measure))

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

                    # Numerical jacobian base class
                    _, loptype = class_type_from_cache("operator")
                    which = ufl_measure_to_pdelab_measure(measure)
                    base_class("Dune::PDELab::NumericalJacobian{}<{}>".format(which, loptype), classtag="operator")

                    # Numerical jacobian initializer list
                    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",
                                     )

                    # In the case of matrix free operator evaluation we need jacobian apply methods
                    if get_option("matrix_free"):
                        from dune.perftool.pdelab.driver import is_linear
                        if is_linear(formdata.original_form):
                            # Numeical jacobian apply base class
                            base_class("Dune::PDELab::NumericalJacobianApply{}<{}>".format(which, loptype), classtag="operator")

                            # Numerical jacobian apply initializer list
                            initializer_list("Dune::PDELab::NumericalJacobianApply{}<{}>".format(which, loptype),
                                             ["{}.get<double>(\"numerical_epsilon.{}\", 1e-9)".format(ini_constructor, ini, which.lower())],
                                             classtag="operator",
                                             )
                        else:
                            # Numerical nonlinear jacobian apply base class
                            base_class("Dune::PDELab::NumericalNonlinearJacobianApply{}<{}>".format(which, loptype), classtag="operator")

                            # Numerical nonlinear jacobian apply initializer list
                            initializer_list("Dune::PDELab::NumericalNonlinearJacobianApply{}<{}>".format(which, loptype),
                                             ["{}.get<double>(\"numerical_epsilon.{}\", 1e-9)".format(ini_constructor, ini, which.lower())],
                                             classtag="operator",
                                             )

                operator_kernels[(measure, 'residual')] = kernel

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

        from dune.perftool.ufl.preprocess import preprocess_form
        jacform = preprocess_form(jacform).preprocessed_form

        with global_context(form_type="jacobian"):
            for measure in set(i.integral_type() for i in jacform.integrals()):
                with global_context(integral_type=measure):
                    kernel = generate_kernel(jacform.integrals_by_type(measure))
                operator_kernels[(measure, '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

        # Jacobian apply methods for matrix-free computations
        if get_option("matrix_free"):
            # The apply vector has reserved index 1 so we directly use Coefficient class from ufl
            from ufl import Coefficient
            apply_coefficient = Coefficient(form.coefficients()[0].ufl_element(), 1)

            # Create application of jacobian on vector
            from ufl import action
            jac_apply_form = action(jacform, apply_coefficient)

            # Create kernel for jacobian application
            with global_context(form_type="jacobian_apply"):
                for measure in set(i.integral_type() for i in jac_apply_form.integrals()):
                    with global_context(integral_type=measure):
                        kernel = generate_kernel(jac_apply_form.integrals_by_type(measure))
                    operator_kernels[(measure, 'jacobian_apply')] = 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_apply_measures = set(i.integral_type() for i in jac_apply_form.integrals())
                    for it in alpha_measures - jacobian_apply_measures:
                        operator_kernels[(it, 'jacobian_apply')] = None

    # Return the set of generated kernels
    return operator_kernels


def generate_localoperator_file(formdata, kernels, filename):
    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(formdata)
            operator_methods.append(AssemblyMethod(signature, kernel, filename))

    if get_option('timer'):
        include_file('dune/perftool/common/timer.hh', filetag='operatorfile')
        operator_methods.append(TimerMethod())

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


def generate_localoperator_basefile(formdatas, data):
    filename = get_option("operator_file")
    for formdata in formdatas:
        lop_filename = name_localoperator_file(formdata, data)
        include_file(lop_filename, filetag="operatorbasefile")

    from dune.perftool.file import generate_file
    generate_file(filename, "operatorbasefile", [])