# -*- coding: utf-8 -*-
"""
Created on Wed Aug 24 10:49:48 2022

@author: jhchen
"""

import sys
sys.path.append('.')

import numpy as np
from sfepy.discrete.fem import Mesh, FEDomain, Field
from sfepy.discrete import (FieldVariable,Material,Integral,
                            Equation,Equations,Problem)
from sfepy.terms import Term


# load mesh
mesh = Mesh.from_file('./segmented_mouse.mesh')

# domain & regions
domain = FEDomain('domain',mesh)
omega = domain.create_region('Omega','all')
surf = domain.create_region('Surf','vertices of surface','facet')
heart = domain.create_region('Heart','cells of group 2')
lung = domain.create_region('Lung','cells of group 3')
liver = domain.create_region('Liver','cells of group 4')
intestine = domain.create_region('Intestine','cells of group 5')
bone = domain.create_region('Bone','cells of group 10')
main_organs = domain.create_region('Main_Organs',"""r.Heart +c r.Lung +c r.Liver
                                   +c r.Intestine +c r.Bone""")
tissue = domain.create_region('Tissue','r.Omega -c r.Main_Organs')
light_source = domain.create_region('Light_Source','cells of group 7')
# domain.save_regions_as_groups('regions.vtk')

# field
field = Field.from_args('fu',np.float64,'scalar',omega,approx_order=1)

# define unknown variable u and test variable v
u = FieldVariable('u', 'unknown', field)
v = FieldVariable('v', 'test', field, primary_var_name='u')

# define optical parameters to different regions
m = Material('m',values={'D':{'Tissue':0.251,'Heart':0.256,'Lung':0.113,
                              'Liver':0.17,'Intestine':0.251,'Bone':0.108},
                         'mua':{'Tissue':0.039,'Heart':0.2,'Lung':0.663,
                              'Liver':1.2,'Intestine':0.039,'Bone':0.206},
                         # 'q':{'Surf':0.177},
                         's':{'Light_Source':1}})
m2 = Material('m2',values={'q':0.177})

# integral
integral = Integral('i', order=2)

# equation terms
t1 = Term.new('dw_laplace(m.D,v,u)',integral,omega,m=m,v=v,u=u)
t2 = Term.new('dw_dot(m.mua,v,u)',integral,omega,m=m,v=v,u=u)
t3 = Term.new('dw_dot(m2.q,v,u)',integral,surf,m2=m2,v=v,u=u)
t4 = Term.new('dw_integrate(m.s,v)',integral,light_source,m=m,v=v)

eq = Equation('balance',t1+t2+t3-t4)
eqs = Equations([eq])

# problem
pb = Problem('photon flux',equations=eqs)

pb.save_regions_as_groups('regions')

pb.time_update()

pb.update_materials()

variables = pb.get_initial_state()
variables.apply_ebc() # Not necessary for a linear problem...

K = pb.equations.eval_tangent_matrices(variables(), pb.mtx_a)
print(K)