from firedrake import *
from create_grid import *
import numpy as np

(Lx, Ly, Lc, Lp, dx, dy, n_Lp, Nn) = create_grid();

mesh_data = np.loadtxt("sweVbuoy2D_mesh.in")

Nn = mesh_data.max().astype(int)

xcoords = mesh_data[:,[0,2,4,6]]
ycoords = mesh_data[:,[1,3,5,7]]
coords = np.empty((Nn, 2), dtype=float)

cells = mesh_data[:,[8,9,10,11]].astype(int) - 1
coords[cells] = mesh_data[:, :8].reshape(-1, 4, 2)

from firedrake.mesh import _from_cell_list
plex = _from_cell_list(2, cells, coords)

# Apply boundary IDs and mark boundary facets
plex.createLabel("boundary_ids")
plex.markBoundaryFaces("boundary_faces")
coordinates = plex.getCoordinates()
coord_sec = plex.getCoordinateSection()
if plex.getStratumSize("boundary_faces", 1) > 0:
    boundary_faces = plex.getStratumIS("boundary_faces", 1).getIndices()
    eps = 1e-5
    for face in boundary_faces:
        face_coords = plex.vecGetClosure(coord_sec, coordinates, face)
        # face_coords has 4 entries x0, y0, x1, y1.
        if abs(face_coords[1]) < eps and abs(face_coords[3]) < eps:
            plex.setLabelValue("boundary_ids", face, 1)

mesh = Mesh(plex)

V = FunctionSpace(mesh, "CG", 1)
f = Function(V)

f_file = File("f.pvd")
f_file << f