import numpy as np
import getfem as gf

def mesh(pts):
    """
    Usage:
    pts = np.array([[0,1,2,3,4],
                    [0,0,0,0,0]])  # along x-axis
    # pts = np.array([[0,0,0,0,0],
    #                 [0,1,2,3,4]])  # along y-axis
    # pts = np.array([[0,1,2,3,4],
    #                 [0,1,2,3,4]])  # along y=x
    # pts = np.array([[0,1],
    #                 [0,0]])  # short, along x-axis
    m = mesh(pts)
    """
    """points"""
    m = gf.Mesh('empty', 2)
    pids = m.add_point(pts)
    print('pids:\n', pids)  # 0 1 2 3 4: 5 points
    """convexes"""
    gt = gf.GeoTrans('GT_PK(1,1)')  # 2-node-line
    ptss = np.zeros((2, 2, pts.shape[1]-1))
    for i in range(pts.shape[1]-1):
        ptss[:,:,i] = pts[:, i:i+2]
    cvids = m.add_convex(gt, ptss)
    print('cvids:\n', cvids)  # 0--1--2--3--4: 4 convexes
    return m

MODES = [0, 1]
if 0 in MODES:
    """
    Default (inspection of gf.asm_linear_elasticity)
    
    all expressions yields the same K:
    [[ 100.    0. -100.    0.    0.    0.    0.    0.    0.    0.]
     [   0.   50.    0.  -50.    0.    0.    0.    0.    0.    0.]
     [-100.    0.  200.    0. -100.    0.    0.    0.    0.    0.]
     [   0.  -50.    0.  100.    0.  -50.    0.    0.    0.    0.]
     [   0.    0. -100.    0.  200.    0. -100.    0.    0.    0.]
     [   0.    0.    0.  -50.    0.  100.    0.  -50.    0.    0.]
     [   0.    0.    0.    0. -100.    0.  200.    0. -100.    0.]
     [   0.    0.    0.    0.    0.  -50.    0.  100.    0.  -50.]
     [   0.    0.    0.    0.    0.    0. -100.    0.  100.    0.]
     [   0.    0.    0.    0.    0.    0.    0.  -50.    0.   50.]]
    This K is strange for a bar or truss because this K has unexpected non-zero stiffness (like 50., -50. or thier sum 100.) in y-direction.
    The K should have non-zero stiffness (like 100., -100. or thier sum 200.) only in x-direction.
    *This K is also strange in having a half stiffness (50. instead of 100.) in y-direction.
    """
    print('MODE 0')
    """expressions"""
    comments = []
    expressions = []
    comments.append(r"from https://lists.nongnu.org/archive/html/getfem-users/2019-09/msg00005.html")
    expressions.append("(Lambda*Trace(Grad_u)*Id(qdim(u)) + Mu*(Grad_u+Grad_u')):Grad_Test_u")
    comments.append(r"from asm_stiffness_matrix_for_linear_elasticity in getfem-5.3/src/getfem/getfem_assembling.h line971")
    expressions.append("((Lambda*Div_Test_u)*Id(meshdim) + (2*Mu)*Sym(Grad_Test_u)):Grad_Test2_u")
    comments.append(r"edited by 'Test_u & Test2_u' -> 'u & Test_u'")
    expressions.append("((Lambda*Div_u)*Id(meshdim) + (2*Mu)*Sym(Grad_u)):Grad_Test_u")
    """constants"""
    E = 100.
    Nu = 0.
    Lambda = E * Nu / ((1. + Nu) * (1. - 2. * Nu))
    Mu = E / (2. * (1. + Nu))
    """assembling"""
    pts = np.array([[0,1,2,3,4],
                    [0,0,0,0,0]])  # along x-axis
    m = mesh(pts)
    mfu = gf.MeshFem(m, 2)  # displacement u, v
    mfu.set_fem(gf.Fem('FEM_PK(1, 1)'))  # 2-node on one convex
    U = np.zeros((mfu.nbdof(),))  # this is for telling getfem the vector length of unknown variable 'u'
    mim = gf.MeshIm(m, gf.Integ('IM_GAUSS1D(3)'))
    order = 2  # 0:scalar 1:vector 2:matrix; 0:potential energy, 1:residual vector, 2:stiffness matrix or tangent term for linear ploblem
    region = -1  # means all
    for expression, comment in zip(expressions, comments):
        K = gf.asm_generic(mim, order, expression, region, # [Model model, [‘Secondary_domain’, ‘name’,]]
                           'u', 1, mfu, U,
                           'Lambda', 0, Lambda,
                           'Mu', 0, Mu)
                           # [string varname, int is_variable[, {MeshFem mf, MeshImd mimd}], value], [‘select_output’, ‘varname1’[, ‘varname2]], …
        np.set_printoptions(suppress=True, precision=1)
        print(comment)
        print(np.array(K.full()))

if 1 in MODES:
    """
    Bar or Truss
    
    yielded K:
    # for first expression: strange K
    [[ 100.    0. -100.    0.    0.    0.    0.    0.    0.    0.]
     [   0.  100.    0. -100.    0.    0.    0.    0.    0.    0.]
     [-100.    0.  200.    0. -100.    0.    0.    0.    0.    0.]
     [   0. -100.    0.  200.    0. -100.    0.    0.    0.    0.]
     [   0.    0. -100.    0.  200.    0. -100.    0.    0.    0.]
     [   0.    0.    0. -100.    0.  200.    0. -100.    0.    0.]
     [   0.    0.    0.    0. -100.    0.  200.    0. -100.    0.]
     [   0.    0.    0.    0.    0. -100.    0.  200.    0. -100.]
     [   0.    0.    0.    0.    0.    0. -100.    0.  100.    0.]
     [   0.    0.    0.    0.    0.    0.    0. -100.    0.  100.]]
     
    # for second expression: expected K
    [[ 100.    0. -100.    0.    0.    0.    0.    0.    0.    0.]
     [   0.    0.    0.    0.    0.    0.    0.    0.    0.    0.]
     [-100.    0.  200.    0. -100.    0.    0.    0.    0.    0.]
     [   0.    0.    0.    0.    0.    0.    0.    0.    0.    0.]
     [   0.    0. -100.    0.  200.    0. -100.    0.    0.    0.]
     [   0.    0.    0.    0.    0.    0.    0.    0.    0.    0.]
     [   0.    0.    0.    0. -100.    0.  200.    0. -100.    0.]
     [   0.    0.    0.    0.    0.    0.    0.    0.    0.    0.]
     [   0.    0.    0.    0.    0.    0. -100.    0.  100.    0.]
     [   0.    0.    0.    0.    0.    0.    0.    0.    0.    0.]]
    """
    print('MODE 1')
    """expressions"""
    comments = []
    expressions = []
    comments.append("yileded K is with transverse stiffness, strange for a bar or truss:")
    expressions.append("A*E*Grad_u:Grad_Test_u")  # : yileded K is with transverse stiffness, strange for a bar or truss
    comments.append("yileded K is without transverse stiffness, expected for a bar or truss:")
    expressions.append("A*E*Grad_u(1,1)*Grad_Test_u(1,1)")  # : yileded K is without transverse stiffness, expected for a bar or truss
    """constants"""
    E = 100.
    A = 1.
    """assembling"""
    pts = np.array([[0,1,2,3,4],
                    [0,0,0,0,0]])  # along x-axis
    m = mesh(pts)
    mfu = gf.MeshFem(m, 2)  # displacement u, v
    mfu.set_fem(gf.Fem('FEM_PK(1, 1)'))  # 2-node on one convex
    U = np.zeros((mfu.nbdof(),))  # this is for telling getfem the vector length of unknown variable 'u'
    mim = gf.MeshIm(m, gf.Integ('IM_GAUSS1D(3)'))
    order = 2  # 0:scalar 1:vector 2:matrix; 0:potential energy, 1:residual vector, 2:stiffness matrix or tangent term for linear ploblem
    region = -1  # means all
    for expression, comment in zip(expressions, comments):
        K = gf.asm_generic(mim, order, expression, region, # [Model model, [‘Secondary_domain’, ‘name’,]]
                           'u', 1, mfu, U,
                           'E', 0, E,
                           'A', 0, A)
                           # [string varname, int is_variable[, {MeshFem mf, MeshImd mimd}], value], [‘select_output’, ‘varname1’[, ‘varname2]], …
        np.set_printoptions(suppress=True, precision=1)
        print(comment)
        print(np.array(K.full()))