import numpy as np
import pandas as pd
from joblib import Parallel, delayed
import capytaine as cpt

from array_tools import make_array_of_random_buoys

assert cpt.__version__ == '2.0'

basic_engine = cpt.BasicMatrixEngine(matrix_cache_size=0)
non_tabulated_gf = cpt.Delhommeau(tabulation_nr=0, tabulation_nz=0)
default_gf = cpt.Delhommeau()

def solve(*, engine, use_tabulation, problem):
    S, K = engine.build_matrices(
            problem.body.mesh, problem.body.mesh,
            problem.free_surface, problem.water_depth, problem.wavenumber,
            default_gf if use_tabulation else non_tabulated_gf,
            )
    sources = engine.linear_solver(K, problem.boundary_condition)
    pressure = 1j * problem.omega * problem.rho * (S @ sources)
    force = np.array([*problem.body.integrate_pressure(pressure).values()])
    return S, K, force


def reference_forces(*, nb_buoys, seed, wavelength, water_depth, use_tabulation):
    body = make_array_of_random_buoys(nb_buoys, seed)
    problem = cpt.DiffractionProblem(body=body, wavelength=wavelength, water_depth=water_depth, wave_direction=0.0)
    _, _, force = solve(engine=basic_engine, use_tabulation=use_tabulation, problem=problem)
    return force


def run_experiment(*, nb_buoys, seed, wavelength, ACA_distance, ACA_tol, water_depth, use_tabulation):
    inputs = {
            'nb_buoys': nb_buoys,
            'seed': seed,
            'wavelength': wavelength,
            'ACA_distance': ACA_distance,
            'ACA_tol': ACA_tol,
            'use_tabulation': use_tabulation,
            }
    print(inputs)
    body = make_array_of_random_buoys(nb_buoys, seed)
    problem = cpt.DiffractionProblem(body=body, wavelength=wavelength, water_depth=water_depth, wave_direction=0.0)
    engine = cpt.HierarchicalToeplitzMatrixEngine(ACA_distance=ACA_distance, ACA_tol=ACA_tol)
    S, K, force = solve(engine=engine, use_tabulation=use_tabulation, problem=problem)
    ref_force = reference_forces(nb_buoys=nb_buoys, seed=seed, wavelength=wavelength, water_depth=water_depth, use_tabulation=use_tabulation)
    return {
            **inputs,
            'nb_faces': body.mesh.nb_faces,
            'density': K.density,
            'error': abs(force - ref_force).max()/abs(ref_force).max()
            }


seed = range(5)
ACA_distance = [1.0]
ACA_tol = [1e-1, 3e-3, 1e-4]
wavelength = np.logspace(0.8, 1.5, 4)
experiments = [dict(nb_buoys=40, seed=s, ACA_distance=d, ACA_tol=t, wavelength=l, water_depth=np.infty, use_tabulation=True)
               for d in ACA_distance for t in ACA_tol for l in wavelength for s in seed]
data = Parallel(n_jobs=6)(delayed(run_experiment)(**exp) for exp in experiments)

res = pd.DataFrame(data)
res.to_csv("results/array_data.csv")
