'''
Here we test a Hertzian contact on a generalized Maxwell material using Tamaas.
Contact with rough surfaces needs to be tested.
'''
import tamaas as tm
import time
import numpy as np


# Set-up of the model
L = 2
Radius = 0.5
S = L**2 

# discretization
n = m = 300
x = np.linspace(0, L, n, endpoint=False, dtype=tm.dtype)
y = np.linspace(0, L, m, endpoint=False, dtype=tm.dtype)
xx, yy = np.meshgrid(x, y, indexing="ij")
# Define the surface
surface = surface = -((xx - L / 2) ** 2 + (yy - L / 2) ** 2) / (2 * Radius)
# Create the model
model = tm.Model(tm.model_type.basic_2d, [L, L], [n, m])

# Defining the elastic branch (i.e. the behavior at t = ∞)
model.E = 3
model.nu = 0.5

# Characteristic times of the relaxation function
times = [0.1, 1]

# Shear moduli for each branch of the model
shear_moduli = [2.75, 2.75]

t0 = 0
t1 = 1
time_steps = 50
# Time step
Δt = (t1 - t0) / time_steps

# Applied load
W = 1.0
load = W / S

# Solver instanciation
solver = tm.MaxwellViscoelastic(model, surface, 1e-10,
                                time_step=Δt,
                                shear_moduli=shear_moduli,
                                characteristic_times=times)

# Solve one timestep with given load
start = time.perf_counter()
solver.solve(load)
end = time.perf_counter()
print(f'Simulation time for one step: {end - start} seconds')

# plot like ub Multi_branches_generalized_Maxwell.py
import matplotlib.pyplot as plt
displacement = model.displacement[:]
pressure = model.traction[:]
plt.figure(figsize=(12, 5))
plt.subplot(1, 2, 1)
plt.imshow(displacement, extent=(0, L, 0, L), origin='lower')
plt.title('Displacement field')
plt.colorbar()
plt.subplot(1, 2, 2)
plt.imshow(pressure, extent=(0, L, 0, L), origin='lower')
plt.title('Pressure field')
plt.colorbar()
plt.show()