# -*- coding: utf-8 -*-
"""
Created on Mon Nov 17 12:10:33 2025

@author: tdror
"""
# This code generates the figures from the main text of 
# "Deforestation Triggers State-dependent Cloud Regimes in the Amazon", 
# submitted to AGU journal Earth's Future.

import os
import pickle
import numpy as np
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import cartopy.feature as cfeature
from cartopy.feature import ShapelyFeature
from cartopy.io.shapereader import Reader
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
from matplotlib.lines import Line2D
import matplotlib.ticker as mticker

# --- Path to your data (.pkl files and shapefile); update this to your local folder ---
data_path = "PATH_TO_YOUR_DATA_FOLDER"

#%% Fig. 1

# Fig1 file path
fig1_file = "Fig1.pkl"
load_path = os.path.join(data_path, fig1_file)

# Load using pickle
with open(load_path, 'rb') as f:
    data = pickle.load(f)

# Extract variables
forest_grid = data['forest_grid']
nonforest_grid = data['nonforest_grid']
dCF_mean_grid = data['dCF_mean_grid']
dCF_std_grid = data['dCF_std_grid']

# --- Load state borders ---
shapefile_file = "BR_UF_2020.shp"
shapefile_path = os.path.join(data_path, shapefile_file)
state_borders = ShapelyFeature(Reader(shapefile_path).geometries(),
                               ccrs.PlateCarree(), edgecolor='black', facecolor='none')

# ============================
# Spatial domain & bin edges
# ============================
lon_min, lon_max = -65.0, -55.0
lat_min, lat_max = -14.0, -6.0
buffer = 2

lon_min_buff = lon_min - buffer
lon_max_buff = lon_max + buffer
lat_min_buff = lat_min - buffer
lat_max_buff = lat_max + buffer

# 1° edges
lon_bins = np.arange(lon_min, lon_max + 1, 1)
lat_bins = np.arange(lat_min, lat_max + 1, 1)

# centers & meshes
lon_centers = 0.5 * (lon_bins[:-1] + lon_bins[1:])
lat_centers = 0.5 * (lat_bins[:-1] + lat_bins[1:])
lon_mesh,  lat_mesh  = np.meshgrid(lon_bins,  lat_bins)     # for pcolormesh edges
lon_mesh_c, lat_mesh_c = np.meshgrid(lon_centers, lat_centers)  # for scatter dots


# ==========================
# Gridlines / tick labeling
# ==========================
lon_grid   = np.arange(np.round(lon_min_buff), np.round(lon_max_buff) + 1, 1)  # grid every 1°
lat_grid   = np.arange(np.round(lat_min_buff), np.round(lat_max_buff) + 1, 1)
lon_labels = np.arange(np.round(lon_min_buff), np.round(lon_max_buff) + 1, 2)  # labels every 2°
lat_labels = np.arange(np.round(lat_min_buff), np.round(lat_max_buff) + 1, 2)


# Plot

fig, axes = plt.subplots(2, 2, figsize=(16, 12), 
                         subplot_kw={"projection": ccrs.PlateCarree()})

axes = axes.flatten()

# Data + settings
grids  = [forest_grid, nonforest_grid, dCF_mean_grid, dCF_std_grid]
cmaps  = ["Greens", "Reds", "RdBu", "viridis"]

grid_names = ["forest_grid", "nonforest_grid", "dCF_mean_grid", "dCF_std_grid"]

titles = [
    r"Forest fraction ($f_{f}$)",
    r"Non-forest fraction ($f_{nf}$)",
    r"$\overline{\Delta CF}$",
    r"$\sigma_{\Delta CF}$"
]

vmins  = [0, 0, -0.1, 0]
vmaxs  = [1.0, 1.0, 0.1, 0.1]

# Colorbar ticks
cb_ticks = [
    [0, 0.5, 1],        # forest
    [0, 0.5, 1],        # nonforest
    [-0.1, 0, 0.1],     # mean ΔCF
    [0, 0.05, 0.1]      # std ΔCF
]

# Panel labels
panel_labels = ["(a)", "(b)", "(c)", "(d)"]

for i, (ax, grid, cmap, title, vmin_val, vmax_val, ticks, label) in enumerate(
    zip(axes, grids, cmaps, titles, vmins, vmaxs, cb_ticks, panel_labels)
):

    ax.set_extent([lon_min_buff, lon_max_buff, lat_min_buff, lat_max_buff], crs=ccrs.PlateCarree())
    ax.add_feature(cfeature.BORDERS, linewidth=1.5)
    # ax.add_feature(cfeature.COASTLINE, linewidth=1)
    ax.add_feature(cfeature.LAND, facecolor="lightgray")
    ax.add_feature(cfeature.RIVERS, edgecolor="blue", alpha=0.3)
    ax.add_feature(state_borders, linewidth=1.5, linestyle=":")

    # Gridlines
    ax.gridlines(draw_labels=False, xlocs=lon_grid, ylocs=lat_grid,
                 linestyle="--", linewidth=1.2, alpha=0.8)

    # Ticks
    ax.set_xticks(lon_labels, crs=ccrs.PlateCarree())
    ax.set_yticks(lat_labels, crs=ccrs.PlateCarree())
    ax.tick_params(axis="x", labelsize=14)
    ax.tick_params(axis="y", labelsize=14)

    # Remove y tick labels for 2nd and 4th subplot
    if i in [1, 3]:  
        ax.set_yticklabels([])

    # Remove x tick labels for 1st and 2nd subplot (upper row)
    if i in [0, 1]:
        ax.set_xticklabels([])

    # Colormap + colorbar
    pcm = ax.pcolormesh(lon_mesh, lat_mesh, grid, cmap=cmap,
                        vmin=vmin_val, vmax=vmax_val, shading="auto")

    cbar = plt.colorbar(pcm, ax=ax, orientation="vertical",
                        fraction=0.035, pad=0.015, ticks=ticks)
    cbar.ax.tick_params(labelsize=14)

    ax.set_title(title, fontsize=18)

    # Panel label
    ax.text(0.02, 0.95, label, transform=ax.transAxes,
            fontsize=18, va="top", ha="left")

# Make rows & cols closer
plt.subplots_adjust(wspace=0.03, hspace=0.13)
plt.show()

# %% Fig. 2

# Fig2 file path
fig2_file = "Fig2.pkl"
load_path = os.path.join(data_path, fig2_file)

# Load Fig2 data
with open(load_path, "rb") as f:
    fig2_data = pickle.load(f)
    
    
# Access variables
dCF_vals = fig2_data['dCF_vals']
mean_dCF = fig2_data['mean_dCF']
eps = fig2_data['eps']
len_neg_dcf = fig2_data['len_neg_dcf']
len_zero_dcf = fig2_data['len_zero_dcf']
len_pos_dcf = fig2_data['len_pos_dcf']
N = fig2_data['N']
images_array = fig2_data['images']  # actual numeric arrays with ΔCF


# Plot
fig, ax = plt.subplots(figsize=(10,6))

# Histogram
n, bins, patches = ax.hist(dCF_vals, bins=100, color="lightsteelblue",
                           edgecolor="black")
# Mean line as Bar
plt.axvline(mean_dCF, color="red", linestyle="--", alpha=0.7, linewidth=2,
            label=rf"$\overline{{\Delta \mathrm{{CF}}}}$ = {mean_dCF:.2f}")
# ΔCF = 0 line
plt.axvline(0, color="black", linestyle="-", linewidth=2, label=r"$\Delta$CF = 0")

# Faint vertical lines at -eps and +eps
ax.axvline(-eps, color="gray", linestyle="--", alpha=0.7, label=rf"$\epsilon = {eps}$")
ax.axvline(eps, color="gray", linestyle="--", alpha=0.7)

# Annotations
ax.text(0.11, 1400, "more clouds over non-forest", color="black", fontsize=14)
ax.text(-0.35, 1400, "more clouds over forest", color="black", fontsize=14)

# Dark green arrow pointing left from -0.1
plt.annotate(
    '', xy=(-0.17, 1300), xytext=(-0.066, 1300),
    arrowprops=dict(facecolor='darkgreen',
                    edgecolor='darkgreen', shrink=0.05, width=3, headwidth=10)
)

# Dark red arrow pointing right from 0.15
plt.annotate(
    '', xy=(0.2090, 1300), xytext=(0.105, 1300),
    arrowprops=dict(facecolor='darkred',
                    edgecolor='darkred', shrink=0.05, width=3, headwidth=10)
)

# Grid, labels, limits
ax.grid(True, color='black', linestyle='--', linewidth=0.5, alpha=0.3)
ax.set_xlabel(r"$\Delta$CF", fontsize=14)
ax.set_ylabel("Counts", fontsize=14)
ax.set_xlim((-0.5, 0.5))
ax.set_ylim((0,1800))
ax.legend(fontsize=11, loc="upper right")

# Custom ticks
plt.xticks(np.arange(-0.5, 0.6, 0.1), fontsize=14)
yticks = np.arange(0, 1900, 200)
ylabels = [str(y) if y % 400 == 0 else '' for y in yticks]
plt.yticks(yticks, ylabels, fontsize=14)

# --- Inset bar plot (normalized counts) ---
# inset_ax = ax.inset_axes([0.78, 0.17, 0.21, 0.3])  # lower-right corner
inset_ax = ax.inset_axes([0.1, 0.17, 0.3, 0.4])  # lower-right corner

groups = [r"$\Delta \mathrm{CF} \leq -\epsilon$",
          r"$-\epsilon < \Delta \mathrm{CF} < \epsilon$",
          r"$\Delta \mathrm{CF} \geq \epsilon$"]
counts = np.array([len_neg_dcf, len_zero_dcf,
                   len_pos_dcf])
counts_norm = counts / N  # normalize 0-1
colors = ["darkgreen", "black", "darkred"]

x_pos = np.arange(len(groups))
bar_width = 0.4  # thinner bars

inset_ax.bar(x_pos, counts_norm, color=colors, width=bar_width)
inset_ax.set_xticks(x_pos)
inset_ax.set_xlim(-0.5, 2.5)  # same as main histogram or whatever range you want
inset_ax.set_xticklabels(groups, rotation=25, ha="right", fontsize=11)
inset_ax.set_ylim(0, 0.8)
# inset_ax.set_yticks([0, 0.5, 1])
inset_ax.set_yticks([0, 0.4, 0.8])
inset_ax.set_title(r"$\Delta$CF group counts", fontsize=12)
inset_ax.tick_params(axis="y", labelsize=12)

# Annotate absolute counts
for i, count in enumerate(counts):
    inset_ax.text(x_pos[i], counts_norm[i] + 0.02, str(count), ha='center', va='bottom', fontsize=8)


plt.show()

# --- Plot images ---
fig, axes = plt.subplots(1, len(images_array), figsize=(5 * len(images_array), 5))
if len(images_array) == 1:
    axes = [axes]

for ax, (rgb_array, dcf_val) in zip(axes, images_array):
    ax.imshow(rgb_array)
    ax.set_title(f"ΔCF={dcf_val:.2f}", fontsize=26)
    ax.axis("off")

plt.tight_layout()
plt.show()


# %% Fig. 3

# Fig3 file path
fig3_file = "Fig3.pkl"
load_path = os.path.join(data_path, fig3_file)

# Load Fig3 data
with open(load_path, "rb") as f:
    precomputed = pickle.load(f)
    
# Access thresholds
low_threshold = precomputed["low_threshold"]
high_threshold = precomputed["high_threshold"]

# Access upper and lower tile data
upper_data = precomputed["upper_data"]
lower_data = precomputed["lower_data"]

# ----------------------
# Plot ΔCF vs CF_forest
# ----------------------
plt.figure(figsize=(9,6))

group_cmaps = {
    "Negative ΔCF": plt.cm.Greens,
    "Near-zero ΔCF": plt.cm.Greys,
    "Positive ΔCF": plt.cm.Reds
}

group_labels = {
    "Negative ΔCF": r"$\Delta \mathrm{CF} \leq -\epsilon$",
    "Near-zero ΔCF": r"$-\epsilon < \Delta \mathrm{CF} < \epsilon$",
    "Positive ΔCF": r"$\Delta \mathrm{CF} \geq \epsilon$"
}

shade_color = "gray"

# Plot each precomputed ΔCF group
for name in ["Negative ΔCF", "Near-zero ΔCF", "Positive ΔCF"]:
    cf, dCF_mean, se, cf_tot = precomputed[name]
    if len(cf) > 0:
        cmap = group_cmaps[name]
        norm = plt.Normalize(vmin=0, vmax=0.4)
        colors = cmap(norm(cf_tot))
        
        plt.fill_between(cf, dCF_mean - 2*se, dCF_mean + 2*se,
                         color=shade_color, alpha=0.4)
        plt.scatter(cf, dCF_mean, c=colors, s=60,
                    edgecolor='black', linewidth=0.5, label=group_labels[name])

# Zero line and labels
plt.axhline(0, color='black', linestyle='--')
plt.xlabel(r"$\mathrm{CF_{forest}}$", fontsize=14)
plt.ylabel("ΔCF (%)", fontsize=14)
plt.ylim((-175, 600))

# Legend using max color
handles, _ = plt.gca().get_legend_handles_labels()
new_handles = []
for name in group_labels:
    color = group_cmaps[name](0.8)
    new_handle = plt.Line2D([], [], marker='o', color='w',
                            markerfacecolor=color, markeredgecolor='black',
                            markersize=8, linestyle='None')
    new_handles.append(new_handle)
plt.legend(new_handles, list(group_labels.values()), fontsize=12)

# CF thresholds
plt.axvline(low_threshold, color='gray', linestyle='--', linewidth=2, zorder=0)
plt.axvline(high_threshold, color='gray', linestyle='--', linewidth=2, zorder=0)

# Arrows & annotations
plt.arrow(0.22, 100, 0, 40, color='darkred', head_width=0.008, head_length=30)
plt.arrow(0.22, -80, 0, -40, color='darkgreen', head_width=0.008, head_length=30)
plt.text(0.225, 100, "more clouds over non-forest", fontsize=12, ha='left', va='bottom')
plt.text(0.225, -80, "more clouds over forest", fontsize=12, ha='left', va='top')

# Grid
ax = plt.gca()
ax.grid(which='major', linestyle='--', alpha=1)
ax.grid(which='minor', linestyle='--', alpha=0.5)

# ----------------------
# Add 3 stacked colorbars
# ----------------------
from mpl_toolkits.axes_grid1 import make_axes_locatable

divider = make_axes_locatable(ax)
cbar_width = 0.015
gap = 0.03
height = ax.get_position().height
height_per_cbar = (height - 2*gap)/3
x0 = ax.get_position().x1 + 0.01

for i, name in enumerate(["Negative ΔCF", "Near-zero ΔCF", "Positive ΔCF"]):
    cmap = group_cmaps[name]
    norm = plt.Normalize(vmin=0, vmax=0.4)
    sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
    sm.set_array([])
    bottom = ax.get_position().y0 + i*(height_per_cbar + gap)
    cbar_ax = plt.gcf().add_axes([x0, bottom, cbar_width, height_per_cbar])
    cbar = plt.colorbar(sm, cax=cbar_ax)
    cbar.set_label(r"$\mathrm{CF_{tot}}$", fontsize=10)
    cbar.set_ticks([0, 0.2, 0.4])

plt.show()

# ----------------------
# Plot tile images
# ----------------------
ncols = max(len(upper_data), len(lower_data))
fig, axes = plt.subplots(2, ncols, figsize=(5 * ncols, 10))

for row_idx, tiles_data in enumerate([upper_data, lower_data]):
    for col_idx in range(ncols):
        ax = axes[row_idx, col_idx]
        if col_idx >= len(tiles_data):
            ax.axis("off")
            continue
        tile_name, dcf_val, rgb_array = tiles_data[col_idx]
        ax.imshow(rgb_array)
        ax.set_title(f"ΔCF% = {dcf_val:.1f}%", fontsize=26)
        ax.set_xticks([])
        ax.set_yticks([])
        for spine in ax.spines.values():
            spine.set_visible(False)
        if col_idx == 0:
            ylabel = r"$\mathrm{CF_{forest}^{-}}$" if row_idx == 0 else r"$\mathrm{CF_{forest}^{+}}$"
            ax.set_ylabel(ylabel, fontsize=26, labelpad=30)

plt.tight_layout(h_pad=2.5)
plt.show()

# %% Fig. 4

# Fig4 file path
fig4_file = "Fig4.pkl"
load_path = os.path.join(data_path, fig4_file)

# Load Fig4 data
with open(load_path, "rb") as f:
    precomputed = pickle.load(f)

all_group_stats = precomputed["all_group_stats"]
ceres_stats = precomputed["ceres_stats"]
group_order = precomputed["group_order"]

colors = {
    "Negative ΔCF": "green",
    "Near-zero ΔCF": "black",
    "Positive ΔCF - Low CF_forest": "lightcoral",
    "Positive ΔCF - High CF_forest": "darkred"
}

# === Labels and formatting ===
xlabels = [
    r"$\Delta \mathrm{CF} \leq -\epsilon$",
    r"$-\epsilon < \Delta \mathrm{CF} < \epsilon$",
    r"$\Delta \mathrm{CF} \geq \epsilon$" "\n" r"(CF$_{\rm forest}^-)$",
    r"$\Delta \mathrm{CF} \geq \epsilon$" "\n" r"(CF$_{\rm forest}^+)$"
]


# Labels, titles, limits
panel_vars_modis = ["cth", "lwp", "cot"]
panel_vars_ceres = ["toa_sw_all_daily", "toa_lw_all_daily", "toa_net_all_daily"]
titles_upper_2 = ['(a)', '(b)', '(c)', '(d)', '(e)', '(f)']
ylabels_modis = ["CTH (m)", "LWP (g m$^{-2}$)", "COT"]
ylabels_ceres = [r"$\mathrm{F_{SW}^\uparrow}$ (W m$^{-2}$)",
                 r"$\mathrm{F_{LW}^\uparrow}$ (W m$^{-2}$)",
                 r"$\mathrm{F_{net}}$ (W m$^{-2}$)"]

# Plot
fig, axes = plt.subplots(2, 3, figsize=(14, 7))
axes = axes.ravel()

# --- Upper row: MODIS ---
for i, (var, ax) in enumerate(zip(panel_vars_modis, axes[:3])):
    x_positions = np.arange(1, len(group_order) + 1)
    for x, g_name in zip(x_positions, group_order):
        stats = all_group_stats[g_name][var]
        color = colors[g_name]
        m, md, se = stats["mean"], stats["median"], stats["se"]
        p25, p75 = stats["p25"], stats["p75"]
        ax.plot(x, m, "o", color=color, markersize=9)
        ax.hlines([m - 2*se, m + 2*se], x - 0.1, x + 0.1, color=color, lw=3)
        ax.scatter(x, md, marker='D', color=color, s=80, zorder=3)
        ax.vlines(x, p25, p75, color=color, lw=2)
    ax.text(0.03, 0.95, titles_upper_2[i], transform=ax.transAxes,
            fontsize=14, va="top", ha="left",
            bbox=dict(facecolor='white', edgecolor='none', alpha=0.7))
    ax.grid(True, linestyle='--', alpha=0.9)
    ax.set_ylabel(ylabels_modis[i], fontsize=12)
    ax.tick_params(axis='y', labelsize=12)
    ax.set_xticks(x_positions)
    ax.set_xticklabels([])

# --- Lower row: CERES ---
for i, (var, ax) in enumerate(zip(panel_vars_ceres, axes[3:])):
    x_positions = np.arange(1, len(group_order) + 1)
    for x, g_name in zip(x_positions, group_order):
        stats = ceres_stats[g_name][var]
        color = colors[g_name]
        m, md, se = stats["mean"], stats["median"], stats["se"]
        p25, p75 = stats["p25"], stats["p75"]
        ax.plot(x, m, "o", color=color, markersize=9)
        ax.hlines([m - 2*se, m + 2*se], x - 0.1, x + 0.1, color=color, lw=3)
        ax.scatter(x, md, marker='D', color=color, s=80, zorder=3)
        ax.vlines(x, p25, p75, color=color, lw=2)
    ax.text(0.03, 0.95, titles_upper_2[i+3], transform=ax.transAxes,
            fontsize=14, va="top", ha="left",
            bbox=dict(facecolor='white', edgecolor='none', alpha=0.7))
    ax.grid(True, linestyle='--', alpha=0.9)
    ax.set_ylabel(ylabels_ceres[i], fontsize=12)
    ax.tick_params(axis='y', labelsize=12)
    ax.set_xticks(x_positions)
    ax.set_xticklabels(xlabels, rotation=25, ha='center', fontsize=12)

# Titles and spacing
fig.text(0.5, 0.96, "MODIS Cloud Properties", ha='center', va='center', fontsize=18)
fig.text(0.5, 0.50, "CERES Radiative Fluxes", ha='center', va='center', fontsize=18)
fig.subplots_adjust(hspace=0.27, wspace=0.35, top=0.92, bottom=0.08)

plt.show()

# %% Fig. 5

# Fig5 file path
fig5_file = "Fig5.pkl"
load_path = os.path.join(data_path, fig5_file)

# Load Fig5 data
with open(load_path, "rb") as f:
    fig5_data = pickle.load(f)

# Extract variables
low_KS_cloud = fig5_data["low_KS_cloud"]
low_KS_void = fig5_data["low_KS_void"]
high_KS_cloud = fig5_data["high_KS_cloud"]
high_KS_void = fig5_data["high_KS_void"]

# --- Plot ---
fig, ax = plt.subplots(figsize=(7, 7))

# Colors and labels
colors = ["lightcoral", "darkred"]
labels = [
    r"$\Delta \mathrm{CF} \geq \epsilon$, (CF$_{\rm forest}^-)$",
    r"$\Delta \mathrm{CF} \geq \epsilon$, (CF$_{\rm forest}^+)$"
]

# Plot subgroups
ax.scatter(low_KS_cloud, low_KS_void,
           color=colors[0], label=labels[0], alpha=0.7, s=100,
           edgecolor='k', linewidth=0.5)
ax.scatter(high_KS_cloud, high_KS_void,
           color=colors[1], label=labels[1], alpha=0.7, s=100,
           edgecolor='k', linewidth=0.25)

# Axes and style
ax.set_xlabel("Cloud deviation from randomness", fontsize=18)
ax.set_ylabel("Void deviation from randomness", fontsize=18)
ax.tick_params(axis='both', labelsize=14)
ax.grid(True, linestyle='--', alpha=0.6, linewidth=1.2)
ax.set_xlim(0.25, 0.95)
ax.set_ylim(0.05, 0.82)

# Legend
ax.legend(fontsize=14, loc='upper left', frameon=False)

# Title
plt.title("LvL Phase Space, Non-forest Preference Subgroups", fontsize=16)

plt.tight_layout()
plt.show()

# %% Fig. 6

# Fig6 file path
fig6_file = "Fig6.pkl"
load_path = os.path.join(data_path, fig6_file)

# Load Fig6 data
with open(load_path, "rb") as f:
    fig6_data = pickle.load(f)

all_group_stats = fig6_data["all_group_stats"]
p_lev = fig6_data["p_lev"]

# --- Configuration ---
sfc_plot_vars = ["sshf", "slhf", "sfc_wind_speed", "tp"]
plev_plot_vars = ["theta", "q", "w", "crwc"]
sfc_ylabels = [r"SH (W$\:m^{-2}$)", r"LH (W$\:m^{-2}$)", r"V$_{\rm sfc}$ (m$\:s^{-1}$)", r"P$_{\rm tot}$ (mm)"]

group_order = [
    "Negative ΔCF",
    "Near-zero ΔCF",
    "Positive ΔCF Low CF Forest",
    "Positive ΔCF High CF Forest"
]
colors = ["darkgreen", "black", "lightcoral", "darkred"]
offsets = [-0.3, -0.15, 0, 0.15]
xlabels = [
    r"$\Delta \mathrm{CF} \leq -\epsilon$",
    r"$-\epsilon < \Delta \mathrm{CF} < \epsilon$",
    r"$\Delta \mathrm{CF} \geq \epsilon$" "\n" r"(CF$_{\rm forest}^-)$",
    r"$\Delta \mathrm{CF} \geq \epsilon$" "\n" r"(CF$_{\rm forest}^+)$"
]

plev_var_info = {
    "crwc": ("Specific rain water content", "g kg⁻¹"),
    "q":    ("Specific humidity", "g kg⁻¹"),
    "theta": ("Potential temperature", "K"),
    "w":    ("Vertical velocity", "Pa s⁻¹"),
}

# --- Figure setup ---
fig, axes = plt.subplots(2, 4, figsize=(16, 8))
axes = axes.flatten()

panel_letters = ["(a)", "(b)", "(c)", "(d)"]
panel_letters_lower = ["(e)", "(f)", "(g)", "(h)"]

# --- Surface variables ---
for i, (ax, var, ylabel) in enumerate(zip(axes[:4], sfc_plot_vars, sfc_ylabels)):
    x = np.arange(len(group_order))
    x_tick_positions = [x_val + offset for x_val, offset in zip(x, offsets)]
    ax.set_xticks(x_tick_positions)
    ax.set_xticklabels(xlabels, rotation=25, ha="center", fontsize=13)

    for j, (g_key, color, offset) in enumerate(zip(group_order, colors, offsets)):
        stats = all_group_stats[g_key][var]
        y_mean = stats["mean"]
        y_median = stats["median"]
        se = stats["se"]
        p25 = stats["p25"]
        p75 = stats["p75"]
        x_pos = x[j] + offset

        if var in ["sshf", "slhf"]:
            y_mean = -y_mean; y_median = -y_median; p25 = -p25; p75 = -p75

        ax.vlines(x_pos, p25, p75, color=color, linewidth=3)
        ax.hlines([y_mean-2*se, y_mean+2*se], x_pos-0.15, x_pos+0.15, color=color, linewidth=3)
        ax.plot(x_pos, y_mean, '.', color=color, markersize=20)
        ax.plot(x_pos, y_median, 'D', color=color, markersize=10, markeredgewidth=0.8)

    ax.set_xlim([-0.75, len(group_order)-0.25])
    ax.set_ylabel(ylabel, fontsize=14)
    ax.tick_params(axis='y', labelsize=14)
    ax.grid(True)
    halign = "right" if var in ["sfc_wind_speed", "sshf"] else "left"
    xpos_letter = 0.97 if halign == "right" else 0.03
    ax.text(xpos_letter, 0.95, panel_letters[i], transform=ax.transAxes,
            fontsize=16, verticalalignment='top', horizontalalignment=halign,
            bbox=dict(facecolor='white', edgecolor='none', alpha=1.0))

# --- Pressure-level variables ---
for i, (ax, var) in enumerate(zip(axes[4:], plev_plot_vars)):
    if i == 0:
        lower_ax0 = ax
        inset_ax = inset_axes(lower_ax0, width="37%", height="60%", loc='lower right',
                              bbox_to_anchor=(-0.4, 0.15, 1.65, 1.25), bbox_transform=lower_ax0.transAxes, borderpad=0)
        inset_ax.set_yticks([1000, 900, 800])
        inset_ax.set_yticklabels([1000, 900, 800], fontsize=12)
    else:
        ax.sharey(lower_ax0)

    for j, (g_key, color) in enumerate(zip(group_order, colors)):
        mean_profile = all_group_stats[g_key][var]["mean"][:16]
        se_profile   = all_group_stats[g_key][var]["se"][:16]
        ax.plot(mean_profile, p_lev[:16], color=color, linewidth=3)
        ax.fill_betweenx(p_lev[:16], mean_profile-2*se_profile, mean_profile+2*se_profile, color=color, alpha=0.5)

        if i == 0:  # inset for theta
            mask = (p_lev[:16]>=800) & (p_lev[:16]<=1000)
            inset_ax.plot(mean_profile[mask], p_lev[:16][mask], color=color, linewidth=2)
            inset_ax.fill_betweenx(p_lev[:16][mask], mean_profile[mask]-2*se_profile[mask],
                                   mean_profile[mask]+2*se_profile[mask], color=color, alpha=0.4)

    if i == 0:
        inset_ax.set_ylim(1000, 800)
        inset_ax.set_xlim(np.min([all_group_stats[g]["theta"]["mean"][:16][mask].min() for g in group_order]),
                          np.max([all_group_stats[g]["theta"]["mean"][:16][mask].max() for g in group_order]))
        inset_ax.tick_params(axis='both', labelsize=10)
        inset_ax.grid(True)

    if var == "w":
        ax.axvline(0, color="gray", linestyle="--", linewidth=3.5, zorder=0)

    full_name, units = plev_var_info.get(var, (var, "(unit unknown)"))
    if var == "crwc":
        units = r"$\mathrm{g\:kg^{-1}}$"
        ax.xaxis.set_major_formatter(mticker.ScalarFormatter(useMathText=True))
        ax.xaxis.get_major_formatter().set_powerlimits((-3, 3))
        ax.locator_params(axis='x', nbins=4)

    if var == "theta":
        ax.set_xlabel(r"$\theta$ (K)", fontsize=16)
    else:
        ax.set_xlabel(f"{var} ({units})", fontsize=16)

    if i == 0:
        ax.set_ylabel("Pressure level (hPa)", fontsize=16)
    else:
        ax.tick_params(labelleft=False)

    ax.set_ylim(1000, 500)
    ax.grid(True)
    ax.tick_params(axis='both', which='major', labelsize=14)

    ha_letter = "right" if var in ["crwc","q"] else "left"
    xpos_letter = 0.97 if ha_letter=="right" else 0.03
    ax.text(xpos_letter, 0.95, panel_letters_lower[i], transform=ax.transAxes,
            fontsize=16, verticalalignment='top', horizontalalignment=ha_letter,
            bbox=dict(facecolor='white', edgecolor='none', alpha=1.0))

# --- Legend ---
legend_lines = [Line2D([0],[0], color=c, lw=3) for c in colors]
leg_txt = [
    '$\\Delta \\mathrm{CF} \\leq -\\epsilon$', '$-\\epsilon < \\Delta \\mathrm{CF} < \\epsilon$',
    '$\\Delta \\mathrm{CF} \\geq \\epsilon$, (CF$_{\\rm forest}^-)$',
    '$\\Delta \\mathrm{CF} \\geq \\epsilon$, (CF$_{\\rm forest}^+)$'
]
fig.legend(legend_lines, leg_txt, loc='lower center', ncol=4, fontsize=16, frameon=False,
           handlelength=2, handleheight=1.5)

plt.tight_layout(rect=[0, 0.05, 1, 0.95])
plt.show()