# %%
# Plots of determinant distribution from TGP invariant data

# %%
import pathlib

# Configure plotting
import numpy as np
import xarray as xr
import matplotlib
from matplotlib import pyplot as plt
from matplotlib.patches import Rectangle
from scipy.stats import special_ortho_group
from mpl_toolkits.axes_grid1.inset_locator import inset_axes

import multiterminal_invariant  # noqa: F401

figwidth = matplotlib.rcParams["figure.figsize"][0]

fig, axs = plt.subplots(
    1,
    3,
    figsize=(figwidth, figwidth / 3),
    sharex=True,
    sharey=True,
    layout="compressed",
)

np.random.seed(0)
size = 6
smatrices = special_ortho_group.rvs(size, size=10000)
axs[0].scatter(
    np.linalg.det(smatrices[:, : size // 2, : size // 2]),
    np.linalg.det(smatrices[:, size // 2 :, size // 2 :]),
    s=0.1,
    rasterized=True,
)
n1, n2 = size // 3, 2 * size // 3
axs[1].scatter(
    np.linalg.det(smatrices[:, :n1, :n1]),
    np.linalg.det(smatrices[:, n1:n2, n1:n2]),
    c=np.linalg.det(smatrices[:, n2:, n2:]),
    s=0.1,
    cmap="coolwarm",
    clim=(-1, 1),
    rasterized=True,
)

cax = inset_axes(
    axs[1],
    width="30%",
    height="5%",
    loc="lower left",
    bbox_to_anchor=(0, 0, 1, 1),
    bbox_transform=axs[1].transAxes,
    borderpad=0.9,
)
colorbar = fig.colorbar(axs[1].collections[0], cax=cax, orientation="horizontal")
colorbar.set_label(r"$\det r_3$")
colorbar.ax.set_xticks([-1, 0, 1])
colorbar.ax.xaxis.set_ticks_position("top")
colorbar.ax.xaxis.set_label_position("top")

blocks = [[(0, size // 2), (size // 2, size)], [(0, n1), (n1, n2)]]
for ax, block in zip(axs[:2], blocks):
    inset_ax = ax.inset_axes([0.05, 0.55, 0.4, 0.4])
    inset_ax.axis("off")
    inset_ax.set_xlim(0, size)
    inset_ax.set_ylim(0, size)
    inset_ax.add_patch(
        Rectangle((0, 0), size, size, edgecolor="black", facecolor="white", linewidth=2)
    )
    for pos, label in zip(block, ["L", "R"]):
        inset_ax.text(
            pos[0] + 0.5 * (pos[1] - pos[0]),
            pos[0] + 0.5 * (pos[1] - pos[0]),
            f"$r_{label}$",
            ha="center",
            va="center",
            fontsize=8,
        )
        inset_ax.add_patch(
            Rectangle(
                (pos[0], pos[0]),
                pos[1] - pos[0],
                pos[1] - pos[0],
                edgecolor="black",
                facecolor=(0.9, 0.9, 0.9, 0.8),
                linewidth=1,
            )
        )
    inset_ax.invert_yaxis()

# Gather all data using a loop comprehension
data = [
    ((ds := xr.load_dataset(filename)).detr_L.values, ds.detr_R.values)
    for filename in pathlib.Path("../data/tgp_determinants").glob("*.nc")
]

# Convert the plot to a 2D histogram
detr_L = np.concatenate([d[0].flatten() for d in data])
detr_R = np.concatenate([d[1].flatten() for d in data])
*_, img = axs[2].hist2d(
    detr_L,
    detr_R,
    bins=300,
    range=[[-1, 1], [-1, 1]],
    norm=plt.cm.colors.LogNorm(),
    cmap="inferno",
    rasterized=True,
)
colorbar = fig.colorbar(img, ax=axs[2], label=r"Counts")

# Find the dataset with the biggest variance of det r_L - det r_R
variances = [np.nanvar(d[0] - d[1]) for d in data]
max_variance_data = data[np.argmax(variances)]

# Scatter plot the data with the biggest variance as an inset
inset_ax = axs[2].inset_axes([0.6, 0.1, 0.3, 0.3])
inset_ax.hist2d(
    max_variance_data[0].flatten(),
    max_variance_data[1].flatten(),
    bins=100,
    range=[[-1, 1], [-1, 1]],
    norm=plt.cm.colors.LogNorm(),
    cmap="inferno",
    rasterized=True,
)
inset_ax.set_xlim(-1, 1)
inset_ax.set_ylim(-1, 1)
inset_ax.set_aspect("equal")
inset_ax.set_xticks([])
inset_ax.set_yticks([])

for ax in axs:
    ax.set_aspect("equal")
    ax.set_xlabel(r"$\det r_L$")
    ax.locator_params(nbins=3)
    ax.set_xlim(-1, 1)
    ax.set_ylim(-1, 1)
axs[0].set_ylabel(r"$\det r_R$")


axs[0].set_title("(a) Two blocks")
axs[1].set_title("(b) Three blocks")
axs[2].set_title("(c) Simulated TGP")
plt.savefig("../publication/figures/fig3.pdf")

# %%