"""This module extends the form algebra with a collection of operators
based on the basic form algebra operations."""

__author__ = "Anders Logg (logg@tti-c.org)"
__date__ = "2005-09-07 -- 2005-12-20"
__copyright__ = "Copyright (c) 2005-2006 Anders Logg"
__license__  = "GNU GPL Version 2"

# Modified by Ola Skavhaug, 2005

# Python modules
import sys
import Numeric

# FFC common modules
sys.path.append("../../")
from ffc.common.exceptions import *

# FFC compiler modules
from index import *
from algebra import *
from projection import *
from finiteelement import *

def Identity(n):
    "Return identity matrix of given size."
    # Let Numeric handle the identity
    return Numeric.identity(n)

def rank(v):
    "Return rank for given object."
    if isinstance(v, BasisFunction):
        return v.element.rank() - len(v.component)
    elif isinstance(v, Product):
        return rank(v.basisfunctions[0])
    elif isinstance(v, Sum):
        return rank(v.products[0])
    elif isinstance(v, Function):
        return rank(Sum(v))
    else:
        return Numeric.rank(v)
    return 0

def vec(v):
    "Create vector of scalar functions from given vector-valued function."
    # Check if we already have a vector
    if isinstance(v, list):
        return v
    # Check if we have an element of the algebra
    if isinstance(v, Element):
        # Check that we have a vector
        if not rank(v) == 1:
            raise FormError, (v, "Cannot create vector from scalar expression.")
        # Get vector dimension
        n = __tensordim(v, 0)
        # Create list of scalar components
        return [v[i] for i in range(n)]        
    # Let Numeric handle the conversion
    if isinstance(v, Numeric.ArrayType) and len(v.shape) == 1:
        return v.tolist()
    # Unable to find a proper conversion
    raise FormError, (v, "Unable to convert given expression to a vector,")

def dot(v, w):
    "Return scalar product of given functions."
    # Check ranks
    if rank(v) == rank(w) == 1:
        # Check dimensions
        if not len(v) == len(w):
            raise FormError, ((v, w), "Dimensions don't match for scalar product.")
        # Use index notation if possible
        if isinstance(v, Element) and isinstance(w, Element):
            i = Index()
            return v[i]*w[i]
        # Otherwise, use Numeric.dot
        return Numeric.dot(vec(v), vec(w))
    elif rank(v) == rank(w) == 2:
        # Check dimensions
        if not len(v) == len(w):
            raise FormError, ((v, w), "Dimensions don't match for scalar product.")
        # Compute dot product (:) of matrices
        return Numeric.sum([v[i][j]*w[i][j] for i in range(len(v)) for j in range(len(v[i]))])

def cross(v, w):
    "Return cross product of given functions."
    # Check dimensions
    if not len(v) == len(w):
        raise FormError, ((v, w), "Cross product only defined for vectors in R^3.")
    # Compute cross product
    return [v[1]*w[2] - v[2]*w[1], v[2]*w[0] - v[0]*w[2], v[0]*w[1] - v[1]*w[0]]

def trace(v):
    "Return trace of given matrix"
    # Let Numeric handle the trace
    return Numeric.trace(v)

def transp(v):
    "Return transpose of given matrix."
    # Let Numeric handle the transpose."
    return Numeric.transpose(v)

def mult(v, w):
    "Compute matrix-matrix product of given matrices."
    # First, convert to Numeric.array (safe for both array and list arguments)
    vv = Numeric.array(v)
    ww = Numeric.array(w)
    if len(vv.shape) == 0 or len(ww.shape) == 0:
        # One argument is a scalar
        return vv*ww
    if len(vv.shape) == len(ww.shape) == 1:
        # Vector times vector
        return Numeric.multiply(vv, ww) 
    elif len(vv.shape) == 2 and (len(ww.shape) == 1 or len(ww.shape) == 2):
        # Matvec or matmat product, use matrixmultiply instead
        return Numeric.matrixmultiply(vv, ww)
    else:
        raise FormError, ((v, w), "Dimensions don't match for multiplication.")

def D(v, i):
    "Return derivative of v in given coordinate direction."
    # Use member function dx() if possible
    if isinstance(v, Element):
        return v.dx(i)
    # Otherwise, apply to each component
    return [D(v[j], i) for j in range(len(v))]
    
def grad(v):
    "Return gradient of given function."
    # Get shape dimension
    d = __shapedim(v)
    # Check if we have a vector
    if rank(v) == 1:
        return [ [D(v[i], j) for j in range(d)] for i in range(len(v)) ]
    # Otherwise assume we have a scalar
    return [D(v, i) for i in range(d)]

def div(v):
    "Return divergence of given function."
    # Check dimensions
    if not len(v) == __shapedim(v):
        raise FormError, (v, "Dimensions don't match for divergence.")
    # Use index notation if possible
    if isinstance(v, Element):
        i = Index()
        return v[i].dx(i)
    # Otherwise, use Numeric.sum
    return Numeric.sum([D(v[i], i) for i in range(len(v))])

def rot(v):
    "Return rotation of given function."
    # Check dimensions
    if not len(v) == __shapedim(v) == 3:
        raise FormError, (v, "Rotation only defined for v : R^3 --> R^3")
    # Compute rotation
    return [D(v[2], 1) - D(v[1], 2), D(v[0], 2) - D(v[2], 0), D(v[1], 0) - D(v[0], 1)]

def curl(v):
    "Alternative name for rot."
    return rot(v)

def mean(v):
    "Return mean value of given Function (projection onto piecewise constants)."
    # Check that we got a Function
    if not isinstance(v, Function):
        raise FormError, (v, "Mean values are only supported for Functions.")
    # Different projections needed for scalar and vector-valued elements
    element = v.e0
    if element.rank() == 0:
        P0 = FiniteElement("Discontinuous Lagrange", element.shape_str, 0)
        pi = Projection(P0)
        return pi(v)
    else:
        P0 = FiniteElement("Discontinuous vector Lagrange", element.shape_str, 0, element.tensordim(0))
        pi = Projection(P0)
        return pi(v)
     
def __shapedim(v):
    "Return shape dimension for given object."
    if isinstance(v, list):
        # Check that all components have the same shape dimension
        for i in range(len(v) - 1):
            if not __shapedim(v[i]) == __shapedim(v[i + 1]):
                raise FormError, (v, "Components have different shape dimensions.")
        # Return length of first term
        return __shapedim(v[0])
    elif isinstance(v, BasisFunction):
        return v.element.shapedim()
    elif isinstance(v, Product):
        return __shapedim(v.basisfunctions[0])
    elif isinstance(v, Sum):
        return __shapedim(v.products[0])
    elif isinstance(v, Function):
        return __shapedim(Sum(v))
    else:
        raise FormError, (v, "Shape dimension is not defined for given expression.")
    return 0

def __tensordim(v, i):
    "Return size of given dimension for given object."
    if i < 0 or i >= rank(v):
        raise FormError, ((v, i), "Tensor dimension out of range.")
    if isinstance(v, BasisFunction):
        return v.element.tensordim(i + len(v.component))
    elif isinstance(v, Product):
        return __tensordim(v.basisfunctions[0], i)
    elif isinstance(v, Sum):
        return __tensordim(v.products[0], i)
    elif isinstance(v, Function):
        return __tensordim(Sum(v), i)
    else:
        raise FormError, ((v, i), "Tensor dimension is not defined for given expression.")
    return 0

if __name__ == "__main__":

    scalar = FiniteElement("Lagrange", "tetrahedron", 2)
    vector = FiniteElement("Vector Lagrange", "tetrahedron", 2)

    i = Index()

    v = BasisFunction(scalar)
    u = BasisFunction(scalar)
    w = Function(scalar)

    V = BasisFunction(vector)
    U = BasisFunction(vector)
    W = Function(vector)
    
    i = Index()
    j = Index()

    dx = Integral()

    print dot(grad(v), grad(u))*dx
    print vec(U)
    print dot(U, V)
    print dot(vec(V), vec(U))
    print dot(U, grad(v))
    print div(U)
    print dot(rot(V), rot(U))
    print div(grad(dot(rot(V), U)))*dx
    print cross(V, U)
    print trace(mult(Identity(len(V)), grad(V)))
