import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
import seaborn as sns

from capytaine.matrices.block import BlockMatrix
from capytaine.matrices.low_rank import LowRankMatrix

full_labels = {
        "fro_error": "Relative Frobenius error",
        "distance": "Distance between bodies",
        "wavelength": "Wavelength (m)",
        "rank": "Rank",
        "green_function": "Green function evaluation",
        "tol": "Error tolerance",
        "resolution": "Mesh resolution (m¯¹)",
        "nb_faces": "Number of faces in mesh",
        "wavenumber": "Wavenumber (m¯¹)"
        }

def four_panels_plot(df, *, figax=None, x, y, color_var, line_kwargs=(lambda c: {}), ax_kwargs={}):
    idf = df.set_index(['matrix', 'method', color_var]).sort_index()
    # fig, axs = plt.subplots(2, 2, sharex=True, sharey=True)
    if figax is None:
        fig, axd = plt.subplot_mosaic([['S_SVD', 'K_SVD'],
                                       ['S_ACA', 'K_ACA']],
                                      sharex=True, sharey=True,
                                      figsize=(8, 6), layout="constrained")
    else:
        fig, axd = figax
    color_var_values = df[color_var].unique()
    for i, method in enumerate(['SVD', 'ACA']):
        for j, matrix in enumerate(['S', 'K']):
            ax = axd[f"{matrix}_{method}"]
            for k, c in enumerate(sorted(color_var_values)):
                sub_df = idf.loc[(matrix, method, c)]
                ax.plot(sub_df[x], sub_df[y], **(line_kwargs(c)))
            ax.set(**ax_kwargs, title=f"Approximating {matrix} with {method}")
    axd["S_SVD"].set(ylabel=full_labels[y])
    axd["S_ACA"].set(xlabel=full_labels[x], ylabel=full_labels[y])
    axd["K_ACA"].set(xlabel=full_labels[x])
    h, l = axd["S_SVD"].get_legend_handles_labels()
    fig.legend(h, l, loc='outside right', title=full_labels[color_var])
    return fig, axd


_red = sns.color_palette("hls", 8)[0]
_green = sns.color_palette("hls", 8)[2]

def plot_patches_of_matrix(ax, matrix, global_frame=(0, 0)):
    all_blocks_in_flat_iterator = (block for line in matrix._stored_blocks for block in line)
    positions_of_all_blocks = matrix._stored_block_positions(global_frame=global_frame)

    for block, positions_of_the_block in zip(all_blocks_in_flat_iterator, positions_of_all_blocks):
        position = positions_of_the_block[0]

        if isinstance(block, BlockMatrix):
            plot_patches_of_matrix(ax, block, position)

        elif isinstance(block, np.ndarray):
            ax.add_patch(Rectangle((position[1], position[0]), block.shape[1], block.shape[0], edgecolor='k', facecolor=_red))

        elif isinstance(block, LowRankMatrix):
            ax.add_patch(Rectangle((position[1], position[0]), block.shape[1], block.shape[0], edgecolor='k', facecolor=_green))
            if block.shape[0]*block.shape[1] > 20_000:
                ax.annotate(block.rank, (position[1]+block.shape[1]/2, position[0]+block.shape[0]/2), color='k', size=8, va="center", ha="center")

