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


def norm(A):
    """Frobenius norm"""
    return np.linalg.norm(A, 'fro')


def LowRankMatrix_from_full_matrix_with_SVD_tol(full_matrix, tol):
    decomp = np.linalg.svd(full_matrix)
    u, s, v = decomp
    assert np.all(s == sorted(s, reverse=True))
    rank = np.count_nonzero(s > tol)
    left_matrix = u[:, 0:rank] @ np.diag(s[0:rank])
    right_matrix = v[0:rank, :]
    return cpt.matrices.LowRankMatrix(left_matrix, right_matrix)


def LowRankMatrix_from_full_matrix_with_ACA_fixed_tol(full_matrix, tol):
   lr = LowRankMatrix.from_full_matrix_with_ACA(full_matrix, tol=tol)
   # Remove one rank because the implementation in Capytaine goes sligtly beyond the tol
   if lr.rank > 1:
       return LowRankMatrix(lr.left_matrix[:, :-1], lr.right_matrix[:-1, :])
   else:
       return lr


def run_case(*, distance, wavelength, depth=np.infty, tol, resolution):

    mesh_1 = cpt.mesh_sphere(radius=1.0/2, resolution=(int(resolution), int(resolution)), center=(0.0, 0.0, 0.0)).immersed_part()
    mesh_2 = mesh_1.copy(name="mesh_2").translated_x(distance)
    wavenumber = 2*np.pi/wavelength

    inputs = {'distance': distance, 'wavelength': wavelength, 'wavenumber': wavenumber, 'depth': depth, 'tol': tol, 'resolution': resolution, 'nb_faces': mesh_1.nb_faces}
    print("Running: ", inputs)

    if 8*mesh_1.faces_radiuses.max() > wavelength:
        # Skip as resolution is insufficient
        return pd.DataFrame([])

    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_tol(full_offdiag_S, tol=tol)
    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_tol(full_offdiag_K, tol=tol)
    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_fixed_tol(full_offdiag_S, tol=tol)
    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_fixed_tol(full_offdiag_K, tol=tol)
    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, 'rank': aca_offdiag_S.rank},
        {**inputs, 'method': 'ACA', 'matrix': 'K', 'fro_error': error_K_aca, 'rank': aca_offdiag_K.rank},
        {**inputs, 'method': 'SVD', 'matrix': 'S', 'fro_error': error_S_svd, 'rank': svd_offdiag_S.rank},
        {**inputs, 'method': 'SVD', 'matrix': 'K', 'fro_error': error_K_svd, 'rank': svd_offdiag_K.rank},
        ])


## CHECK RANK WRT RESOLUTION
print("Compute Figure 7: rank as function of mesh resolution and wavelength")
resolution_range = np.arange(8, 40, 2)
wavelength_range = np.logspace(-0.5, 1.5, 4)
data = [run_case(distance=2.0, wavelength=l, tol=1e-3, resolution=r) for r in resolution_range for l in wavelength_range]
df = pd.concat(data, ignore_index=True)

cmap = sns.color_palette("flare", as_cmap=True)
fig, axs = four_panels_plot(
    df,
    x="resolution", y="rank", color_var="wavelength",
    line_kwargs= lambda l: dict(color=cmap((np.log10(l) + 0.5)/2.0), label=f"{l:.2f}", marker='.'),
    ax_kwargs=dict(),
)
output_path = "results/rank_of_resolution_and_wavelength.pdf"
plt.savefig(output_path)
print(f"Done! Saved in {output_path}")
