"""This module is a simple implementation of a general FEM assembly algorithm for educational purposes. The implementation is partially based on Puffin."""

__author__ = "Johan Jansson (jjan@csc.kth.se)"
__date__ = "2009-10-05"
__copyright__ = "Copyright (C) 2009 Johan Jansson"
__license__  = "GNU LGPL Version 2.1"

from dolfin import *
from numpy import *
from math import *

def myassemble(mesh):

    # Define basis functions and their gradients on the reference elements.
    def phi0(x):
        return 1.0 - x[0] - x[1]

    def phi1(x):
        return x[0]

    def phi2(x):
        return x[1]

    def f(x):
        return 1.0

    phi = [phi0, phi1, phi2]
    dphi = array([[-1.0, -1.0], [1.0, 0.0], [0.0, 1.0]])

    # Define quadrature points
    midpoints = array([[0.5, 0.0], [0.5, 0.5], [0.0, 0.5]])

    N = mesh.num_vertices()

    A = Matrix(N, N)
    A.apply()

    b = Vector(N)
    b.apply()

    coord = zeros([3, 2])
    nodes = zeros(3)

    # Iterate over all cells, adding integral contribution to matrix/vector
    for c in cells(mesh):

        # Extract node numbers and vertex coordinates
        nodes = c.entities(0).astype('int')
        for i in range(0, 3):
            v = Vertex(mesh, int(nodes[i]))
            for j in range(0, 2):
                coord[i][j] = v.point()[j]

        # Compute Jacobian of map and area of cell
        J = outer(coord[0, :], dphi[0]) + \
            outer(coord[1, :], dphi[1]) + \
            outer(coord[2, :], dphi[2])
        dx = 0.5 * abs(linalg.det(J))

        # Iterate over quadrature points
        for p in midpoints:
            # Map p to physical cell
            x = coord[0, :] * phi[0](p) + \
                coord[1, :] * phi[1](p) + \
                coord[2, :] * phi[2](p)

            # Iterate over test functions
            for i in range(0, 3):
                v = phi[i](p)
                dv = linalg.solve(J.transpose(), dphi[i])

                # Assemble vector (linear form)
                integral = f(x)*v*dx / 3.0
                b[nodes[i]] += integral

                # Iterate over trial functions
                for j in range(0, 3):
                    u = phi[j](p)
                    du = linalg.solve(J.transpose(), dphi[j])

                    # Assemble matrix (bilinear form)
                    integral = (inner(du, dv)) * dx / 3.0
                    A[nodes[i], nodes[j]] += integral

    A.apply()
    b.apply()
    print A.array()
    print b.array()
    return (A, b)

def solve(mesh):

    (A, b) = myassemble(mesh)

    # Apply boundary condition weakly by adding contribution computed
    # with DOLFIN
    class Kappa(Expression):
        def eval(self, values, x):
            values[0] = 1.0e3
            if(x[0] < 1.0e-5):
                values[0] = 0.0

    V = FunctionSpace(mesh, "Lagrange", 1)
    u = TrialFunction(V)
    v = TestFunction(V)
    kappa = Kappa(V)
    a = kappa*u*v*ds
    A2 = assemble(a)

    Atot = A + A2

    # Construct U function
    V = FunctionSpace(mesh, "CG", 1)
    U = Function(V)
    xi = U.vector()

    # Solve for xi
    solver = LinearSolver()
    solver.solve(Atot, xi, b)

    return U