#! /usr/bin/env python3

import numpy as np
from numpy.random import default_rng

import capytaine as cpt


def make_array_of_random_buoys(nb_buoys, seed):

    def one_buoy(x, y, z, r):
        # print(f"New buoy at {(x,y,z)} with radius {r}")
        buoy_mesh = cpt.mesh_sphere(radius=r, center=(x, y, z), resolution=(int(r*20), int(r*20)), name="buoy_at_{:.2f}_{:.2f}".format(x, y))
        buoy = cpt.FloatingBody(mesh=buoy_mesh.immersed_part())
        buoy.add_translation_dof(name="Heave")
        return buoy

    x, y = random_in_square(r=1.0, L=5*np.sqrt(nb_buoys), N=nb_buoys, seed=seed) # checks that there are no overlaps
    buoys_positions = np.array((x, y)).T

    # mean_dist_to_nearest_neighbours = np.mean(np.sort([np.hypot(xi, xj) for xi in buoys_positions for xj in buoys_positions])[:5])
    # print(mean_dist_to_nearest_neighbours)

    rng = default_rng(seed=seed)
    all_buoys = [one_buoy(x, y, 0.0, 0.5 + 0.5*rng.random()) for (x, y) in buoys_positions]
    all_buoys = cpt.FloatingBody.join_bodies(*all_buoys, name="many buoys")

    return all_buoys


def random_in_square(r, L, N, maxit=100, seed=None):
    """
    Puts N circles of radius r in a square of edge L, without overlaps
    kwarg: maxit (maximum number of total inner iterations)

    Returns positions x, y and status 0 for not converged, 1 for converged
    """

    max_packing = np.pi * np.sqrt(3) / 6
    if (N * np.pi*r**2/L**2) > max_packing:
        raise ValueError('Impossible packing!')

    rng = default_rng(seed=seed)

    xc = rng.uniform(-L/2, L/2, N)
    yc = rng.uniform(-L/2, L/2, N)

    nit = 0
    reiterate = True
    outer_iter = 0

    while reiterate:
        reiterate = False
        for ii in range(N):
            overlaps = False
            # Check if cylinder ii has overlaps with any other
            for jj in range(N):
                if ii!=jj:
                    dist2 = ((xc[jj] - xc[ii])**2 + (yc[jj] - yc[ii])**2)
                    if dist2 < 4*r**2:
                        overlaps = True
            nit = 0
            while overlaps and nit<maxit:
                nit += 1
                xc[ii] = rng.uniform(-L/2, L/2)
                yc[ii] = rng.uniform(-L/2, L/2)
                overlaps = False
                # Check if cylinder ii has overlaps with any other
                for jj in range(N):
                    if ii!=jj:
                        dist2 = ((xc[jj] - xc[ii])**2 + (yc[jj] - yc[ii])**2)
                        if dist2 < 4*r**2:
                            overlaps = True
            if nit==maxit:
                reiterate = True

        outer_iter += 1
    return xc, yc

if __name__ == "__main__":
    mesh = make_array_of_random_buoys(40, 0).mesh
    gf = cpt.Delhommeau()
    engine = cpt.HierarchicalToeplitzMatrixEngine(ACA_distance=1.5, ACA_tol=1e-2)
    S, K = engine.build_matrices(mesh, mesh, 0.0, np.infty, 1.0, gf)

    from plot_tools import plot_patches_of_matrix
    import matplotlib.pyplot as plt
    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("BLR_matrix_of_array.pdf")
