import sfepy

from sfepy.base.base import Struct
from sfepy.discrete import (FieldVariable, Material, Integral, Integrals,
                            Equation, Equations, Problem)
from sfepy.discrete.fem import FEDomain, Field,Mesh
from sfepy.terms import Term
from sfepy.discrete.conditions import Conditions, EssentialBC
from sfepy.mechanics.matcoefs import stiffness_from_youngpoisson,lame_from_youngpoisson
from sfepy.mechanics.tensors import get_von_mises_stress


import numpy as np
from IPython.display import Image

print sfepy.__version__

Image('/Users/Sukhbinder/Documents/sfepy-simple_tutorial/simple_plate_geometry.png')

# Load the Mesh
mesh=Mesh.from_file('/Users/Sukhbinder/sqarePlate.vtk')

# Setup the Domain and Regions
domain = FEDomain('domain', mesh)

# We call the entire domain omega
omega = domain.create_region('Omega', 'all')

#  We define bottom
bottom = domain.create_region('Bottom','vertices in (y < 0.001)','facet')

# and Top
top = domain.create_region('Top','vertices in (y > 0.99)','facet')

young=1.0
poisson=0.0
density=1.0
loadval=1.0

lam,mu = lame_from_youngpoisson(young,poisson,plane='stress')
mtx_d = stiffness_from_youngpoisson(2, young, poisson, plane='stress' )

m = Material('m',lam=lam,mu=mu, D=mtx_d, rho=density)

f = Material('f', val=loadval)

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

# Define the integral type Volume/Surface and quadrature rule
integral = Integral('i', order=2)

# The weak formulation of the linear elastic problem.
t1 = Term.new('dw_lin_elastic_iso(m.lam, m.mu,v, u)', integral, omega, m=m, v=v, u=u)
t2 = Term.new('dw_surface_ltr(f.val, v)', integral, top, f=f, v=v)
eq1 = Equation('balance_of_forces', t1-t2)

lhs_eqs = Equations([eq1])



pb = Problem('plate', equations=lhs_eqs)

fixed_b = EssentialBC('FixedB', bottom, {'u.all' : 0.0})

pb.time_update(ebcs=Conditions([fixed_b])) 

state = pb.solve()

print pb

out = state.create_output_dict()
pb.save_state('plate.vtk', out=out)

# Check outout
print out

# Check solve state
print state

from sfepy.postprocess.viewer import Viewer
view = Viewer('plate.vtk')


view()

Image('/Users/Sukhbinder/Documents/sfepy-simple_tutorial/sfepy_simple_plate_simple.png')

view(vector_mode='warp_norm', rel_scaling=.1,is_scalar_bar=True, is_wireframe=True ,\
     rel_text_width=.1)

Image('/Users/Sukhbinder/Documents/sfepy-simple_tutorial/sfepy_result_simple_plate.png')

pb.save_regions_as_groups("plateregions")
view = Viewer('plateregions.vtk')
view(is_scalar_bar=True, is_wireframe=True ,rel_text_width=.2)

Image('/Users/Sukhbinder/Documents/sfepy-simple_tutorial/sfepy_regions_plot.png')


strain = pb.evaluate('ev_cauchy_strain.2.Omega(u)', mode='el_avg')
stress = pb.evaluate('ev_cauchy_stress.2.Omega(m.D, u)', mode='el_avg')

# Strain
out['cauchy_strain'] = Struct(name='output_data', mode='cell',data=strain, dofs=None)

# Stress
out['cauchy_stress'] = Struct(name='output_data', mode='cell',data=stress, dofs=None)

# Von mises Stress
vms = get_von_mises_stress(stress.squeeze())
vms.shape = (vms.shape[0], 1, 1, 1)
out['von_mises_stress'] = Struct(name='output_data', mode='cell',data=vms, dofs=None)

print out 

pb.save_state('plate_stress_strain.vtk', out=out,extend=True)

view = Viewer('plate_stress_strain.vtk')
view(vector_mode='warp_norm', rel_scaling=.1,is_scalar_bar=True, is_wireframe=True ,rel_text_width=.1)

Image('/Users/Sukhbinder/Documents/sfepy-simple_tutorial/Sfepy_Stress_Strain.png')

u=state()

print u.reshape((121,2))[:100,1]

print u.reshape((121,2))[101:121,1]

print state

from IPython.core.display import HTML
def css_styling():
    styles = open("custom.css", "r").read()
    return HTML(styles)
css_styling()

t1 = Term.new('dw_lin_elastic(m.D, v, u)', integral, omega, m=m, v=v, u=u)
t2 = Term.new('dw_volume_dot(m.rho, v, u)', integral, omega, m=m, v=v, u=u)
eq1 = Equation('stiffness', t1)
eq2 = Equation('mass', t2)
lhs_eqs = Equations([eq1, eq2])
pbm = Problem('massstiff', equations=lhs_eqs)
pbm.time_update()
pbm.update_materials()


# Assemble stiffness and mass matrices.                                                         
mtx_k = eq1.evaluate(mode='weak', dw_mode='matrix', asm_obj=pbm.mtx_a)
mtx_m = mtx_k.copy()
mtx_m.data[:] = 0.0
mtx_m = eq2.evaluate(mode='weak', dw_mode='matrix', asm_obj=mtx_m)
print mtx_k
print mtx_m
print mesh.coors
print mesh.conns