from rich.progress import track

import numpy as np
import xarray as xr
import matplotlib.pyplot as plt

import capytaine as cpt
from capytaine.io.xarray import problems_from_dataset

from array_tools import make_array_of_random_buoys

nb_buoys = 40
seed = 0

body = make_array_of_random_buoys(nb_buoys, seed)
# body.show_matplotlib()
print("nb_faces:", body.mesh.nb_faces)

test_matrix = xr.Dataset(coords={
    "wavelength": np.linspace(6, 30, 100),
    "wave_direction": [0.0],
    })

engine = cpt.BasicMatrixEngine(linear_solver="lu_decomposition")
solver = cpt.BEMSolver(engine=engine)
problems = problems_from_dataset(test_matrix, body)
results = [solver.solve(pb, keep_details=False) for pb in track(sorted(problems))]
ds = cpt.assemble_dataset(results)

(
    np.real(ds["diffraction_force"])
    .sel(wave_direction=0.0)
    .isel(influenced_dof=range(4))
    .plot(x="wavelength", hue="influenced_dof")
)
plt.xlabel("Wavelength (m)")
plt.ylabel("Real part of the diffraction force (N)")
plt.title("")
plt.tight_layout()
plt.savefig("results/diff_force_example.pdf")
plt.show()

