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")

from matplotlib.gridspec import GridSpec

# DATA SECTOR: same x/t grid, spreading sigma, field U, and sampled grid as original.py.
x = np.linspace(0, 1, 400)
t = np.linspace(0, 1, 400)
X, T = np.meshgrid(x, t)
sigma0 = 0.02
sigma1 = 0.24
Sigma = sigma0 + (sigma1 - sigma0) * T
U = np.exp(-((X - 0.5) ** 2) / (2 * Sigma ** 2))

fig = plt.figure(figsize=(6.1, 4.75))
gs = GridSpec(4, 1, figure=fig, hspace=0.08)
ax_profile = fig.add_subplot(gs[0, 0])
ax_contour = fig.add_subplot(gs[1:, 0], sharex=ax_profile)
levels_filled = np.linspace(U.min(), U.max(), 50)
cf = ax_contour.contourf(X, T, U, levels=levels_filled, cmap="viridis", extend="both")
ax_contour.contour(X, T, U, levels=np.linspace(U.min(), U.max(), 10), colors="white", linestyles="--", linewidths=0.35, alpha=0.8)
sample_rate = 40
ax_contour.scatter(X[::sample_rate, ::sample_rate], T[::sample_rate, ::sample_rate], s=7, c=COL_RED, alpha=0.58, linewidth=0, label="sampled grid")
ax_contour.legend(frameon=False, loc="upper left", handletextpad=0.3)
ax_contour.set_ylabel("t")
ax_contour.set_xlabel("x")
ax_contour.set_yticks(np.linspace(0, 1, 6))
ax_contour.tick_params(length=0, colors=COL_MUTED)
for spine in ax_contour.spines.values():
    spine.set_color("#303642")
    spine.set_linewidth(0.7)
cbar = fig.colorbar(cf, ax=ax_contour, fraction=0.035, pad=0.018)
cbar.set_label("u(x,t)", color=COL_MUTED)
cbar.ax.tick_params(length=0, colors=COL_MUTED)
cbar.outline.set_linewidth(0.55)

t_slice_index = np.argmin(np.abs(t - 0.5))
ax_profile.plot(x, U[t_slice_index, :], color="#20252e", linewidth=1.5)
ax_profile.fill_between(x, 0, U[t_slice_index, :], color=COL_BLUE, alpha=0.16, linewidth=0)
ax_profile.set_title(f"Profile of U at t={t[t_slice_index]:.2f}", loc="left", color=COL_INK, fontweight=600, pad=3)
ax_profile.set_ylabel("u")
ax_profile.set_ylim(0, 1.05)
polish_axes(ax_profile, "y")
plt.setp(ax_profile.get_xticklabels(), visible=False)
fig.suptitle("Contour field with marginal profile and sampled grid", x=0.08, y=0.985, ha="left", color=COL_INK, fontweight=600)
fig.subplots_adjust(left=0.09, right=0.89, bottom=0.10, top=0.90)
save(fig)