# -*- 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, EssentialBC
from sfepy.solvers.ls import ScipyDirect
from sfepy.solvers.nls import Newton
import sfepy.discrete.fem.periodic as per
from sfepy.base.base import Struct
from sfepy.discrete.variables import Variables
from sfepy.discrete.probes import Probe
import matplotlib.font_manager as fm
import mayavi.mlab as mlab
from sfepy.discrete.probes import LineProbe
import matplotlib.pyplot as plt

####################################################################################################################
#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 = 129
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)
R = 0.5*1./10.*nx*dx
hole = ( np.sqrt( (fvm_x-0.5*nx*dx)**2 + (fvm_y-0.5*ny*dy)**2 ) <= R )
c.setValue(1., where = hole )

#define stiffness tensor and eigenstress
stiffness_gene = np.zeros((nx*ny,3,3))
C_matrix_11 = 1400.
C_matrix_12 = 600.
C_matrix_44 = 400.
stiffness_gene[:,0,0] = C_matrix_11
stiffness_gene[:,0,1] = C_matrix_12
stiffness_gene[:,1,0] = C_matrix_12
stiffness_gene[:,1,1] = C_matrix_11
stiffness_gene[:,2,2] = C_matrix_44
strain00_11_value = 0.1
strain00_22_value = strain00_11_value
strain00_12_value = 0.0
strain00_21_value = strain00_12_value
stress00 = np.zeros((nx*ny,3,))
for i in range(nx*ny):
    stiffness_gene[i] = stiffness_gene[i]*(1.00001-c.value[i])
    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-2
min_y, max_y = domain.get_mesh_bounding_box()[:,1]  

lef = domain.create_region('Lef', 'vertices in x < %f' % (min_x + eps), 'facet')
rig = domain.create_region('Rig', 'vertices in x > %f' % (max_x - eps), 'facet')     
bot = domain.create_region('Bot', 'vertices in y < %f' % (min_y + eps), 'facet')
top = domain.create_region('Top', 'vertices in y > %f' % (max_y - eps), 'facet')        
lef_corner = domain.create_region('Lef_corner', 'vertices in (y < 0.01) & (x < 0.01)', 'vertex')    
rig_corner = domain.create_region('Rig_corner', 'vertices in (y < 0.01) & (x > 128.9)', 'vertex')  


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_lef = Material('traction_lef', val=[[1.], [0.], [0.]]) # sigma11, sigma22, sigma12.
traction_rig = Material('traction_rig', val=[[1.], [0.], [0.]]) # sigma11, sigma22, sigma12.
prestress_value = c.value[:, None] * stress00
prestress_value.shape = (-1, 3, 1)
prestress = Material('prestress', function=prestress_function)

integral = Integral('i', order=2)
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_surface_ltr(traction_lef.val, v)', integral, lef, traction_lef=traction_lef, v=v)
t4 = Term.new('dw_surface_ltr(traction_rig.val, v)', integral, rig, traction_rig=traction_rig, v=v)
fem_eq = Equation('balance', t1 - t2 - t3 - t4 )
eqs = Equations([fem_eq])

fix_u_lef_corner = EssentialBC('fix_u_lef_corner', lef_corner, {'u.all' : 0.0})
fix_u_rig_corner = EssentialBC('fix_u_rig_corner', rig_corner, {'u.1' : 0.0})
match_x_line = Function('match_x_line', per.match_x_line)
periodic_y = PeriodicBC('periodic_y', [top, bot], {'u.all' : 'u.all'}, match='match_x_line')     
functions_y = Functions([match_x_line]) 
match_y_line = Function('match_y_line', per.match_y_line)
periodic_x = PeriodicBC('periodic_x', [lef, rig], {'u.all' : 'u.all'}, match='match_y_line')     
functions_x = Functions([match_y_line]) 
functions = Functions([match_x_line, match_y_line]) 

ls = ScipyDirect({})
nls_status = IndexedStruct()
nls = Newton({'i_max' : 2,'problem' : 'nonlinear'}, lin_solver=ls, status=nls_status)
pb = Problem('elasticity', equations=eqs, nls=nls, ls=ls)

fix_u_lef_corner = EssentialBC('fix_u_lef_corner', lef_corner, {'u.all' : 0.0})
fix_u_rig_corner = EssentialBC('fix_u_rig_corner', rig_corner, {'u.1' : 0.0})
#stress free    
pb.time_update(ebcs=Conditions([fix_u_lef_corner, fix_u_rig_corner])) 

#top-bot periodic
#pb.time_update(ebcs=Conditions([fix_u_lef_corner, fix_u_rig_corner]), epbcs=Conditions([periodic_y]), functions=functions_y) 

pb.update_materials()
vec = pb.solve()
state = pb.create_state()
out = state.create_output_dict()   
    
def stress_strain(out, pb, state, extend=False):    
    sfepy_strain = pb.evaluate('ev_cauchy_strain.2.Omega(u)', mode='el_avg')  
    sfepy_stress = pb.evaluate('ev_cauchy_stress.2.Omega(m.function, u)', mode='el_avg')
    out['cauchy_strain'] = Struct(name='output_data', mode='cell', data=sfepy_strain, dofs=None)
    out['cauchy_stress'] = Struct(name='output_data', mode='cell', data=sfepy_stress, dofs=None)
    return out
   

def gen_lines(pbroblem):    
    ps0 = [[0.0, 0.0],   [ 0.0, 0.0]]
    ps1 = [[129.0, 0.0], [ 0.0, 129.0]]
    n_point = -10
    labels = ['%s -> %s' % (p0, p1) for p0, p1 in zip(ps0, ps1)]
    probes = []
    for ip in xrange(len(ps0)):
        p0, p1 = ps0[ip], ps1[ip]
        probes.append(LineProbe(p0, p1, n_point))
        return probes, labels
def probe_hook(data, probe, label, problem):
    def get_it(name, var_name):
        var = problem.create_variables([var_name])[var_name]
        var.set_data(data[name].data)
        pars, vals = Probe.probe(var)
        vals = vals.squeeze()
        return pars, vals
    results = {}
    results['cauchy_strain'] = get_it('cauchy_strain', 's')
    results['cauchy_stress'] = get_it('cauchy_stress', 's')
    fig = plt.figure()
    plt.clf()
    fig.subplots_adjust(hspace=0.4)
   
    sym_indices = ['11', '22', '12']

    plt.subplot(211)
    pars, vals = results['cauchy_strain']
    for ic in range(vals.shape[1]):
        plt.plot(pars, vals[:,ic], label=r'$e_{%s}$' % sym_indices[ic], lw=1, ls='-', marker='+', ms=3)
    plt.ylabel('Cauchy strain')
    plt.xlabel('probe %s' % label, fontsize=8)
    plt.legend(loc='best', prop=fm.FontProperties(size=8))
    
    plt.subplot(212)
    pars, vals = results['cauchy_stress']
    for ic in range(vals.shape[1]):
        plt.plot(pars, vals[:,ic], label=r'$\sigma_{%s}$' % sym_indices[ic], lw=1, ls='-', marker='+', ms=3)
    plt.ylabel('Cauchy stress')
    plt.xlabel('probe %s' % label, fontsize=8)
    plt.legend(loc='best', prop=fm.FontProperties(size=8))
    
    return plt.gcf(), results
data = stress_strain(out, pb, state)
probes, labels = gen_lines(pb)
probe_hook(data, probes, labels, pb)