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

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()
non_tabulated_gf = cpt.Delhommeau(tabulation_nr=0, tabulation_nz=0)


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 > 0:
        return pd.DataFrame([])

    if wavelength == 0.0:
        wavenumber = np.infty
    else:
        wavenumber = 2*np.pi/wavelength

    if depth == 'rankine':
        free_surface, water_depth = np.infty, np.infty
    else:
        free_surface, water_depth = 0.0, depth

    full_diag_S, full_diag_K = basic_engine.build_matrices(
        mesh_1, mesh_1, free_surface, water_depth, wavenumber, non_tabulated_gf
    )
    full_offdiag_S, full_offdiag_K = basic_engine.build_matrices(
        mesh_1, mesh_2, free_surface, water_depth, wavenumber, non_tabulated_gf
    )

    svd_offdiag_S = LowRankMatrix.from_full_matrix_with_SVD(full_offdiag_S, max_rank=rank)
    error_S_svd = norm(svd_offdiag_S.full_matrix() - full_offdiag_S)/np.sqrt(norm(full_offdiag_S)**2 + norm(full_diag_S)**2)

    svd_offdiag_K = LowRankMatrix.from_full_matrix_with_SVD(full_offdiag_K, max_rank=rank)
    error_K_svd = norm(svd_offdiag_K.full_matrix() - full_offdiag_K)/np.sqrt(norm(full_offdiag_K)**2 + norm(full_diag_K)**2)

    aca_offdiag_S = LowRankMatrix.from_full_matrix_with_ACA(full_offdiag_S, max_rank=rank)
    error_S_aca = norm(aca_offdiag_S.full_matrix() - full_offdiag_S)/np.sqrt(norm(full_offdiag_S)**2 + norm(full_diag_S)**2)

    aca_offdiag_K = LowRankMatrix.from_full_matrix_with_ACA(full_offdiag_K, max_rank=rank)
    error_K_aca = norm(aca_offdiag_K.full_matrix() - full_offdiag_K)/np.sqrt(norm(full_offdiag_K)**2 + norm(full_diag_K)**2)

    return pd.DataFrame([
        {**inputs, 'method': 'ACA', 'matrix': 'S', 'fro_error': error_S_aca},
        {**inputs, 'method': 'ACA', 'matrix': 'K', 'fro_error': error_K_aca},
        {**inputs, 'method': 'SVD', 'matrix': 'S', 'fro_error': error_S_svd},
        {**inputs, 'method': 'SVD', 'matrix': 'K', 'fro_error': error_K_svd},
        ])


# ROLE OF THE RANK OF THE APPROXIMATION
print("Compute Figure 3: error as function of distance and rank")
distance_range = np.logspace(0, 1.5, 30)
ranks_range = [0, 1, 2, 3, 4, 5]
data = [run_case(distance=d, wavelength=1.0, depth=np.infty, rank=r) for d in distance_range for r in ranks_range]
df = pd.concat(data, ignore_index=True)
cmap = sns.color_palette("Blues", as_cmap=True)
fig, axs = four_panels_plot(
    df,
    x="distance", y="fro_error", color_var="rank",
    line_kwargs= lambda r: dict(color=cmap((r+2)/7), label=f"{r}", marker='.'),
    ax_kwargs=dict(xscale="log", yscale="log"),
)
output_path = "results/error_of_distance_and_rank.pdf"
plt.savefig(output_path)
print(f"Done! Saved in {output_path}")

## ROLE OF THE WAVELENGTH
print("Compute Figure 4: error as function of distance and wavelength")
distance_range = np.logspace(0, 1, 20)  # Do not go higher than 10m because insufficient number of integration points for higher wavenumbers
wavelength_range = np.logspace(-0.5, 1.5, 8)
data = [run_case(distance=d, wavelength=l, depth=np.infty, rank=2) for d in distance_range for l in wavelength_range]
data.extend([run_case(distance=d, wavelength=0, depth='rankine', rank=2) for d in distance_range])
df = pd.concat(data, ignore_index=True)
cmap = sns.color_palette("flare", as_cmap=True)
fig, axs = four_panels_plot(
    df,
    x="distance", y="fro_error", color_var="wavelength",
    line_kwargs= lambda l: (
                        dict(color=cmap((np.log10(l) + 0.5)/2.0), label=f"{l:.2f}", marker='.') if l > 0.0 else
                        dict(color="grey", linestyle="--", label="Rankine only")
                        ),
    ax_kwargs=dict(xscale="log", yscale="log"),
)
output_path = "results/error_of_distance_and_wavelength.pdf"
plt.savefig(output_path)
print(f"Done! Saved in {output_path}")
