import tamaas as tm
import sys
import matplotlib.pyplot as plt

# Initialize threads and fftw
tm.initialize()

# Surface size
n = 512

# Surface generator
sg = tm.SurfaceGeneratorFilter2D()
sg.setSizes([n, n])
sg.random_seed = 0

# Spectrum
spectrum = tm.Isopowerlaw2D()
sg.setFilter(spectrum)

# Parameters
spectrum.q0 = 16
spectrum.q1 = 16
spectrum.q2 = 64
spectrum.hurst = 0.8

# Generating surface
surface = sg.buildSurface()
#surface /= tm.SurfaceStatistics.computeSpectralRMSSlope(surface)
surface /= n
#print(spectrum.rmsSlopes())
#print(tm.SurfaceStatistics.computeRMSSlope(surface))

plt.imshow(surface)

# Creating model
model = tm.ModelFactory.createModel(tm.model_type_basic_2d, [1., 1.], [n, n])

# Solver
solver = tm.PolonskyKeerRey(model, surface, 1e-12,
                            tm.PolonskyKeerRey.pressure,
                            tm.PolonskyKeerRey.pressure)

# Solve for target pressure
p_target = 0.1
solver.solve(p_target)

plt.figure()
plt.imshow(model.getTraction())
print(model.getTraction().mean())

# Cleanup threads
tm.finalize()

plt.show()