# Copyright (c) 2017, Lawrence Livermore National Security, LLC. Produced at
# the Lawrence Livermore National Laboratory. LLNL-CODE-734707. All Rights
# reserved. See files LICENSE and NOTICE for details.
#
# This file is part of CEED, a collection of benchmarks, miniapps, software
# libraries and APIs for efficient high-order finite element and spectral
# element discretizations for exascale applications. For more information and
# source code availability see http://github.com/ceed.
#
# The CEED research is supported by the Exascale Computing Project 17-SC-20-SC,
# a collaborative effort of two U.S. Department of Energy organizations (Office
# of Science and the National Nuclear Security Administration) responsible for
# the planning and preparation of a capable exascale ecosystem, including
# software, applications, hardware, advanced system engineering and early
# testbed platforms, in support of the nation's exascale computing imperative.

from _ceed_cffi import ffi, lib
import ctypes
import tempfile
from abc import ABC

# ------------------------------------------------------------------------------


class _QFunctionBase(ABC):
    """Ceed QFunction: point-wise operation at quadrature points for evaluating
         volumetric terms."""

    # Attributes
    _ceed = ffi.NULL
    _pointer = ffi.NULL

    # Destructor
    def __del__(self):
        # libCEED call
        lib.CeedQFunctionDestroy(self._pointer)

    # Representation
    def __repr__(self):
        return "<CeedQFunction instance at " + hex(id(self)) + ">"

    # String conversion for print() to stdout
    def __str__(self):
        """View a QFunction via print()."""

        # libCEED call
        with tempfile.NamedTemporaryFile() as key_file:
            with open(key_file.name, 'r+') as stream_file:
                stream = ffi.cast("FILE *", stream_file)

                lib.CeedQFunctionView(self._pointer[0], stream)

                stream_file.seek(0)
                out_string = stream_file.read()

        return out_string

    # Apply CeedQFunction
    def apply(self, q, inputs, outputs):
        """Apply the action of a QFunction.

           Args:
             q: number of quadrature points
             *inputs: array of input data vectors
             *outputs: array of output data vectors"""

        # Array of vectors
        invecs = ffi.new("CeedVector[16]")
        for i in range(min(16, len(inputs))):
            invecs[i] = inputs[i]._pointer[0]
        outvecs = ffi.new("CeedVector[16]")
        for i in range(min(16, len(outputs))):
            outvecs[i] = outputs[i]._pointer[0]

        # libCEED call
        lib.CeedQFunctionApply(self._pointer[0], q, invecs, outvecs)

        # Clean-up
        ffi.release(invecs)
        ffi.release(outvecs)


# ------------------------------------------------------------------------------
class QFunction(_QFunctionBase):
    """Ceed QFunction: point-wise operation at quadrature points for evaluating
         volumetric terms."""

    # Constructor
    def __init__(self, ceed, vlength, f, source):
        # libCEED object
        self._pointer = ffi.new("CeedQFunction *")

        # Reference to Ceed
        self._ceed = ceed

        # Function pointer
        fpointer = ffi.cast(
            "CeedQFunctionUser", ctypes.cast(
                f, ctypes.c_void_p).value)

        # libCEED call
        sourceAscii = ffi.new("char[]", source.encode('ascii'))
        lib.CeedQFunctionCreateInterior(self._ceed._pointer[0], vlength, fpointer,
                                        sourceAscii, self._pointer)

    # Set context data
    def set_context(self, ctx):
        """Set global context for a QFunction.

           Args:
             *ctx: Numpy array holding context data to set"""

        # Setup the numpy array for the libCEED call
        ctx_pointer = ffi.new("CeedScalar *")
        ctx_pointer = ffi.cast("void *", ctx.__array_interface__['data'][0])

        # libCEED call
        lib.CeedQFunctionSetContext(self._pointer[0], ctx_pointer, len(ctx))

    # Add fields to CeedQFunction
    def add_input(self, fieldname, size, emode):
        """Add a QFunction input.

           Args:
             fieldname: name of QFunction field
             size: size of QFunction field, (ncomp * dim) for CEED_EVAL_GRAD or
                     (ncomp * 1) for CEED_EVAL_NONE and CEED_EVAL_INTERP
             **emode: CEED_EVAL_NONE to use values directly,
                      CEED_EVAL_INTERP to use interpolated values,
                      CEED_EVAL_GRAD to use gradients."""

        # libCEED call
        fieldnameAscii = ffi.new("char[]", fieldname.encode('ascii'))
        lib.CeedQFunctionAddInput(
            self._pointer[0], fieldnameAscii, size, emode)

    def add_output(self, fieldname, size, emode):
        """Add a QFunction output.

           Args:
             fieldname: name of QFunction field
             size: size of QFunction field, (ncomp * dim) for CEED_EVAL_GRAD or
                     (ncomp * 1) for CEED_EVAL_NONE and CEED_EVAL_INTERP
             **emode: CEED_EVAL_NONE to use values directly,
                      CEED_EVAL_INTERP to use interpolated values,
                      CEED_EVAL_GRAD to use gradients."""

        # libCEED call
        fieldnameAscii = ffi.new("char[]", fieldname.encode('ascii'))
        lib.CeedQFunctionAddOutput(
            self._pointer[0], fieldnameAscii, size, emode)

# ------------------------------------------------------------------------------


class QFunctionByName(_QFunctionBase):
    """Ceed QFunction By Name: point-wise operation at quadrature points
         from a given gallery, for evaluating volumetric terms."""

    # Constructor
    def __init__(self, ceed, name):
        # libCEED object
        self._pointer = ffi.new("CeedQFunction *")

        # Reference to Ceed
        self._ceed = ceed

        # libCEED call
        nameAscii = ffi.new("char[]", name.encode('ascii'))
        lib.CeedQFunctionCreateInteriorByName(self._ceed._pointer[0], nameAscii,
                                              self._pointer)

# ------------------------------------------------------------------------------


class IdentityQFunction(_QFunctionBase):
    """Ceed Identity QFunction: identity qfunction operation."""

    # Constructor
    def __init__(self, ceed, size, inmode, outmode):
        # libCEED object
        self._pointer = ffi.new("CeedQFunction *")

        # Reference to Ceed
        self._ceed = ceed

        # libCEED call
        lib.CeedQFunctionCreateIdentity(self._ceed._pointer[0], size, inmode,
                                        outmode, self._pointer)

# ------------------------------------------------------------------------------