#!/usr/bin/env python3

import vtk
from vtk.util import numpy_support

import numpy as np

import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable

plt.rc("text", usetex=True)

reader = vtk.vtkXMLPolyDataReader()

reader.SetFileName("postProcessing/surfacesf/200/centerPlane.vtp")
reader.Update()

data = reader.GetOutput()

triangle_filter = vtk.vtkTriangleFilter()
triangle_filter.SetInputData(data)
triangle_filter.Update()
data = triangle_filter.GetOutput()

vtk_points = data.GetPoints().GetData()
points = numpy_support.vtk_to_numpy(vtk_points)
x = points[:, 0]
y = points[:, 1]

vectors_name = "U"
vtk_vectors = data.GetPointData().GetArray(vectors_name)

U = numpy_support.vtk_to_numpy(vtk_vectors)
magU = np.linalg.norm(U, axis=1)

vtk_cells = data.GetPolys().GetData()
cells_raw = numpy_support.vtk_to_numpy(vtk_cells)
triangles = cells_raw.reshape(-1, 4)[:, 1:]

fig, ax = plt.subplots(figsize=(8, 6))

contour_filled = ax.tricontourf(x, y, triangles, magU, levels=16, cmap="turbo")

ax.tricontour(
    x, y, triangles, magU, levels=16, colors="white", linewidths=0.5, alpha=0.5
)

ax.set_xticks([-2, 0, 2, 4, 6, 8, 10, 12])
ax.tick_params(labelsize=18)

divider = make_axes_locatable(ax)
cax = divider.append_axes("right", size="5%", pad=0.1)

cbar = fig.colorbar(contour_filled, cax=cax)
ax.set_aspect("equal")

ax.grid(ls=":")

for f in ["svg", "pdf", "eps", "png"]:
    plt.savefig("velocity.{}".format(f), dpi=300, bbox_inches="tight", pad_inches=0.5)

plt.show()
plt.close(fig)