import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

import capytaine as cpt
from capytaine.matrices import LowRankMatrix

from plot_tools import four_panels_plot

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 norm(A):
    """Frobenius norm"""
    return np.linalg.norm(A, 'fro')


def run_case(*, distance, wavelength, depth, rank):
    inputs = {'distance': distance, 'wavelength': wavelength, 'depth': depth, 'rank': rank}

    mesh_1 = cpt.mesh_sphere(radius=1.0/2, resolution=(20, 20), center=(0.0, 0.0, 0.0)).immersed_part()
    mesh_2 = mesh_1.copy(name="mesh_2").translated_x(distance)
    inputs['nb_faces'] = mesh_1.nb_faces
    print("Running: ", inputs)

    if 8*mesh_1.faces_radiuses.max() > wavelength:
        return pd.DataFrame([])

    wavenumber = 2*np.pi/wavelength
    free_surface, water_depth = 0.0, depth

    full_diag_exact_S, full_diag_exact_K = basic_engine.build_matrices(
        mesh_1, mesh_1, free_surface, water_depth, wavenumber, non_tabulated_gf
    )
    full_offdiag_exact_S, full_offdiag_exact_K = basic_engine.build_matrices(
        mesh_1, mesh_2, free_surface, water_depth, wavenumber, non_tabulated_gf
    )
    full_diag_S, full_diag_K = basic_engine.build_matrices(
        mesh_1, mesh_1, free_surface, water_depth, wavenumber, default_gf
    )
    full_offdiag_S, full_offdiag_K = basic_engine.build_matrices(
        mesh_1, mesh_2, free_surface, water_depth, wavenumber, default_gf
    )

    svd_offdiag_exact_S = LowRankMatrix.from_full_matrix_with_SVD(full_offdiag_exact_S, max_rank=rank)
    svd_offdiag_exact_K = LowRankMatrix.from_full_matrix_with_SVD(full_offdiag_exact_K, max_rank=rank)
    aca_offdiag_exact_S = LowRankMatrix.from_full_matrix_with_ACA(full_offdiag_exact_S, max_rank=rank)
    aca_offdiag_exact_K = LowRankMatrix.from_full_matrix_with_ACA(full_offdiag_exact_K, max_rank=rank)
    svd_offdiag_S       = LowRankMatrix.from_full_matrix_with_SVD(full_offdiag_S,       max_rank=rank)
    svd_offdiag_K       = LowRankMatrix.from_full_matrix_with_SVD(full_offdiag_K,       max_rank=rank)
    aca_offdiag_S       = LowRankMatrix.from_full_matrix_with_ACA(full_offdiag_S,       max_rank=rank)
    aca_offdiag_K       = LowRankMatrix.from_full_matrix_with_ACA(full_offdiag_K,       max_rank=rank)

    errorwrtfull_exact_S_svd = norm(svd_offdiag_exact_S.full_matrix() - full_offdiag_exact_S)/np.sqrt(norm(full_offdiag_exact_S)**2 + norm(full_diag_exact_S)**2)
    errorwrtfull_exact_K_svd = norm(svd_offdiag_exact_K.full_matrix() - full_offdiag_exact_K)/np.sqrt(norm(full_offdiag_exact_K)**2 + norm(full_diag_exact_K)**2)
    errorwrtfull_exact_S_aca = norm(aca_offdiag_exact_S.full_matrix() - full_offdiag_exact_S)/np.sqrt(norm(full_offdiag_exact_S)**2 + norm(full_diag_exact_S)**2)
    errorwrtfull_exact_K_aca = norm(aca_offdiag_exact_K.full_matrix() - full_offdiag_exact_K)/np.sqrt(norm(full_offdiag_exact_K)**2 + norm(full_diag_exact_K)**2)
    errorwrtfull_S_svd       = norm(svd_offdiag_S.full_matrix()       - full_offdiag_S)      /np.sqrt(norm(full_offdiag_S)**2       + norm(full_diag_S)**2)
    errorwrtfull_K_svd       = norm(svd_offdiag_K.full_matrix()       - full_offdiag_K)      /np.sqrt(norm(full_offdiag_K)**2       + norm(full_diag_K)**2)
    errorwrtfull_S_aca       = norm(aca_offdiag_S.full_matrix()       - full_offdiag_S)      /np.sqrt(norm(full_offdiag_S)**2       + norm(full_diag_S)**2)
    errorwrtfull_K_aca       = norm(aca_offdiag_K.full_matrix()       - full_offdiag_K)      /np.sqrt(norm(full_offdiag_K)**2       + norm(full_diag_K)**2)
    errorwrtexactfull_S_svd  = norm(svd_offdiag_S.full_matrix()       - full_offdiag_exact_S)/np.sqrt(norm(full_offdiag_exact_S)**2 + norm(full_diag_exact_S)**2)
    errorwrtexactfull_K_svd  = norm(svd_offdiag_K.full_matrix()       - full_offdiag_exact_K)/np.sqrt(norm(full_offdiag_exact_K)**2 + norm(full_diag_exact_K)**2)
    errorwrtexactfull_S_aca  = norm(aca_offdiag_S.full_matrix()       - full_offdiag_exact_S)/np.sqrt(norm(full_offdiag_exact_S)**2 + norm(full_diag_exact_S)**2)
    errorwrtexactfull_K_aca  = norm(aca_offdiag_K.full_matrix()       - full_offdiag_exact_K)/np.sqrt(norm(full_offdiag_exact_K)**2 + norm(full_diag_exact_K)**2)

    return pd.DataFrame([
        {**inputs, 'method': 'ACA', 'matrix': 'S', 'green_function': 'exact' ,                    'fro_error': errorwrtfull_exact_S_aca},
        {**inputs, 'method': 'ACA', 'matrix': 'K', 'green_function': 'exact' ,                    'fro_error': errorwrtfull_exact_K_aca},
        {**inputs, 'method': 'SVD', 'matrix': 'S', 'green_function': 'exact' ,                    'fro_error': errorwrtfull_exact_S_svd},
        {**inputs, 'method': 'SVD', 'matrix': 'K', 'green_function': 'exact' ,                    'fro_error': errorwrtfull_exact_K_svd},
        {**inputs, 'method': 'ACA', 'matrix': 'S', 'green_function': 'tabulation_wrt_tabulation', 'fro_error': errorwrtfull_S_aca},
        {**inputs, 'method': 'ACA', 'matrix': 'K', 'green_function': 'tabulation_wrt_tabulation', 'fro_error': errorwrtfull_K_aca},
        {**inputs, 'method': 'SVD', 'matrix': 'S', 'green_function': 'tabulation_wrt_tabulation', 'fro_error': errorwrtfull_S_svd},
        {**inputs, 'method': 'SVD', 'matrix': 'K', 'green_function': 'tabulation_wrt_tabulation', 'fro_error': errorwrtfull_K_svd},
        {**inputs, 'method': 'ACA', 'matrix': 'S', 'green_function': 'tabulation_wrt_exact' ,     'fro_error': errorwrtexactfull_S_aca},
        {**inputs, 'method': 'ACA', 'matrix': 'K', 'green_function': 'tabulation_wrt_exact' ,     'fro_error': errorwrtexactfull_K_aca},
        {**inputs, 'method': 'SVD', 'matrix': 'S', 'green_function': 'tabulation_wrt_exact' ,     'fro_error': errorwrtexactfull_S_svd},
        {**inputs, 'method': 'SVD', 'matrix': 'K', 'green_function': 'tabulation_wrt_exact' ,     'fro_error': errorwrtexactfull_K_svd},
        ])


print("Compute Figure 6: error as function of distance and GF evaluation method")
distance_range = np.logspace(0, 1.5, 30)
data = [run_case(distance=d, wavelength=1.0, depth=np.infty, rank=5) for d in distance_range]
df = pd.concat(data, ignore_index=True)

fig, axs = four_panels_plot(
    df,
    x="distance", y="fro_error", color_var="green_function",
    line_kwargs= lambda gf: (
        dict(label="Error 1", color="C1", marker='.') if gf == 'exact' else
        dict(label="Error 2", color="C2", marker='.') if gf == 'tabulation_wrt_tabulation' else
        dict(label="Error 3", color="C3", marker='.')
        ),
    ax_kwargs=dict(xscale="log", yscale="log"),
)
output_path = "results/error_of_distance_and_gf.pdf"
plt.savefig(output_path)
print(f"Done! Saved in {output_path}")
