# -*- coding: utf-8 -*-
"""
Created on Tue Feb 24 18:02:29 2015

@author: ronghaiwu
"""

#!/usr/bin/env python
import fipy as fp 
import numpy as np
import sys
sys.path.append('.')
from sfepy import data_dir 
from sfepy.base.base import IndexedStruct
from sfepy.discrete import (FieldVariable, Material, Integral, Function, Equation, Equations, Problem, Functions)
from sfepy.discrete.fem import Mesh, FEDomain, Field
from sfepy.terms import Term
from sfepy.discrete.conditions import Conditions, PeriodicBC
from sfepy.solvers.ls import ScipyDirect
from sfepy.solvers.nls import Newton
import sfepy.discrete.fem.periodic as per

####################################################################################################################
#this is the funtion to generate mesh for fem, the fem_mesh should be the same as fvm_mesh, especially the order.
####################################################################################################################  
def fem_mesh_generation(nx, ny, dx, dy, path=data_dir, name='2D_mesh.mesh'):
    outfile = open(path + name,'w')
    outfile.write('MeshVersionFormatted 1\t\t\t\t\t\t\nDimension\t2\t\t\t\t\t\t\nVertices\t\t\t\t\t\t\n')
    outfile.write('%4d'%((nx+1)*(ny+1)))
    outfile.write('\n')
    for j in range(ny+1):
        for i in range(nx+1):
            outfile.write('\t\t')
            outfile.write('%8f'%(i*dx))
            outfile.write('\t\t')
            outfile.write('%8f'%(j*dy))
            outfile.write('\t\t')
            outfile.write('%8i'%(0))
            outfile.write('\n')
    outfile.write('\t\t\t\t\t\t\n\t\t\t\t\t\t\nQuadrilaterals\t\t\t\t\t\t\n')
    outfile.write('%4d'%(nx*ny))
    outfile.write('\n')
    for q in range(1,nx*ny + ny+1):
        if (q % (nx+1)) != 0:
            outfile.write('%4d'%(q))
            outfile.write('\t')
            outfile.write('%4d'%(q+1))
            outfile.write('\t')
            outfile.write('%4d'%(q+nx+2))
            outfile.write('\t')
            outfile.write('%4d'%(q+nx+1))
            outfile.write('\t')
            outfile.write('%4d'%(2))
            outfile.write('\n')
    outfile.close()
    return path + name

#######################################################################################
#The prestress value is already calculated at each element. 
#This function is to copy value to all the four verties in each element 
#######################################################################################    
def prestress_function(ts, coors, mode, term=None, **kwargs):
    if mode != 'qp': return

    _, weights = term.integral.get_qp('2_4')
    n_qp = weights.shape[0]

    val = np.repeat(prestress_value, n_qp, axis=0)
    print coors.shape
    print val.shape

    return {'function' : val}

#########################################################################################
#The stiffness tensor is already calculated at each element
#This function is to copy value to all the four verties in each element
##########################################################################################
def D_function(ts, coors, mode, term=None, **kwargs):
    if mode != 'qp': return

    _, weights = term.integral.get_qp('2_4')
    n_qp = weights.shape[0]
    val = np.repeat(stiffness_gene, n_qp, axis=0)
    return {'function' : val} 



dx = 1.
dy = dx
nx = 128
ny = nx

#define a indicator(scalar field) to indicate where there is a inclusion
#one inclusion is prescribed at mesh center
fvm_mesh = fp.Grid2D(dx=dx, dy=dy, nx=nx, ny=ny)
fvm_x = fvm_mesh.cellCenters[0]
fvm_y = fvm_mesh.cellCenters[1]
c = fp.CellVariable(mesh=fvm_mesh, hasOld=1)
c.setValue(1., where = (fvm_x > 0.5*nx*dx-4.*dx) & (fvm_x < 0.5*nx*dx+4.*dx) & (fvm_y > 0.5*ny*dy-4.*dy) & (fvm_y < 0.5*ny*dy+4.*dy) )

#define stiffness tensor and eigenstress
stiffness_gene = np.zeros((nx*ny,3,3))
C_incl_11 = 193.
C_incl_12 = 131.
C_incl_44 = 97.
stiffness_gene[:,0,0] = C_incl_11
stiffness_gene[:,0,1] = C_incl_12
stiffness_gene[:,1,0] = C_incl_12
stiffness_gene[:,1,1] = C_incl_11
stiffness_gene[:,2,2] = C_incl_44
strain00_11_value = 0.01
strain00_22_value = strain00_11_value
strain00_12_value = 0.
strain00_21_value = strain00_12_value
stress00 = np.zeros((nx*ny,3,))
for i in range(nx*ny):
    stress00[i] = np.dot(stiffness_gene[i], np.array([strain00_11_value, strain00_22_value, 2.*strain00_12_value]))

 
fem_mesh = Mesh.from_file(fem_mesh_generation(nx, ny, dx, dy))
domain = FEDomain('domain', fem_mesh)

omega = domain.create_region('Omega', 'all')
min_x, max_x = domain.get_mesh_bounding_box()[:,0]
eps = 1e-8 * (max_x - min_x)
min_y, max_y = domain.get_mesh_bounding_box()[:,1]  
gamma_left  = domain.create_region('Gamma_left', 'vertices in x < %.10f' % (min_x + eps), 'facet')
gamma_right = domain.create_region('Gamma_right', 'vertices in x > %.10f' % (max_x - eps), 'facet')     
gamma_bottom = domain.create_region('Gamma_bottom', 'vertices in y < %.10f' % (min_y + eps), 'facet')
gamma_top = domain.create_region('Gamma_top', 'vertices in y > %.10f' % (max_y - eps), 'facet')    

field = Field.from_args('displacement', np.float64, 'vector', omega, approx_order=1)
u = FieldVariable('u', 'unknown', field)
v = FieldVariable('v', 'test', field, primary_var_name='u')

m = Material('m', function=D_function) 
traction_top = Material('traction_top', val=[[1.0], [0.]])
traction_bottom = Material('traction_bottom', val=[[-1.0], [0.]])
prestress_value = c.value[:, None] * stress00
prestress_value.shape = (-1, 3, 1)
prestress = Material('prestress', function=prestress_function)

integral = Integral('i', order=3)
t1 = Term.new('dw_lin_elastic(m.function, v, u )', integral, omega, m=m, v=v, u=u) 
t2 = Term.new('dw_lin_prestress(prestress.function, v )',integral, omega, prestress=prestress, v=v)
t3 = Term.new('dw_tl_surface_traction(traction_top.val, v)', integral, gamma_top, traction_top=traction_top, v=v)
t4 = Term.new('dw_tl_surface_traction(traction_bottom.val, v)', integral, gamma_bottom, traction_bottom=traction_bottom, v=v)
fem_eq = Equation('balance', t1 - t2 + t3 + t4 )
eqs = Equations([fem_eq])

match_y_line = Function('match_y_line', per.match_y_line)
periodic_x = PeriodicBC('periodic_x', [gamma_left, gamma_right], {'u.all' : 'u.all'}, match='match_y_line')     
functions = Functions([match_y_line]) 
ls = ScipyDirect({})
nls_status = IndexedStruct()
nls = Newton({}, lin_solver=ls, status=nls_status)
pb = Problem('elasticity', equations=eqs, nls=nls, ls=ls)
pb.save_regions_as_groups('regions')
pb.time_update(ebcs=Conditions([periodic_x]), functions=functions) 
vec = pb.solve()