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.collections import PolyCollection

# DATA SECTOR: same pressure/wind simulation, 6-hour bins, and gradient fill as original.py.
np.random.seed(42)
time_step = np.linspace(0, 24, 60)
pressure = 1010 + 25 * np.sin(np.pi * time_step / 6) + 0.5 * np.random.normal(0, 5, 60)
wind_speed = 7 + 3 * np.sin(np.pi * time_step / 12) ** 2 + np.cos(np.pi * time_step / 12) + 0.05 * np.random.normal(0, 2, 60)
bins = np.arange(0, 25, 6)
inds = np.digitize(time_step, bins) - 1
group_p = [pressure[inds == i].mean() for i in range(len(bins) - 1)]
group_w = [wind_speed[inds == i].mean() for i in range(len(bins) - 1)]
group_centers = (bins[:-1] + bins[1:]) / 2
fig, axes = plt.subplots(3, 1, figsize=(6.7, 7.25))
ax_a = axes[0]
bar_width = 1.75
ax_a.bar(group_centers - bar_width/2, group_p, width=bar_width, color=COL_BLUE, alpha=0.82, edgecolor="white", linewidth=0.65, label="Avg pressure")
ax_a.set_ylabel("Pressure (hPa)", color=COL_BLUE)
ax_a.tick_params(axis="y", labelcolor=COL_BLUE)
polish_axes(ax_a, "y")
ax_a2 = ax_a.twinx()
ax_a2.plot(group_centers, group_w, "-o", color=COL_TEAL, linewidth=1.55, markersize=4.2, label="Avg wind speed")
ax_a2.set_ylabel("Wind speed (km/h)", color=COL_TEAL)
ax_a2.tick_params(length=0, colors=COL_TEAL)
ax_a2.spines["top"].set_visible(False)
ax_a2.spines["right"].set_color("#303642")
handles1, labels1 = ax_a.get_legend_handles_labels()
handles2, labels2 = ax_a2.get_legend_handles_labels()
ax_a.legend(handles1 + handles2, labels1 + labels2, frameon=False, loc="upper center", ncol=2)
ax_a.set_title("6-hour grouped averages", loc="left", color=COL_INK, fontweight=600, pad=4)

ax_b = axes[1]
ax_b2 = ax_b.twinx()
ax_b.fill_between(time_step, pressure, color=COL_BLUE, alpha=0.22, linewidth=0)
ax_b.plot(time_step, pressure, color=COL_BLUE, linewidth=1.25)
ax_b2.fill_between(time_step, wind_speed, color=COL_TEAL, alpha=0.18, linewidth=0)
ax_b2.plot(time_step, wind_speed, color=COL_TEAL, linewidth=1.25)
ax_b.set_ylabel("Pressure (hPa)", color=COL_BLUE)
ax_b2.set_ylabel("Wind speed", color=COL_TEAL)
ax_b.tick_params(axis="y", labelcolor=COL_BLUE)
ax_b2.tick_params(length=0, colors=COL_TEAL)
ax_b.set_xlim(0, 24)
polish_axes(ax_b, "y")
ax_b2.spines["top"].set_visible(False)
ax_b2.spines["right"].set_color("#303642")
ax_b.set_title("Raw data with dual area encodings", loc="left", color=COL_INK, fontweight=600, pad=4)

ax_c = axes[2]
verts = []
for i in range(len(time_step) - 1):
    verts.append([(time_step[i], 980), (time_step[i], pressure[i]), (time_step[i+1], pressure[i+1]), (time_step[i+1], 980)])
poly = PolyCollection(verts, array=pressure[:-1], cmap="viridis", edgecolors="none", alpha=0.92)
ax_c.add_collection(poly)
ax_c.autoscale_view()
imax = np.argmax(pressure)
ax_c.annotate("Peak", xy=(time_step[imax], pressure[imax]), xytext=(22, -18), textcoords="offset points", arrowprops=dict(arrowstyle="->", color=COL_RED, lw=0.8), color=COL_RED, fontweight=600)
ax_c.set_xlabel("Time (hours)")
ax_c.set_ylabel("Pressure (hPa)")
ax_c.set_title("Pressure gradient fill", loc="left", color=COL_INK, fontweight=600, pad=4)
polish_axes(ax_c, "y")
fig.subplots_adjust(left=0.10, right=0.90, bottom=0.07, top=0.97, hspace=0.45)
save(fig)