from grains.dic import DIC
from grains.geometry import TriMesh
x_grid, y_grid = np.mgrid[-1:1:100j, 1:2:50j]
exact_solution = lambda x, y: 1 - (x + y**2) * np.sign(x)
grid = DIC(exact_solution(x_grid, y_grid).T, exact_solution(x_grid,y_grid).T)
grid.set_transformation((-1, 2), 50)

# The finite element solution is obtained at the nodes (upper right figure) of the mesh.
# Here, we assumed that the nodes are sampled from a uniformly random distribution on
# :math:`[-1,1] \times [1, 2]`.

n_nodes = 100  # modify this to see how good the interpolated solution is

# For the sake of this example, we also assume that the nodal values are "exact", i.e.
# they admit the function values at the nodes. In reality, of course, this will
# not be the case, but this allows us to investigate the effect moving from the FE mesh
# to the DIC grid.

mesh = TriMesh(*TriMesh.sample_mesh(1, n_nodes))
fe_values = exact_solution(mesh.vertices[:, 0], mesh.vertices[:, 1])
mesh.associate_field(fe_values)
interpolated = grid.project_onto_grid(mesh.vertices, fe_values, method='linear')

# The FE field available at the nodes are interpolated at the DIC grid points, as shown
# in the lower right figure. Note that no extrapolation is performed, so values are not
# computed at the grid points lying outside the convex hull of the finite element nodes.

ax = []
ax.append(plt.subplot(221))
plt.plot(x_grid, y_grid, 'k.', ms=1)
plt.title('DIC grid')
ax.append(plt.subplot(222))
plt.plot(mesh.vertices[:, 0], mesh.vertices[:, 1], 'k.', ms=1)
plt.title('FE nodes')
ax.append(plt.subplot(223))
mesh.plot_field(0, ax=ax[-1])
plt.title('FE field')
ax.append(plt.subplot(224))
plt.imshow(interpolated, extent=(-1, 1, 1, 2), origin='lower', vmin=-4, vmax = 5)
plt.title('FE field interpolated at the DIC grid')
plt.tight_layout()
for a in ax:
    a.set_aspect('equal', 'box')
plt.show()