import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
import seaborn as sns

res = pd.read_csv("results/array_data.csv", index_col=0)

cmap = sns.color_palette("flare", as_cmap=True)
markers = {1e-1: "^", 3e-3: "s", 1e-4: "p"}

def color_of_wavelength(l):
    lmin = res.wavelength.min()
    lmax = res.wavelength.max()
    return cmap((l - lmin)/(lmax - lmin))

## PLOTS
fig, ax = plt.subplots(tight_layout=True, figsize=(6, 4))
for group_values, group in res.groupby(["ACA_tol", "wavelength"]):
    ax.scatter(group['density'], group['error'], marker=markers[group_values[0]], color=color_of_wavelength(group_values[1]))

## LEGEND
blank = Line2D([], [], color='white')
legend_entries = []
legend_entries += [(blank, "ACA tolerance")]
legend_entries += [(Line2D([], [], linestyle='none', marker=m, color='gray'), "{:.0e}".format(val)) for (val, m) in markers.items()]
legend_entries += [(blank, "")]
legend_entries += [(blank, "Wavelength (m)")]
legend_entries += [(Line2D([], [], linestyle='none', marker='o', color=color_of_wavelength(l)), "{:.1f}".format(l)) for l in res.wavelength.unique()]
ax.legend(*zip(*legend_entries), loc='upper right', ncols=2)

ax.set(
       # xscale='log',
       yscale='log',
       ylabel="Relative maximal error on diffraction forces",
       xlabel="Matrix density",)
fig.savefig("results/density_vs_error.pdf")
# plt.show()
