# -*- coding: utf-8 -*-
"""
    Created on Mon Nov 28 2016
    
    @author: karthik
"""

#######################################################################

from firedrake import *
from firedrake.petsc import PETSc
import numpy as np


# Create mesh
parameters["matnest"] = False

mesh = Mesh('try.msh')
#====================================================;
#  Define function spaces and mixed (product) space  ;
#====================================================;
velSpace = VectorFunctionSpace(mesh,"CG",2)
pSpace = FunctionSpace(mesh, "CG", 1)
dSpace = TensorFunctionSpace(mesh, "CG", 1)
W = velSpace * pSpace

#===================================;
#  Define trial and test functions  ;
#===================================;
(v,p) = TrialFunctions(W)
(w,q) = TestFunctions(W)
solution = Function(W)

# BOUNDS
vimin = 0.0
vimax = 1.0

#=====================;
#  Define body force  ;
#=====================;
rhob = interpolate(Expression(('0','0','0')),velSpace)

#============================;
#  Define medium properties  ;
#============================; 

theta = pi / 3.0
c = cos(theta)
s = sin(theta)
d1 = 1000.0
d2 = 1.0
d3 = 1.0

#D0 = as_matrix([[d1,0.,0.],[0.,d2,0.],[0.,0.,d3]])
#R = as_matrix([[c,-s, 0.],[s,c, 0.],[0.,0.,1.]])
#alpha = R* (D0 * R.T)
#invalpha = inv(alpha)
alpha = interpolate(Expression(("d1*c*c+d2*s*s","c*s*(d2-d1)","0","c*s*(d2-d1)","d1*s*s+d2*c*c","0","0","0","d3"),d1=d1,d2=d2,d3=d3,c=c,s=s),dSpace)
invalpha = inv(alpha)
#alpha = Constant(1.0)
#invalpha = inv(alpha)

#===========================;
#  Define variational form  ;
#===========================;
n = FacetNormal(mesh)

solver_type = "VMS"

p_in = Constant(vimax)
p_out = Constant(vimin)
LSW = Constant(0.001)

if solver_type == "VMS":
    a = inner(alpha*v,w)*dx - div(w)*p*dx - div(v)*q*dx \
        - 0.5*inner(invalpha*(alpha*v + grad(p)), alpha*w + grad(q))*dx 
    L = inner(rhob,w)*dx - 0.5*inner(invalpha*rhob, alpha*w + grad(q))*dx \
        - inner(w,n)*p_in*(ds(43)+ds(44)) - dot(w,n)*p_out*(ds(41)+ds(42))
elif solver_type == "Galerkin":
    a = dot(alpha*v,w)*dx - div(w)*p*dx - div(v)*q*dx 
    L = dot(rhob,w)*dx 
elif solver_type == "LS":
    a = LSW*dot(alpha*w + grad(q),alpha*v + grad(p))*dx + div(w)*div(v)*dx
    L = LSW*dot(alpha*w+grad(q),rhob)*dx 
else:
    print "The given option is not available."

#========================;
#  VI solver parameters  ;
#========================;
solver_parameters = {
   "ksp_type": "gmres",
   "pc_type": "fieldsplit",
   "pc_fieldsplit_type": "schur",
   "pc_fieldsplit_schur_fact_type":"full",
   "pc_fieldsplit_schur_precondition": "selfp",
   "fieldsplit_0_ksp_type": "preonly",
   "fieldsplit_0_pc_type": "bjacobi",
   "fieldsplit_1_ksp_type": "preonly",
   "fieldsplit_1_pc_type": "hypre",
   "fieldsplit_1_pc_hypre_type": "boomeramg",
   "fieldsplit_1_pc_hypre_boomeramg_strong_threshold": 0.75,
   "fieldsplit_1_pc_hypre_boomeramg_agg_nl": 2,
   "ksp_monitor": True
}
petsc_options = PETSc.Options()
petsc_options.prefixPush("rs_")
petsc_options.setValue("ksp_type", "gmres")
#petsc_options.setValue("pc_type", "gamg")
petsc_options.setValue("pc_type", "fieldsplit")
petsc_options.setValue("pc_fieldsplit_type", "schur")
petsc_options.setValue("pc_fieldsplit_schur_fact_type","full")
petsc_options.setValue("pc_fieldsplit_schur_precondition", "selfp")
petsc_options.setValue("fieldsplit_0_ksp_type", "preonly")
petsc_options.setValue("fieldsplit_0_pc_type", "bjacobi")
petsc_options.setValue("fieldsplit_1_ksp_type", "preonly")
petsc_options.setValue("fieldsplit_1_pc_type", "hypre")
petsc_options.setValue("fieldsplit_1_pc_hypre_type", "boomeramg")
petsc_options.setValue("fieldsplit_1_pc_hypre_boomeramg_strong_threshold", 0.75)
petsc_options.setValue("fieldsplit_1_pc_hypre_boomeramg_agg_nl", 2)
petsc_options.prefixPop()

#==========================;
#  Variational Inequality  ;
#==========================;
rs_solver = PETSc.SNES().create(PETSc.COMM_WORLD)
rs_solver.setOptionsPrefix("rs_")

A = assemble(a)
A.M._force_evaluation()
b = assemble(L)
b.dat._force_evaluation()

lb = Function(W)
ub = Function(W)

lb_vel,lb_pres = lb.split()
ub_vel,ub_pres = ub.split()
lb_vel.assign(PETSc.NINFINITY)
ub_vel.assign(PETSc.INFINITY)
lb_pres.assign(vimin)
ub_pres.assign(vimax)

def rs_formJac(snes, petsc_x, petsc_J, petsc_JP):
  pass
def rs_formFunc(snes, petsc_x, petsc_g, A=None, b=None):
  with b.dat.vec as b_vec:
    A.M.handle.mult(petsc_x, petsc_g)
    petsc_g.axpy(-1.0,b_vec)
rs_conc = Function(W)
def virs():
  with solution.dat.vec as sol_vec, lb.dat.vec as lb_vec, ub.dat.vec as ub_vec:
    with rs_conc.dat.vec as con_vec:
      rs_solver.setFunction(rs_formFunc,con_vec, kargs={'A':A,'b':b})
    rs_solver.setJacobian(rs_formJac,A.M.handle)
    rs_solver.setType(PETSc.SNES.Type.VINEWTONRSLS)
    rs_solver.setVariableBounds(lb_vec,ub_vec)

    # MODIFY
    rs_solver.setTolerances(rtol=1e-3,atol=1e-1)
    rs_solver.setFromOptions()
    rs_solver.solve(None,sol_vec)
    rs_solver.reset()



#========================
#  Boundary conditions
#========================
#bc1 = DirichletBC(W.sub(1), Constant(1000.0), 'bottom')
#bc2 = DirichletBC(W.sub(1), Constant(0.0), 'top')
#bc_all = [bc1,bc2]

#========================
#  Solve problem
#========================
#solve(a==L,solution,bcs=bc_all)
#solve(a==L,solution,options_prefix='vms_',solver_parameters=solver_parameters)
virs()
vel,pres = solution.split()

outfile = File('test.pvd')
outfile.write(vel,pres)