import pytest

import numpy as np
import xarray as xr

import capytaine as cpt
from capytaine import __version__

from capytaine.meshes.symmetric_meshes import ReflectionSymmetricMesh

@pytest.fixture
def sphere():
    mesh = cpt.mesh_sphere(radius=1.0, resolution=(4, 4)).immersed_part()
    sphere = cpt.FloatingBody(mesh=mesh)
    sphere.add_translation_dof(direction=(1, 0, 0), name="Surge")
    return sphere


def test_exportable_settings():
    gf = cpt.Delhommeau(
            tabulation_nr=10, tabulation_nz=10,
            tabulation_grid_shape="legacy",
            tabulation_nb_integration_points=50,
            finite_depth_prony_decomposition_method="fortran"
            )
    assert gf.exportable_settings['green_function'] == 'Delhommeau'
    assert gf.exportable_settings['tabulation_nb_integration_points'] == 50
    assert gf.exportable_settings['tabulation_grid_shape'] == "legacy"
    assert gf.exportable_settings['finite_depth_prony_decomposition_method'] == 'fortran'

    gf2 = cpt.XieDelhommeau()
    assert gf2.exportable_settings['green_function'] == 'XieDelhommeau'

    engine = cpt.BasicMatrixEngine(green_function=gf)
    assert engine.exportable_settings['engine'] == 'BasicMatrixEngine'
    assert engine.exportable_settings['linear_solver'] == 'lu_decomposition'

    solver = cpt.BEMSolver(engine=engine)
    assert solver.exportable_settings['green_function'] == 'Delhommeau'
    assert solver.exportable_settings['tabulation_nb_integration_points'] == 50
    assert solver.exportable_settings['finite_depth_prony_decomposition_method'] == 'fortran'
    assert solver.exportable_settings['engine'] == 'BasicMatrixEngine'
    assert solver.exportable_settings['linear_solver'] == 'lu_decomposition'

    solver = cpt.BEMSolver(green_function=gf)
    assert solver.exportable_settings['green_function'] == 'Delhommeau'
    assert solver.exportable_settings['tabulation_nb_integration_points'] == 50
    assert solver.exportable_settings['finite_depth_prony_decomposition_method'] == 'fortran'

def test_cannot_define_gf_and_engine_in_solver():
    with pytest.raises(ValueError):
        cpt.BEMSolver(engine=cpt.BasicMatrixEngine(), green_function=cpt.Delhommeau())

def test_solver_has_initialized_timer():
    s = cpt.BEMSolver()
    assert s.timer.total == 0.0

def test_solver_update_timer(sphere):
    problem = cpt.DiffractionProblem(body=sphere, omega=1.0)
    s = cpt.BEMSolver()
    s.solve(problem)
    assert s.timer.total > 0.0

def test_direct_solver(sphere):
    problem = cpt.DiffractionProblem(body=sphere, omega=1.0)
    direct_solver = cpt.BEMSolver(method='direct')
    direct_result = direct_solver.solve(problem)
    indirect_solver = cpt.BEMSolver(method='indirect')
    indirect_result = indirect_solver.solve(problem)
    assert direct_result.forces["Surge"] == pytest.approx(indirect_result.forces["Surge"], rel=1e-1)


@pytest.mark.parametrize("method", ["direct", "indirect"])
def test_same_result_with_symmetries(method):
    solver = cpt.BEMSolver(method=method)
    sym_mesh = ReflectionSymmetricMesh(cpt.mesh_sphere(center=(0, 2, 0)).immersed_part(), plane='xOz')
    sym_body = cpt.FloatingBody(mesh=sym_mesh, dofs=cpt.rigid_body_dofs())
    sym_result = solver.solve(cpt.DiffractionProblem(body=sym_body, omega=1.0))
    mesh = sym_mesh.merged()
    body = cpt.FloatingBody(mesh=mesh, dofs=cpt.rigid_body_dofs())
    result = solver.solve(cpt.DiffractionProblem(body=body, omega=1.0))
    assert sym_result.forces["Surge"] == pytest.approx(result.forces["Surge"], rel=1e-4)


def test_parallelization(sphere):
    pytest.importorskip("joblib")
    solver = cpt.BEMSolver()
    test_matrix = xr.Dataset(coords={
        'omega': np.linspace(0.1, 4.0, 3),
        'radiating_dof': list(sphere.dofs.keys()),
    })
    solver.fill_dataset(test_matrix, sphere, n_jobs=2)


@pytest.mark.parametrize("n_jobs", [1, 2])
@pytest.mark.parametrize("n_threads", [1, 2])
def test_control_threads(sphere, n_jobs, n_threads):
    pytest.importorskip("joblib")
    pytest.importorskip("threadpoolctl")
    solver = cpt.BEMSolver()
    test_matrix = xr.Dataset(coords={
        'omega': np.linspace(0.1, 4.0, 3),
        'radiating_dof': list(sphere.dofs.keys()),
    })
    solver.fill_dataset(test_matrix, sphere, n_jobs=n_jobs, n_threads=n_threads)


def test_nb_timer(sphere):
    pytest.importorskip("joblib")
    solver = cpt.BEMSolver()
    n_jobs = 3
    problems = [
            cpt.RadiationProblem(body=sphere, radiating_dof="Surge", omega=omega)
            for omega in np.linspace(0.1, 3.0, 5)
            ]
    solver.solve_all(problems, n_jobs=n_jobs)
    assert len(solver.timer_summary().columns) == n_jobs


def test_float32_solver(sphere):
    solver = cpt.BEMSolver(green_function=cpt.Delhommeau(floating_point_precision="float32"))
    pb = cpt.RadiationProblem(body=sphere, radiating_dof="Surge", omega=1.0)
    result = solver.solve(pb)
    assert result.pressure.dtype == 'complex64' and result.potential.dtype == 'complex64'


def test_LiangWuNoblesseGF(sphere):
    test_matrix = xr.Dataset(coords={
        'omega': np.linspace(0.1, 4.0, 3),
        'radiating_dof': list(sphere.dofs),
    })
    solver = cpt.BEMSolver(green_function=cpt.LiangWuNoblesseGF())
    ref_solver = cpt.BEMSolver(green_function=cpt.Delhommeau())
    ds = solver.fill_dataset(test_matrix, sphere)
    ref_ds = ref_solver.fill_dataset(test_matrix, sphere)
    assert np.allclose(ds.added_mass.values, ref_ds.added_mass.values, rtol=1e-2)


def test_fill_dataset(sphere):
    solver = cpt.BEMSolver()
    test_matrix = xr.Dataset(coords={
        'omega': np.linspace(0.1, 4.0, 3),
        'wave_direction': np.linspace(0.0, np.pi, 3),
        'radiating_dof': list(sphere.dofs.keys()),
        'rho': [1025.0],
        'water_depth': [np.inf, 30.0],
        'g': [9.81]
    })
    dataset = solver.fill_dataset(test_matrix, sphere, n_jobs=1)

    # Tests on the coordinates
    assert list(dataset.coords['influenced_dof']) == list(dataset.coords['radiating_dof']) == list(sphere.dofs.keys())
    assert dataset.body_name == sphere.name
    assert dataset.rho == test_matrix.rho
    assert dataset.g == test_matrix.g

    # Tests on the results
    assert 'added_mass' in dataset
    assert 'radiation_damping' in dataset
    assert 'Froude_Krylov_force' in dataset
    assert 'diffraction_force' in dataset

    # Test the attributes
    assert dataset.attrs['capytaine_version'] == __version__
    assert 'start_of_computation' in dataset.attrs

    # Try to strip out the outputs and recompute
    naked_data = dataset.drop_vars(["added_mass", "radiation_damping", "diffraction_force", "Froude_Krylov_force"])
    recomputed_dataset = solver.fill_dataset(naked_data, [sphere])
    assert recomputed_dataset.rho == dataset.rho
    assert recomputed_dataset.g == dataset.g
    assert "added_mass" in recomputed_dataset
    assert np.allclose(recomputed_dataset["added_mass"].data, dataset["added_mass"].data)


def test_warning_mesh_resolution(sphere, caplog):
    solver = cpt.BEMSolver()
    pb = cpt.RadiationProblem(body=sphere, wavelength=0.1*sphere.minimal_computable_wavelength)
    with caplog.at_level("WARNING"):
        solver.solve(pb)
    assert "resolution " in caplog.text
