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)

# 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
mesh = TriMesh(*TriMesh.sample_mesh(1, n_nodes))

# For the sake of this example, we also assume that the nodal values are "exact", i.e. they
# are 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.

grid.set_transformation((-1, 2), 50)
fe_values = exact_solution(mesh.vertices[:, 0], mesh.vertices[:, 1])
interpolated = grid.project_onto_mesh(mesh.vertices, 'nearest')

# The FE solution available at the nodes are interpolated at the DIC grid points, as shown
# in the lower right figure. Interpolation with continuous functions cannot resolve well the
# discontinuity present in the "exact solution", i.e. in the measurement data. A discontinuous
# manufactured solution was intentionally prepared to illustrate this.

ax = []
ax.append(plt.subplot(221))
grid.plot_physicalgrid(ax[0])
plt.title('DIC grid')
ax.append(plt.subplot(222))
mesh.plot('k.', ax=ax[1], markersize=1)
plt.title('FE nodes')
ax.append(plt.subplot(223))
plt.imshow(exact_solution(x_grid, y_grid).T, extent=(-1, 1, 1, 2), origin='lower',
           vmin=-4, vmax = 5)
plt.title('Exact solution on the DIC grid')
ax.append(plt.subplot(224))
interpolated[0][np.isnan(interpolated[0])] = 0
mesh.associate_field(interpolated[0])
mesh.plot_field(0, show_mesh=False, ax=ax[3])
plt.title('FE solution interpolated at the DIC grid')
plt.show()