from __future__ import annotations

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt

mpl.rcParams.update({
    "pdf.fonttype": 42,
    "ps.fonttype": 42,
    "figure.dpi": 170,
    "savefig.dpi": 240,
    "font.family": "DejaVu Sans",
    "font.size": 8.2,
    "axes.titlesize": 9.6,
    "axes.labelsize": 8.4,
    "xtick.labelsize": 7.2,
    "ytick.labelsize": 7.2,
    "legend.fontsize": 7.2,
    "axes.linewidth": 0.75,
})

COL_INK = "#172033"
COL_MUTED = "#667085"
COL_GRID = "#d9dee7"
COL_BLUE = "#356ca5"
COL_TEAL = "#2f9c95"
COL_ORANGE = "#d8863b"
COL_RED = "#c75756"
COL_PURPLE = "#8066a8"
COL_GREEN = "#5f9b68"


def polish_axes(ax, grid_axis="y"):
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    ax.spines["left"].set_color("#303642")
    ax.spines["bottom"].set_color("#303642")
    ax.tick_params(length=0, colors=COL_MUTED, pad=3)
    if grid_axis:
        ax.grid(True, axis=grid_axis, color=COL_GRID, linewidth=0.55, alpha=0.85)
        ax.set_axisbelow(True)


def save(fig):
    fig.savefig("augmented.png", bbox_inches="tight", facecolor="white")
    fig.savefig("augmented.pdf", bbox_inches="tight", facecolor="white")

# DATA SECTOR: same scalar field, extrema, and gradient as original.py.
x = np.linspace(-10, 10, 400)
y = np.linspace(-10, 10, 400)
X, Y = np.meshgrid(x, y)


def gauss(X, Y, mu_x, mu_y, sx, sy):
    return np.exp(-(((X - mu_x) ** 2) / (2 * sx ** 2) + ((Y - mu_y) ** 2) / (2 * sy ** 2)))


Z1 = gauss(X, Y, -5, 5, 4, 4)
Z2 = gauss(X, Y, 3, 3, 1.5, 1.5)
Z3 = gauss(X, Y, -2, -2, 2.5, 2.5)
Z4 = gauss(X, Y, 5, -4, 3, 2)
Z5 = gauss(X, Y, 0, -6, 2, 2)
Z = (Z1 + Z2 + Z3 + Z4) - 1.5 * Z5
Z = Z / np.abs(Z).max()
max_idx = np.unravel_index(np.argmax(Z), Z.shape)
min_idx = np.unravel_index(np.argmin(Z), Z.shape)
max_loc = (x[max_idx[1]], y[max_idx[0]])
min_loc = (x[min_idx[1]], y[min_idx[0]])
dy, dx = np.gradient(Z, y, x)
magnitude = np.sqrt(dx ** 2 + dy ** 2)

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(7.4, 3.15), constrained_layout=True)
fig.suptitle("Complex field analysis", x=0.02, y=1.03, ha="left", color=COL_INK, fontweight=600)
levels = np.linspace(-np.abs(Z).max(), np.abs(Z).max(), 30)
cf = ax1.contourf(X, Y, Z, levels=levels, cmap="coolwarm", extend="both")
ax1.contour(X, Y, Z, levels=levels[levels > 0], colors="#20252e", linewidths=0.42, alpha=0.65)
ax1.contour(X, Y, Z, levels=levels[levels < 0], colors="#20252e", linewidths=0.42, linestyles="dashed", alpha=0.55)
ax1.scatter(*max_loc, marker="X", s=52, color="#f2bf3d", edgecolor="#20252e", linewidth=0.55, zorder=5, label="max")
ax1.scatter(*min_loc, marker="P", s=52, color="#40b6c4", edgecolor="#20252e", linewidth=0.55, zorder=5, label="min")
ax1.annotate("max", xy=max_loc, xytext=(-32, 16), textcoords="offset points", arrowprops=dict(arrowstyle="-", lw=0.7, color=COL_MUTED), color=COL_INK)
ax1.annotate("min", xy=min_loc, xytext=(16, -18), textcoords="offset points", arrowprops=dict(arrowstyle="-", lw=0.7, color=COL_MUTED), color=COL_INK)
ax1.set_title("Scalar field", loc="left", color=COL_INK, pad=4)
ax1.set_xlabel("X-axis")
ax1.set_ylabel("Y-axis")
ax1.set_aspect("equal", adjustable="box")
ax1.legend(frameon=False, loc="lower left", ncol=2, handlelength=1.0, columnspacing=0.9)
cbar1 = fig.colorbar(cf, ax=ax1, fraction=0.046, pad=0.025)
cbar1.set_label("Normalized value", color=COL_MUTED)
cbar1.ax.tick_params(length=0, colors=COL_MUTED)
cbar1.outline.set_linewidth(0.55)

im = ax2.imshow(magnitude, extent=[-10, 10, -10, 10], origin="lower", cmap="inferno")
ax2.streamplot(X, Y, dx, dy, color="white", linewidth=0.48, density=1.22, arrowstyle="->", arrowsize=0.68)
ax2.set_title("Gradient magnitude and direction", loc="left", color=COL_INK, pad=4)
ax2.set_xlabel("X-axis")
ax2.set_ylabel("Y-axis")
ax2.set_xlim(-10, 10)
ax2.set_ylim(-10, 10)
ax2.set_aspect("equal", adjustable="box")
cbar2 = fig.colorbar(im, ax=ax2, fraction=0.046, pad=0.025)
cbar2.set_label("Gradient magnitude", color=COL_MUTED)
cbar2.ax.tick_params(length=0, colors=COL_MUTED)
cbar2.outline.set_linewidth(0.55)
for ax in (ax1, ax2):
    for spine in ax.spines.values():
        spine.set_color("#303642")
        spine.set_linewidth(0.7)
    ax.tick_params(length=0, colors=COL_MUTED)
save(fig)