#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import cxxfem as fem
import gmsh


def build_square(cx = 0.0, cy = 0.0,
                 Lx = 1.0, Ly = 1.0,
                 nx = 2, ny = 2,
                 structured = True, recombine = True):

    gmsh.model.add("square")
    lc = 1.0 # no impact
    gmsh.model.geo.addPoint(cx, cy, 0, lc, 1)
    gmsh.model.geo.addPoint(cx + Lx, cy, 0, lc, 2)
    gmsh.model.geo.addPoint(cx + Lx, cy + Ly, 0, lc, 3)
    gmsh.model.geo.addPoint(cx, cy + Ly, 0, lc, 4)

    gmsh.model.geo.addLine(1, 2, 1)
    gmsh.model.geo.addLine(2, 3, 2)
    gmsh.model.geo.addLine(4, 3, 3)
    gmsh.model.geo.addLine(1, 4, 4)

    gmsh.model.geo.addCurveLoop([1, 2, -3, -4], 1)

    gmsh.model.geo.addPlaneSurface([1], 1)

    gmsh.model.geo.synchronize() # sinon warning "unknown entity..." lors du addPhysicalGroup

    # mesh lines
    gmsh.model.mesh.setTransfiniteCurve(1, nx + 1)
    gmsh.model.mesh.setTransfiniteCurve(3, nx + 1)
    gmsh.model.mesh.setTransfiniteCurve(2, ny + 1)
    gmsh.model.mesh.setTransfiniteCurve(4, ny + 1)

    # recombine
    if structured:
        gmsh.model.mesh.setTransfiniteSurface(1)
    if recombine:
        gmsh.model.mesh.setRecombine(2, 1) # quads

    # physical groups - corners
    pt1 = gmsh.model.addPhysicalGroup(0, [1])
    pt2 = gmsh.model.addPhysicalGroup(0, [2])
    pt3 = gmsh.model.addPhysicalGroup(0, [3])
    pt4 = gmsh.model.addPhysicalGroup(0, [4])
    gmsh.model.setPhysicalName(0, pt1, "pt1")
    gmsh.model.setPhysicalName(0, pt2, "pt2")
    gmsh.model.setPhysicalName(0, pt3, "pt3")
    gmsh.model.setPhysicalName(0, pt4, "pt4")

    # physical groups - lines
    bottom = gmsh.model.addPhysicalGroup(1, [1], 1000)
    right = gmsh.model.addPhysicalGroup(1, [2])
    top = gmsh.model.addPhysicalGroup(1, [3])
    left = gmsh.model.addPhysicalGroup(1, [4])
    gmsh.model.setPhysicalName(1, bottom, "bottom")
    gmsh.model.setPhysicalName(1, right, "right")
    gmsh.model.setPhysicalName(1, top, "top")
    gmsh.model.setPhysicalName(1, left, "left")

    # physical groups - surfaces
    interior = gmsh.model.addPhysicalGroup(2, [1])
    gmsh.model.setPhysicalName(2, interior, "interior")

    # finalize and mesh the geometry
    gmsh.model.mesh.generate(2)

    # gmsh.write("square.msh")


if __name__ == "__main__":

    pbl = fem.Problem()    #gmsh.initialize()
    build_square(0.0, 0.0, 10.0, 1.0, 80, 8, True, True)
    # build_square(0.0, 0.0, 10.0, 1.0, 2, 1, True, True)

    m = fem.Medium(pbl, "interior", 100.0, 0.3)

    d1 = fem.Dirichlet(pbl, "left", "X", 0.0)
    d2 = fem.Dirichlet(pbl, "left", "Y", 0.0)
    # d3 = fem.Dirichlet(pbl, "pt3", "Y", -0.1) # option 1
    n3 = fem.Neumann(pbl, "top", "Y", -1e-3)    # option 2

    d4 = fem.Dirichlet(pbl, "interior", "Z", 0.0)


    solver = fem.Solver(pbl)
    solver.solve()

    post = fem.Post(solver)
    post.build_views()
    post.show_views( [ "stress_tensor", "force_vector" ] )
    post.write()
    post.deform(5)
     
    forces = post.probe("force_vector", "left")
    print(f'Rx = {sum(forces[::3])} N')
    print(f'Ry = {sum(forces[1::3])} N')
    disp = post.probe("Y", "pt2")
    print(f'tip displacement = {disp[0]} mm')
    
    args = fem.parseargs()
    if not args.nogui:    
        post.view()