# -*- coding: utf-8 -*-
"""This module provides a simple way to compute various norms of
:py:class:`Vectors <dolfin.cpp.Vector>` and :py:class:`Functions
<dolfin.functions.function.Function>`, including the standard
:math:`L^2`-norm and other norms."""

# Copyright (C) 2008-2017 Anders Logg
#
# This file is part of DOLFIN.
#
# DOLFIN is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# DOLFIN is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with DOLFIN. If not, see <http://www.gnu.org/licenses/>.
#
# Modified by: Jørgen S. Dokken, 2017
# First added: 2007
# Last changed: 2017-03-17

from six import string_types
import ufl
from ufl import inner, grad, div, curl, dx, FiniteElement, VectorElement, Coefficient, TensorElement
from math import sqrt

import dolfin.cpp as cpp

from dolfin.cpp import GenericVector, GenericFunction, Function, Mesh, error, Vector
from dolfin.fem.assembling import assemble, assemble_multimesh
from dolfin.fem.interpolation import interpolate
from dolfin.functions.function import Function
from dolfin.functions.multimeshfunction import MultiMeshFunction
from dolfin.functions.functionspace import FunctionSpace, \
    VectorFunctionSpace, TensorFunctionSpace, MultiMeshFunctionSpace

__all__ = ["norm", "errornorm"]


def norm(v, norm_type="L2", mesh=None):
    """
    Return the norm of a given vector or function.

    *Arguments*
        v
            a :py:class:`Vector <dolfin.cpp.Vector>` or
            a :py:class:`Function <dolfin.functions.function.Function>` or
            a :py:class:`MultiMeshFunction <dolfin.functions.multimeshfunction.MultiMeshFunction>`
        norm_type
            see below for alternatives.
        mesh
            optional :py:class:`Mesh <dolfin.cpp.Mesh>` or
            a :py:class:`MultiMesh <dolfin.cpp.MultiMesh>` on
            which to compute the norm.

    If the norm type is not specified, the standard :math:`L^2` -norm
    is computed. Possible norm types include:

    *Vectors*

    ================   =================  ================
    Norm               Usage
    ================   =================  ================
    :math:`l^2`        norm(x, 'l2')      Default
    :math:`l^1`        norm(x, 'l1')
    :math:`l^\infty`   norm(x, 'linf')
    ================   =================  ================

    *Functions*

    ================  =================  =================================
    Norm              Usage              Includes the :math:`L^2` -term
    ================  =================  =================================
    :math:`L^2`       norm(v, 'L2')      Yes
    :math:`H^1`       norm(v, 'H1')      Yes
    :math:`H^1_0`     norm(v, 'H10')     No
    :math:`H` (div)   norm(v, 'Hdiv')    Yes
    :math:`H` (div)   norm(v, 'Hdiv0')   No
    :math:`H` (curl)  norm(v, 'Hcurl')   Yes
    :math:`H` (curl)  norm(v, 'Hcurl0')  No
    ================  =================  =================================

    *Examples of usage*

    .. code-block:: python

        v = Function(V)
        x = v.vector()

        print norm(x, 'linf')   # print the infinity norm of vector x

        n = norm(v)             # compute L^2 norm of v
        print norm(v, 'Hdiv')   # print H(div) norm of v
        n = norm(v, 'H1', mesh) # compute H^1 norm of v on given mesh

    """

    if not isinstance(v, (GenericVector, GenericFunction, MultiMeshFunction)):
        cpp.dolfin_error("norms.py",
                         "compute norm",
                         "expected a GenericVector ,GenericFunction or a MultiMeshFunction")
    # Check arguments
    if not isinstance(norm_type, string_types):
        cpp.dolfin_error("norms.py",
                         "compute norm",
                         "Norm type must be a string, not " +
                         str(type(norm_type)))
    if mesh is not None and not isinstance(mesh, (cpp.Mesh, cpp.MultiMesh)):
        cpp.dolfin_error("norms.py",
                         "compute norm",
                         "Expecting a Mesh or a MultiMesh, not " + str(type(mesh)))

    # Get mesh from function
    if isinstance(v, Function) and mesh is None:
        mesh = v.function_space().mesh()
    elif isinstance(v, MultiMeshFunction) and mesh is None:
        mesh = v.function_space().multimesh()

    # Define integration measure and domain
    multimesh = False
    if isinstance(v, MultiMeshFunction):
        multimesh = True
        dc = ufl.dx(mesh) + ufl.dC(mesh)
    else:
        dc = ufl.dx(mesh)
    # Select norm type
    if isinstance(v, GenericVector):
        return v.norm(norm_type.lower())
    elif (isinstance(v, Coefficient) and isinstance(v, (GenericFunction,
                                                        MultiMeshFunction))):
        if norm_type.lower() == "l2":
            M = v**2*dc
        elif norm_type.lower() == "h1":
            M = (v**2 + grad(v)**2)*dc
        elif norm_type.lower() == "h10":
            M = grad(v)**2*dc
        elif norm_type.lower() == "hdiv":
            M = (v**2 + div(v)**2)*dc
        elif norm_type.lower() == "hdiv0":
            M = div(v)**2*dc
        elif norm_type.lower() == "hcurl":
            M = (v**2 + curl(v)**2)*dc
        elif norm_type.lower() == "hcurl0":
            M = curl(v)**2*dc
        else:
            cpp.dolfin_error("norms.py",
                             "compute norm",
                             "Unknown norm type (\"%s\") for functions"
                             % str(norm_type))
    else:
        cpp.dolfin_error("norms.py",
                         "compute norm",
                         "Unknown object type. Must be a vector, function or a multimeshfunction")

    # Assemble value
    if multimesh:
        r = assemble_multimesh(M, form_compiler_parameters={"representation": "quadrature"})
    else:
        r = assemble(M)

    # Check value
    if r < 0.0:
        cpp.dolfin_error("norms.py",
                         "compute norm",
                         "Square of norm is negative, might be a round-off error")
    elif r == 0.0:
        return 0.0
    else:
        return sqrt(r)


def errornorm(u, uh, norm_type="l2", degree_rise=3, mesh=None):
    """
    Compute and return the error :math:`e = u - u_h` in the given norm.

    *Arguments*
        u, uh
            :py:class:`Functions <dolfin.functions.function.Function>`
        norm_type
            Type of norm. The :math:`L^2` -norm is default.
            For other norms, see :py:func:`norm <dolfin.fem.norms.norm>`.
        degree_rise
            The number of degrees above that of u_h used in the
            interpolation; i.e. the degree of piecewise polynomials used
            to approximate :math:`u` and :math:`u_h` will be the degree
            of :math:`u_h` + degree_raise.
        mesh
            Optional :py:class:`Mesh <dolfin.cpp.Mesh>` on
            which to compute the error norm.

    In simple cases, one may just define

    .. code-block:: python

        e = u - uh

    and evalute for example the square of the error in the :math:`L^2` -norm by

    .. code-block:: python

        assemble(e**2*dx(mesh))

    However, this is not stable w.r.t. round-off errors considering that
    the form compiler may expand(#) the expression above to::

        e**2*dx = u**2*dx - 2*u*uh*dx + uh**2*dx

    and this might get further expanded into thousands of terms for
    higher order elements. Thus, the error will be evaluated by adding
    a large number of terms which should sum up to something close to
    zero (if the error is small).

    This module computes the error by first interpolating both
    :math:`u` and :math:`u_h` to a common space (of high accuracy),
    then subtracting the two fields (which is easy since they are
    expressed in the same basis) and then evaluating the integral.

    (#) If using the tensor representation optimizations.
    The quadrature represenation does not suffer from this problem.
    """

    # Check argument
    if not isinstance(u, (cpp.GenericFunction, cpp.MultiMeshFunction)):
        cpp.dolfin_error("norms.py",
                         "compute error norm",
                         "Expecting a Function, MultiMeshFunction or Expression for u")
    if not (isinstance(uh, cpp.Function) or (uh, cpp.MultiMeshFunction)):
        cpp.dolfin_error("norms.py",
                         "compute error norm",
                         "Expecting a Function for uh")
    # Test if errornorm is for multimesh
    multimesh=False

    # Get mesh
    if isinstance(u, Function) and mesh is None:
        mesh = u.function_space().mesh()
    if isinstance(uh, Function) and mesh is None:
        mesh = uh.function_space().mesh()
    if isinstance(uh, MultiMeshFunction) and mesh is None:
        mesh = uh.function_space().multimesh()
        multimesh=True

    if mesh is None:
        cpp.dolfin_error("norms.py",
                         "compute error norm",
                         "Missing mesh")

    # Get rank
    if not u.ufl_shape == uh.ufl_shape:
        cpp.dolfin_error("norms.py",
                         "compute error norm",
                         "Value shapes don't match")
    shape = u.ufl_shape
    rank = len(shape)

    # Check that uh is associated with a finite element
    if uh.ufl_element().degree() is None:
        cpp.dolfin_error("norms.py",
                         "compute error norm",
                         "Function uh must have a finite element")

    # Degree for interpolation space. Raise degree with respect to uh.
    degree = uh.ufl_element().degree() + degree_rise

    # Check degree of 'exact' solution u
    degree_u = u.ufl_element().degree()
    if degree_u is not None and degree_u < degree:
        cpp.warning("Degree of exact solution may be inadequate for accurate result in errornorm.")

    # Create function space
    if multimesh:
        if rank == 0:
            element = FiniteElement("Discontinuous Lagrange", mesh.ufl_cell(),
                                    degree)
        elif rank == 1:
            element = VectorElement("Discontinuous Lagrange", mesh.ufl_cell(),
                                    degree)
        elif rank == 2:
            element = TensorElement("Discontinuous Lagrange", mesh.ufl_cell(),
                                    degree)
        V = MultiMeshFunctionSpace(mesh, element)
    else:
        if rank == 0:
            V = FunctionSpace(mesh, "Discontinuous Lagrange", degree)
        elif rank == 1:
            V = VectorFunctionSpace(mesh, "Discontinuous Lagrange", degree,
                                    dim=shape[0])
        elif rank > 1:
            V = TensorFunctionSpace(mesh, "Discontinuous Lagrange", degree,
                                    shape=shape)

    # Interpolate functions into finite element space
    pi_u = interpolate(u, V)
    pi_uh = interpolate(uh, V)

    # Compute the difference
    if multimesh:
        e = MultiMeshFunction(V)
    else:
        e = Function(V)

    e.assign(pi_u)
    e.vector().axpy(-1.0, pi_uh.vector())

    # Compute norm
    return norm(e, norm_type=norm_type, mesh=mesh)
