#
#    CAUTION
#
# While possible, solving structural mechanics problems
# for truss-type or frame-type structures in GetFEM++
# does not look very natural. For such problems one often
# uses precomputed stiffness matrices, making the use
# of variational formulation, shape functions, integration
# rules fully implicit. In GetFEM++ in turn, such elements
# of FEM algorithm are exposed in order to parameterize 
# the algorithm. By adding suitable GetFEM++ bricks, solving
# trusses and frames could be streamlined, for now however,
# one have to "twist" a little the API. 
#
#
"""Solution of a simple truss structure:
         |       |
         V P     V P
         1-------3
        / \     / \
       /   \   /   \           ^ v
      /     \ /     \          |
     0-------2-------4         ----> u
                    
     all bars are the same:   stiffness EA = 1
                              length    a  = 1
     Load: P = 1
     Constraints  (displacements: u,v)   v_0 = 0
                                         u_2 = 0
                                         v_4 = 0
 """ 
import numpy
import getfem
import math
a = 1.0;
h = a*math.sqrt(3)/2

# array of node coordinates -- each row corresponds to a node 
xy = numpy.array([[0,0],
                      [0.5*a, h],
                      [a, 0.0],
                      [1.5*a, h],
                      [2*a, 0.0]], dtype=numpy.float32)

# The pts array defined below is 3-dimensional array
# The valuses inside the inner-most square bracket
# are related to the first index, and so on. Thus
# 
#   pts_ijk    i -- related to convex index
#              j -- related to coordinate index
#              k -- related to the vertex index in a convex
#
# This array is created to add all convexes at once.
pts = numpy.array([[xy[0],xy[1]],
                   [xy[0],xy[2]],
                   [xy[1],xy[2]],
                   [xy[1],xy[3]],
                   [xy[2],xy[3]],
                   [xy[2],xy[4]],
                   [xy[3],xy[4]]
                  ], dtype=numpy.float32);
print pts.shape
# The method add_convex used below expects that the last index
# of the pts array is related to convex number -- thus the need
# to swap axes
pts=pts.swapaxes(0,2)
mesh = getfem.Mesh('empty', 2)
mesh.add_convex(getfem.GeoTrans('GT_PK(1,1)'), pts)

fem = getfem.MeshFem(mesh);
# We solve planar truss thus our nodal unknown U is two dimensional vector.
fem.set_qdim(2);
# We use linear shape functions (the most simple truss element).
fem.set_classical_fem(1);

mim = getfem.MeshIm(mesh, getfem.Integ('IM_GAUSS1D(2)'))

model = getfem.Model('real')

EA = 1.0

model.add_fem_variable('U', fem)


# Here is the twist -- by setting appropriately elasticity constants
# we introduce here the relation between internal forces and displacements
# F = EA u
model.add_initialized_data('lambda', EA )
model.add_initialized_data('mu', 0);
model.add_isotropic_linearized_elasticity_brick(mim, 'U', 'lambda', 'mu') 


# Dirichlet boundary conditions
# ------------------------------
yfix = 1 # region number for y-support nodes
xfix = 2 # region number for x-support nodes

#yfix_cvid -- is the array specifying y-supported region
#             each column corresponds to a face of a convex
#             (as we have 1D element the face reduces to node)
#             first row is convex index, second row is face index
#  CAUTION: mid the face numbering rule for simplex elements:
#           i-th face is opposite i-th node
#
yfix_cvid = numpy.array([[1,5],[1,0]]);
xfix_cvid = numpy.array([[1],[0]])

mesh.set_region(yfix, yfix_cvid);
mesh.set_region(xfix, xfix_cvid);

# Arrays Hy, and Hx are used to pick the direction of the supports
model.add_initialized_data('Hy', [0,0,0,1]);
model.add_initialized_data('Ry', [0,0]);
model.add_initialized_data('Hx', [1,0,0,0]);
model.add_initialized_data('Rx', [0,0]);

model.add_generalized_Dirichlet_condition_with_penalization(mim, 'U', 1e20,
yfix, 'Ry', 'Hy');
model.add_generalized_Dirichlet_condition_with_penalization(mim, 'U', 1e20,
xfix, 'Rx', 'Hx');

# Assembling nodal loads
#-----------------------
# ledof -- an array of DOFs numbers in convex of index 3
ledof, dummy = fem.basic_dof_from_cvid([3])
nodal_load = numpy.zeros(fem.nbdof(), dtype=numpy.float)
nodal_load[ledof[1]] = -1
nodal_load[ledof[3]] = -1
model.add_explicit_rhs('U', nodal_load)


# Solving and exporting the solution
# ----------------------------------
model.solve()

U = model.variable('U')
print U
fem.export_to_vtk("bridge.vtk", 'ascii', U, 'displacement')