"""
Reference: Nagarajan and Nakshatrala, IJNMF, 2011.
Problem 4.3: Two-dimensional problem with anisotropic medium

-div[D(x) grad[u]] = f(x) on the unit square.
u = u0 on the boundary.
"""

import numpy as np
from dolfin import *
import scipy.sparse as sps
import scipy.sparse.linalg as splin
import scipy.optimize as so 
import matplotlib.pyplot as plt

parameters.linear_algebra_backend = "uBLAS"


#=========================================;
#  Create mesh and define function space  ;
#=========================================;
mesh = UnitSquareMesh(20,20)
V = FunctionSpace(mesh, 'Lagrange', 1)

#=====================;
#  Medium properties  ;
#=====================; 
alpha=1.0
d1 = 10000.0 
d2 = 1.0
theta = -pi / 6.0
c = cos(theta)
s = sin(theta)
D = as_matrix([[c*c*d1+s*s*d2,-c*s*(d1-d2)],[-c*s*(d1-d2),s*s*d1+c*c*d2]])

#=============================;
#  Define boubdary conditions ;
#=============================;

u_B = Expression('sin(pi * x[0])')

def bottom_boundary(x,on_boundary):
    tol = 1E-14
    return on_boundary and abs(x[1]) < tol

Gamma_0 = DirichletBC(V,u_B,bottom_boundary)

u_T = Constant(0)

def other_boundary(x,on_boundary):
    tol = 1E-14
    return on_boundary and abs(x[1]) > tol

Gamma_1 = DirichletBC(V,u_T,other_boundary)

bc = [Gamma_0, Gamma_1]

#============================;
#  Define volumetric source  ;
#============================; 
f = Constant(0)

#==============================;
#  Define variational problem  ;
#==============================;
u = TrialFunction(V)
v = TestFunction(V)

a = alpha*inner(u,v)*dx + inner(D*nabla_grad(u), nabla_grad(v))*dx
L = f*v*dx

#===================;
# Compute solution  ;
#===================;
u = Function(V)
# solve(a == L, u, bc)

A, b = assemble_system(a, L, bc)
solve(A,u.vector(),b,'lu')

#=======================================================;
#  Extract stiffness and load vector into numpy format  ;
#=======================================================;
rows, cols, values = A.data()
KMat = sps.csr_matrix((values, cols, rows))
fvec = b.data()

# scisol = splin.spsolve(KMat,fvec)

#===========================;
#  Calling scipy optimizer  ;
#===========================;
def objFunc(x,Aa,fvec):
    return 0.5 * x .dot(Aa.dot(x)) - fvec.dot(x)

def derfunc(x,Aa,fvec):
    return Aa.dot(x) - fvec

neqns = len(fvec)

g = np.zeros(neqns)

bndcons = tuple([(0, 1) for i in range(neqns)])

optsol = so.fmin_tnc(objFunc,g,fprime=derfunc,bounds=bndcons,args=(KMat,fvec,),disp=0)

optimsol = Function(V)
optimsol.vector()[:] = optsol[0]

print "============================================"
print "   Diagnostic report on Galerkin solution   "
print "============================================"
L2_norm_error = sqrt(assemble(dot(u-u_B,u-u_B) * dx))
print "L2 error = ", L2_norm_error
print "Min. concentration =", u.vector().min()
print "Max. concentration =", u.vector().max()

print "==============================================="
print "   Diagnostic report on Nonnegative solution   "
print "==============================================="
L2_norm_error = sqrt(assemble(dot(optimsol-u_B,optimsol-u_B) * dx))
print "L2 error = ", L2_norm_error

print "Min. concentration =", optimsol.vector().min()
print "Max. concentration =", optimsol.vector().max()

# Dump solution to file in VTK format
file = File('Galerkin_P2.pvd')
file << u

file = File('Nonnegative_P2.pvd')
file << optimsol

#===============================;
#  Post-processing the results  ;
#===============================;
plot(u,title="Galerkin formulation")
plot(optimsol,title="Non-negative solution")
plot(mesh,title="Computational mesh")
interactive()