import gzip, tempfile, shutil
import numpy as np
import capytaine as cpt

def norm(A):
    """Frobenius norm"""
    return np.linalg.norm(A, 'fro')

wavenumber = 2*np.pi/50.0
engine = cpt.HierarchicalToeplitzMatrixEngine(ACA_distance=1.0, ACA_tol=1e-2)
gf = cpt.Delhommeau()
solver = cpt.BEMSolver(engine=engine, green_function=gf)


## LOAD MESH
with tempfile.NamedTemporaryFile('wb') as tmp_f:
    with gzip.open("assets/KCS.mar.gz") as f:
        shutil.copyfileobj(f, tmp_f)
        mesh = cpt.load_mesh(tmp_f.name, file_format="mar", name=f.name)

body = cpt.FloatingBody(mesh=mesh.immersed_part())
body.add_all_rigid_body_dofs()


## MAKE CLUSTERS

minced_body = body.minced((16, 2, 1))

def get_leaves(tree):
    if isinstance(tree, cpt.CollectionOfMeshes):
        return sum([get_leaves(branch) for branch in tree._meshes], [])
    else:
        return [tree]
flat_mesh = cpt.CollectionOfMeshes(get_leaves(minced_body.mesh))
print("Nb clusters:", len(flat_mesh))

## SHOW MESH CLUSTERS

from capytaine.ui.vtk import mesh_viewer as mv
import seaborn as sns
from itertools import cycle
mv.COLORS = cycle(sns.color_palette("hls", 11))
# flat_mesh.show()

## SHOW MATRIX SHAPE
import matplotlib.pyplot as plt
from plot_tools import plot_patches_of_matrix

for mesh in [
    # flat_mesh,
    minced_body.mesh,
]:
    S, K = engine.build_matrices(mesh, mesh, 0.0, np.infty, wavenumber, gf)

    fig, ax = plt.subplots(1, 1, tight_layout=True)
    plot_patches_of_matrix(ax, K)
    ax.set(aspect="equal", xlim=(0, K.shape[1]), ylim=(0, K.shape[0]))
    ax.invert_yaxis()
    ax.xaxis.tick_top()

plt.savefig("results/H_matrix_of_ship.pdf")


print("density:", S.density, K.density)


## SOLVER PROBLEMS

print("Start full resolution")
pb = cpt.DiffractionProblem(body=body, wavenumber=wavenumber, wave_direction=0.0)
res = solver.solve(pb)
print("Done")
print("Start H-matrix resolution")
mpb = cpt.DiffractionProblem(body=minced_body, wavenumber=wavenumber, wave_direction=0.0)
mres = solver.solve(mpb)
print("Done")
print("error:", abs(res.forces["Heave"] - mres.forces["Heave"])/abs(res.forces["Heave"]))
