# -*- coding: utf-8 -*-
"""
Unified plotting for Figures from the main text of Dror et al., 2026

Author: Tom Dror
"""

import pickle
from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
from matplotlib.colors import LinearSegmentedColormap
import cartopy.crs as ccrs
import cartopy.feature as cfeature
from cartopy.io.shapereader import Reader
from cartopy.feature import ShapelyFeature
import matplotlib.ticker as mticker
import cmocean
import matplotlib.lines as mlines



# =============================
# USER: set this to your local data directory
# =============================
base_path = Path("/PATH/TO/YOUR/DATA/REPO")

fig1_upper_file = base_path / "Fig1_upper_row.pkl"
fig1_lower_file = base_path / "Fig1_lower_row.pkl"
fig2_upper_file = base_path / "Fig2_upper_row.pkl"
fig2_lower_file = base_path / "Fig2_lower_row.pkl"
fig3_upper_file = base_path / "Fig3_upper_row.pkl"
fig3_lower_file = base_path / "Fig3_lower_row.pkl"
fig4_file = base_path / "Fig4.pkl"

shapefile_path = r'C:\Users\tdror\Downloads\BR_UF_2020\BR_UF_2020.shp'

# =====================================================
# GLOBAL SETTINGS
# =====================================================
mpl.rcParams.update({
    "axes.labelsize": 9,
    "xtick.labelsize": 9,
    "ytick.labelsize": 9,
})

state_borders = ShapelyFeature(Reader(shapefile_path).geometries(),
                               ccrs.PlateCarree(), edgecolor='black', facecolor='none')

# =====================================================
# COLORMAPS
# =====================================================
gray_to_red = LinearSegmentedColormap.from_list(
    "gray_to_red",
    [
        (0.0, "#D3D3D3"),
        (0.15, "#D3D3D3"),
        (0.1501, "#FAA0A0"),
        (0.4, "#fb6a4a"),
        (0.7, "#de2d26"),
        (1.0, "#880808"),
    ],
    N=1024,
)


# -----------------------------
# Fig. 1, uppeer row 
# -----------------------------
# -----------------------------
# Load 
# -----------------------------

with open(fig1_upper_file, "rb") as f:
    data = pickle.load(f)

lon = data["lon"]
lat = data["lat"]
tree_cover = data["tree_cover"]
forest_loss = data["forest_loss"]
bbox_amazon = data["bbox_amazon"]
bbox_rondonia = data["bbox_rondonia"]

fig_width_in = 18.3 / 2.54
fig_height_in = fig_width_in / 2
fig, axes = plt.subplots(1, 2, figsize=(fig_width_in, fig_height_in),
                         subplot_kw={"projection": ccrs.PlateCarree()})

# Panel A — Tree cover
ax = axes[0]
im = ax.pcolormesh(lon, lat, tree_cover, cmap="Greens", vmin=0, vmax=1, transform=ccrs.PlateCarree())
ax.set_extent(bbox_amazon, crs=ccrs.PlateCarree())
ax.coastlines(linewidth=0.8)
ax.add_feature(cfeature.BORDERS, linewidth=0.8)
ax.add_feature(cfeature.RIVERS, alpha=0.35)
ax.add_feature(cfeature.OCEAN, alpha=0.30, zorder=0)
ax.add_patch(plt.Rectangle((bbox_rondonia[0], bbox_rondonia[2]),
                           bbox_rondonia[1] - bbox_rondonia[0],
                           bbox_rondonia[3] - bbox_rondonia[2],
                           edgecolor="gold", facecolor="none", linewidth=1.0,
                           transform=ccrs.PlateCarree()))
gl = ax.gridlines(draw_labels=True, linewidth=0.4, linestyle="--", alpha=0.5)
gl.top_labels = False
gl.right_labels = False
gl.xlocator = plt.FixedLocator(np.arange(-75, -45, 5))
gl.ylocator = plt.FixedLocator(np.arange(-15, 5, 5))
cbar = plt.colorbar(im, ax=ax, orientation="horizontal", pad=0.12, shrink=0.6)
cbar.set_label("Tree cover")
ax.text(0.02, 0.95, "A", transform=ax.transAxes, fontsize=10, fontweight="bold", va="top")

# Panel B — Forest loss
ax = axes[1]
im = ax.pcolormesh(lon, lat, forest_loss, cmap=gray_to_red, vmin=0, vmax=1, transform=ccrs.PlateCarree())
ax.set_extent(bbox_amazon, crs=ccrs.PlateCarree())
ax.coastlines(linewidth=0.8)
ax.add_feature(cfeature.BORDERS, linewidth=0.8)
ax.add_feature(cfeature.RIVERS, alpha=0.35)
ax.add_feature(cfeature.OCEAN, alpha=0.30, zorder=0)
ax.add_patch(plt.Rectangle((bbox_rondonia[0], bbox_rondonia[2]),
                           bbox_rondonia[1] - bbox_rondonia[0],
                           bbox_rondonia[3] - bbox_rondonia[2],
                           edgecolor="gold", facecolor="none", linewidth=1.0,
                           transform=ccrs.PlateCarree()))
gl = ax.gridlines(draw_labels=True, linewidth=0.4, linestyle="--", alpha=0.5)
gl.top_labels = False
gl.right_labels = False
gl.left_labels = False
gl.xlocator = plt.FixedLocator(np.arange(-75, -45, 5))
gl.ylocator = plt.FixedLocator(np.arange(-15, 5, 5))
cbar = plt.colorbar(im, ax=ax, orientation="horizontal", pad=0.12, shrink=0.6)
cbar.set_label(r"Loss fraction ($f_{loss}$)")
ax.text(0.02, 0.95, "B", transform=ax.transAxes, fontsize=10, fontweight="bold", va="top")

plt.tight_layout()
plt.show()


# -----------------------------
# Fig. 1, lower row (Rondônia zoom)
# -----------------------------
# -----------------------------
# Load 
# -----------------------------

with open(fig1_lower_file, "rb") as f:
    data = pickle.load(f)

ron_rgb = data["ron_rgb"]
lon_ron = data["lon_ron"]
lat_ron = data["lat_ron"]
cf_ron_dry = data["cf_ron_dry"]
cth_ron_dry = data["cth_ron_dry"]
bbox_rondonia = data["bbox_rondonia"]

fig_width_in = 7.25
fig_height_in = fig_width_in / 2.5
fig, axes = plt.subplots(1, 3, figsize=(fig_width_in, fig_height_in), subplot_kw={"projection": ccrs.PlateCarree()})
fig.subplots_adjust(left=0.05, right=0.97, top=0.92, bottom=0.18, wspace=0.05)

for i, ax in enumerate(axes):
    ax.set_extent(bbox_rondonia, crs=ccrs.PlateCarree())
    ax.add_feature(cfeature.RIVERS, linewidth=0.8, alpha=0.35)
    ax.add_feature(cfeature.BORDERS, linewidth=1.0)
    ax.add_feature(state_borders, linewidth=0.6, linestyle=":")
    gl = ax.gridlines(draw_labels=True, linewidth=0.4, color="gray", alpha=0.5, linestyle="--")
    gl.top_labels = False
    gl.right_labels = False
    if i > 0:
        gl.left_labels = False
    gl.xlabel_style = {"size": 9}
    gl.ylabel_style = {"size": 9}
    gl.xlocator = mticker.FixedLocator(np.arange(-68, -59, 2))
    gl.ylocator = mticker.FixedLocator(np.arange(-13, -8, 2))

axes[0].imshow(ron_rgb, origin="upper", extent=bbox_rondonia, transform=ccrs.PlateCarree())
axes[0].text(0.02, 0.95, "C", transform=axes[0].transAxes, fontsize=10, fontweight="bold", va="top", color="white")

im_cf = axes[1].pcolormesh(lon_ron, lat_ron, cf_ron_dry, cmap="Blues", vmin=0, vmax=0.6, transform=ccrs.PlateCarree())
axes[1].text(0.02, 0.95, "D", transform=axes[1].transAxes, fontsize=10, fontweight="bold", va="top", color="white")
cb_cf = fig.colorbar(im_cf, ax=axes[1], orientation="horizontal", pad=0.12, shrink=0.6, fraction=0.045)
cb_cf.set_label("Cloud fraction", fontsize=9)
cb_cf.set_ticks([0, 0.3, 0.6])
cb_cf.ax.tick_params(labelsize=9)

im_cth = axes[2].pcolormesh(lon_ron, lat_ron, cth_ron_dry, cmap="plasma", vmin=1000, vmax=2500, transform=ccrs.PlateCarree())
axes[2].text(0.02, 0.95, "E", transform=axes[2].transAxes, fontsize=10, fontweight="bold", va="top", color="white")
cb_cth = fig.colorbar(im_cth, ax=axes[2], orientation="horizontal", pad=0.12, shrink=0.6, fraction=0.045)
cb_cth.set_label("Cloud-top height (m)", fontsize=9)
cb_cth.set_ticks([1000, 1500, 2000, 2500])

plt.show()


# -----------------------------
# Fig. 2, upper row
# -----------------------------
# -----------------------------
# Load 
# -----------------------------

with open(fig2_upper_file, "rb") as f:
    data = pickle.load(f)

lon = data["lon"]
lat = data["lat"]
bbox_am_sub = data["bbox_am_sub"]
dCF = data["dCF_fcc_dry"]
dCTH = data["dCTH_fcc_dry"]

fig_width_in = 7.25
fig_height_in = fig_width_in / 2
fig, axes = plt.subplots(1, 2, figsize=(fig_width_in, fig_height_in), subplot_kw={"projection": ccrs.PlateCarree()})

# Panel A — ΔCF_fcc
ax = axes[0]
im = ax.pcolormesh(lon, lat, dCF, cmap=cmocean.cm.diff_r, vmin=-0.06, vmax=0.06, transform=ccrs.PlateCarree())
ax.set_extent(bbox_am_sub, crs=ccrs.PlateCarree())
ax.coastlines(linewidth=1)
ax.add_feature(cfeature.BORDERS, linewidth=1)
ax.add_feature(cfeature.RIVERS, alpha=0.35)
ax.add_feature(cfeature.OCEAN, alpha=0.3, zorder=0)
ax.add_feature(state_borders, linewidth=1, linestyle=":")
gl = ax.gridlines(draw_labels=True, linestyle="--", alpha=0.5)
gl.xlocator = plt.FixedLocator(np.arange(-75, -45, 5))
gl.ylocator = plt.FixedLocator(np.arange(-15, 5, 5))
gl.top_labels = False
gl.right_labels = False
cbar = plt.colorbar(im, ax=ax, orientation="horizontal", pad=0.12, shrink=0.6, fraction=0.03)
cbar.set_ticks([-0.05, 0, 0.05])
cbar.set_label(r"$\Delta \mathrm{CF}_{\mathrm{fcc}}$", fontsize=9)
ax.text(0.95, 0.93, "A", transform=ax.transAxes, fontsize=10, fontweight="bold", ha="right", va="top")

# Panel B — ΔCTH_fcc
ax = axes[1]
im = ax.pcolormesh(lon, lat, dCTH, cmap="RdYlBu", vmin=-300, vmax=300, transform=ccrs.PlateCarree())
ax.set_extent(bbox_am_sub, crs=ccrs.PlateCarree())
ax.coastlines(linewidth=1)
ax.add_feature(cfeature.BORDERS, linewidth=1)
ax.add_feature(cfeature.RIVERS, alpha=0.35)
ax.add_feature(cfeature.OCEAN, alpha=0.3, zorder=0)
ax.add_feature(state_borders, linewidth=1, linestyle=":")
gl = ax.gridlines(draw_labels=True, linestyle="--", alpha=0.5)
gl.xlocator = plt.FixedLocator(np.arange(-75, -45, 5))
gl.ylocator = plt.FixedLocator(np.arange(-15, 5, 5))
gl.top_labels = False
gl.right_labels = False
gl.left_labels = False
cbar = plt.colorbar(im, ax=ax, orientation="horizontal", pad=0.12, shrink=0.6, fraction=0.03)
cbar.set_ticks([-300, 0, 300])
cbar.set_label(r"$\Delta \mathrm{CTH}_{\mathrm{fcc}}$ (m)", fontsize=9)
ax.text(0.95, 0.93, "B", transform=ax.transAxes, fontsize=10, fontweight="bold", ha="right", va="top")

plt.tight_layout()
plt.show()

# -----------------------------
# Fig. 2, lower row
# -----------------------------
# -----------------------------
# Load 
# -----------------------------
with open(fig2_lower_file, "rb") as f:
    data = pickle.load(f)

x_cf_bin = data["x_cf_bin"]
y_cf_bin = data["y_cf_bin"]
y_cf_se = data["y_cf_se"]
y_cf_box = data["y_cf_box"]

x_cth_bin = data["x_cth_bin"]
y_cth_bin = data["y_cth_bin"]
y_cth_se = data["y_cth_se"]
y_cth_box = data["y_cth_box"]

# Colors
cf_blue = (0.3, 0.5, 0.8)
cth_red = (0.8, 0.2, 0.2)
black = (0, 0, 0)
gray_fill = (0.8, 0.8, 0.8, 0.6)

# Figure
fig_width = 9
fig_height = 3.5
fig = plt.figure(figsize=(fig_width, fig_height))

ax_cf_scatter = fig.add_axes([0.05, 0.15, 0.32, 0.75])
ax_cf_box     = fig.add_axes([0.39, 0.15, 0.05, 0.75])
ax_cth_scatter= fig.add_axes([0.52, 0.15, 0.32, 0.75])
ax_cth_box    = fig.add_axes([0.88, 0.15, 0.05, 0.75])

# dCF scatter
ax_cf_scatter.fill_between(x_cf_bin, y_cf_bin - 2*np.array(y_cf_se), y_cf_bin + 2*np.array(y_cf_se), color=gray_fill)
ax_cf_scatter.scatter(x_cf_bin, y_cf_bin, color=cf_blue, edgecolor='black', s=15)
ax_cf_scatter.axhline(0, color=black, linestyle='--', linewidth=1)
ax_cf_scatter.axvline(0.5, color='red', linestyle='-.', linewidth=1)
ax_cf_scatter.set_xlabel(r"$f_{loss}$", fontsize=9)
ax_cf_scatter.set_ylabel(r"$\Delta \rm{CF_{fcc}}$", fontsize=9)
ax_cf_scatter.set_ylim(-0.05, 0.1)
ax_cf_scatter.set_yticks([-0.05, 0, 0.05, 0.1])
ax_cf_scatter.set_xticks([0, 0.2, 0.4, 0.6, 0.8])
ax_cf_scatter.grid(True, linestyle='--', linewidth=0.7)
ax_cf_scatter.text(0.05, 0.95, "C", transform=ax_cf_scatter.transAxes, fontsize=10, fontweight='bold', va='top', ha='left')

# dCF box
ax_cf_box.boxplot(y_cf_box, vert=True, patch_artist=True,
                  boxprops=dict(facecolor=cf_blue+(0.5,), color=black),
                  medianprops=dict(color=cf_blue, linewidth=2),
                  flierprops=dict(marker='o', color='green', markersize=5, visible=False),
                  widths=1.0)
ax_cf_box.set_ylim(ax_cf_scatter.get_ylim())
ax_cf_box.axhline(0, color=black, linestyle='--', linewidth=1)
ax_cf_box.set_xticklabels([])
ax_cf_box.set_yticklabels([])

# dCTH scatter
ax_cth_scatter.fill_between(x_cth_bin, y_cth_bin - 2*np.array(y_cth_se), y_cth_bin + 2*np.array(y_cth_se), color=gray_fill)
ax_cth_scatter.scatter(x_cth_bin, y_cth_bin, color=cth_red, edgecolor='black', s=15)
ax_cth_scatter.axhline(0, color=black, linestyle='--', linewidth=1)
ax_cth_scatter.axvline(0.5, color='red', linestyle='-.', linewidth=1)
ax_cth_scatter.set_xlabel(r"$f_{loss}$", fontsize=9)
ax_cth_scatter.set_ylabel(r"$\Delta \rm{CTH_{fcc}}$ (m)", fontsize=9)
ax_cth_scatter.set_ylim(-450, 350)
ax_cth_scatter.set_yticks([-300, -150, 0, 150, 300])
ax_cth_scatter.set_xticks([0, 0.2, 0.4, 0.6, 0.8])
ax_cth_scatter.grid(True, linestyle='--', linewidth=0.7)
ax_cth_scatter.text(0.05, 0.95, "D", transform=ax_cth_scatter.transAxes, fontsize=10, fontweight='bold', va='top', ha='left')

# dCTH box
ax_cth_box.boxplot(y_cth_box, vert=True, patch_artist=True,
                   boxprops=dict(facecolor=cth_red+(0.5,), color=black),
                   medianprops=dict(color=cth_red, linewidth=2),
                   flierprops=dict(marker='o', color='green', markersize=5, visible=False),
                   widths=1.0)
ax_cth_box.set_ylim(ax_cth_scatter.get_ylim())
ax_cth_box.axhline(0, color=black, linestyle='--', linewidth=1)
ax_cth_box.set_xticklabels([])
ax_cth_box.set_yticklabels([])
ax_cth_box.set_yticks(ax_cth_scatter.get_yticks())

plt.show()

# -----------------------------
# Fig. 3, upper row
# -----------------------------
# -----------------------------
# Load 
# -----------------------------
with open(fig3_upper_file, "rb") as f:
    data = pickle.load(f)

lon = data["lon"]
lat = data["lat"]
Fsw = data["Fsw"]
Flw = data["Flw"]

fields = ["Fsw", "Flw"]
values = [Fsw, Flw]

# -----------------------------
# Figure settings
# -----------------------------
fig_width_in = 7.25
fig_height_in = fig_width_in * 0.55  # keep ratio consistent
fig, axes = plt.subplots(1, 2, figsize=(fig_width_in, fig_height_in),
                         subplot_kw={"projection": ccrs.PlateCarree()})

cmaps = {"Fsw": "coolwarm_r", "Flw": "BrBG"}
vmins = {"Fsw": -20, "Flw": -7}
vmaxs = {"Fsw": 20, "Flw": 7}
ticks_dict = {"Fsw": [-20, -10, 0, 10, 20], "Flw": [-6, -3, 0, 3, 6]}
ylabel_dict = {
    "Fsw": r"$\Delta \rm{F_{SW,fcc}^{\uparrow}}$ (W m$^{-2}$)",
    "Flw": r"$\Delta \rm{F_{LW,fcc}^{\uparrow}}$ (W m$^{-2}$)"
}

# -----------------------------
# Plot panels
# -----------------------------
for i, field_name in enumerate(fields):
    ax = axes[i]
    val = values[i]

    sc = ax.scatter(lon, lat, c=val, s=.1, cmap=cmaps[field_name],
                    vmin=vmins[field_name], vmax=vmaxs[field_name],
                    transform=ccrs.PlateCarree())

    ax.set_extent([-75, -45, -15, 5], crs=ccrs.PlateCarree())
    ax.add_feature(cfeature.LAND, facecolor="lightgray", alpha=0.2)
    ax.add_feature(cfeature.COASTLINE)
    ax.add_feature(cfeature.BORDERS, linestyle=":")
    ax.add_feature(cfeature.RIVERS, edgecolor='blue', alpha=0.35)
    ax.add_feature(cfeature.OCEAN, facecolor='lightblue', alpha=0.4, zorder=0)

    # Gridlines
    gl = ax.gridlines(draw_labels=True, linewidth=0.5, color='gray', alpha=0.5, linestyle='--')
    gl.xlocator = plt.FixedLocator(np.arange(-75, -45+1, 5))
    gl.ylocator = plt.FixedLocator(np.arange(-15, 5+1, 5))
    gl.top_labels = False
    gl.right_labels = False
    gl.xlabel_style = {'fontsize': 9}
    if i == 0:
        gl.ylabel_style = {'fontsize': 9}
    else:
        gl.ylabel_style = {'fontsize': 0}

    # Colorbar
    cbar = plt.colorbar(sc, ax=ax, orientation="horizontal", pad=0.08, fraction=0.05, shrink=0.6)
    cbar.set_label(ylabel_dict[field_name], fontsize=9)
    cbar.set_ticks(ticks_dict[field_name])
    cbar.ax.tick_params(labelsize=9)

    # Subplot label
    ax.text(0.93, 0.95, chr(65+i), transform=ax.transAxes,
            fontsize=10, fontweight='bold', va='top', ha='left')

plt.tight_layout()
plt.show()


# -----------------------------
# Fig. 3, lower row
# -----------------------------

with open(fig3_lower_file, "rb") as f:
    data = pickle.load(f)

x_Fsw = np.array(data["x_Fsw"])
y_Fsw = np.array(data["y_Fsw"])
c_Fsw = np.array(data["c_Fsw"])
y_se_Fsw = np.array(data["y_se_Fsw"])

x_Flw = np.array(data["x_Flw"])
y_Flw = np.array(data["y_Flw"])
c_Flw = np.array(data["c_Flw"])
y_se_Flw = np.array(data["y_se_Flw"])

# -----------------------------
# Figure settings
# -----------------------------
fields = ["Fsw", "Flw"]
x_vals = [x_Fsw, x_Flw]
y_vals = [y_Fsw, y_Flw]
c_vals = [c_Fsw, c_Flw]
y_se_vals = [y_se_Fsw, y_se_Flw]

ylabel_dict = {
    "Fsw": r"$\Delta \rm{F_{SW,fcc}^{\uparrow}}$ (W m$^{-2}$)",
    "Flw": r"$\Delta \rm{F_{LW,fcc}}$ (W m$^{-2}$)"
}
ylim_dict = {"Fsw": (-4.5, 16), "Flw": (-0.7, 2)}
yticks_dict = {"Fsw": [0,5,10,15], "Flw": [0,1,2]}
cmap_dict = {"Fsw": cmocean.cm.diff_r, "Flw": "RdYlBu"}
vmin_dict = {"Fsw": -0.06, "Flw": -100}
vmax_dict = {"Fsw": 0.06, "Flw": 100}

fig, axes = plt.subplots(1, 2, figsize=(7.25, 3.25), sharey=False)

# -----------------------------
# Plot each panel
# -----------------------------
for i, field_name in enumerate(fields):
    ax = axes[i]
    x = x_vals[i]
    y = y_vals[i]
    c = c_vals[i]
    y_se = y_se_vals[i]

    # Shading 2-sigma
    ax.fill_between(x, y - 2*y_se, y + 2*y_se, color='gray', alpha=0.5, zorder=1)

    # Scatter colored by cf/cth
    sc = ax.scatter(x, y, c=c, cmap=cmap_dict[field_name],
                    vmin=vmin_dict[field_name], vmax=vmax_dict[field_name],
                    edgecolor='black', linewidth=0.5, s=15, zorder=2)

    # Vertical line at f_loss = 0.5
    ax.axvline(0.5, color='red', linestyle='-.', linewidth=1.5, zorder=1.5)
    # Horizontal line at y=0
    ax.axhline(0, color='black', linestyle='dashed', linewidth=1, zorder=0)

    # Labels
    ax.set_xlabel(r"$f_{loss}$", fontsize=9)
    ax.set_ylabel(ylabel_dict[field_name], fontsize=9)
    ax.set_ylim(ylim_dict[field_name])
    ax.set_yticks(yticks_dict[field_name])
    ax.grid(True, linestyle='--', linewidth=0.7)
    ax.set_axisbelow(True)
    ax.tick_params(axis='both', labelsize=9)

    # Subplot label
    ax.text(0.05, 0.95, chr(99+i).upper(), transform=ax.transAxes,
            fontsize=10, fontweight='bold', va='top', ha='left')

    # Colorbar
    cbar = fig.colorbar(sc, ax=ax, orientation='horizontal', pad=0.17, shrink=0.6)
    if field_name == "Fsw":
        cbar.set_label(r"$\Delta \rm{CF_{fcc}}$", fontsize=9)
        cbar.set_ticks([-0.06, -0.03, 0, 0.03, 0.06])
    elif field_name == "Flw":
        cbar.set_label(r"$\Delta \rm{CTH_{fcc}}$ (m)", fontsize=9)
        cbar.set_ticks([-100, 0, 100])
    cbar.ax.tick_params(labelsize=9)

plt.tight_layout()
plt.show()

# -----------------------------
# Fig. 4
# -----------------------------
# -----------------------
# Load 
# -----------------------
with open(fig4_file, "rb") as f:
    points = pickle.load(f)


# -----------------------
# Figure setup
# -----------------------
fig, ax_flux = plt.subplots(figsize=(4.75, 3))
ax_alb = ax_flux.twinx()  # right y-axis for albedo

# Flux limits
ax_flux.set_ylim(-0.5, 7.5)
for spine in ax_flux.spines.values():
    spine.set_linewidth(0.5)
for spine in ax_alb.spines.values():
    spine.set_linewidth(0.5)

# Albedo limits
ax_alb.set_ylim(-0.001, 0.015)
ax_alb.set_yticks(np.arange(0, 0.016, 0.002))
ax_alb.set_yticklabels([f"{t:.3f}" for t in np.arange(0, 0.016, 0.002)])

# Plotting parameters
sz = 60
x_full = 0.75
x_mid  = 1.25
x_shift = 4
colors_clear = '#8b4513'
colors_shc   = '#1f77b4'
alpha_comp = 0.6
marker_main = "o"
markers_comp = {"year":"^"}

seasons = ["dry", "year"]

# -----------------------
# Plot Fsw & Flw
for i, field in enumerate(["Fsw","Flw"]):
    for xx, s_type in zip([x_full, x_mid], ["full","mid"]):
        for regime, color in zip(["shallow_clear","clear"], [colors_shc, colors_clear]):
            # Dry season
            val = points["dry"][field][regime][f"{s_type}_mean"]
            se  = points["dry"][field][regime][f"{s_type}_se"]
            ax_flux.errorbar([xx + i*2], [val], yerr=[2*se], fmt='none', linewidth=0.75,
                             ecolor='black', capsize=2, zorder=2)
            ax_flux.scatter([xx + i*2], [val], color=color, s=sz, edgecolors='black',
                            linewidth=0.75, marker=marker_main, zorder=3)

            # Full year
            val_y = points["year"][field][regime][f"{s_type}_mean"]
            se_y  = points["year"][field][regime][f"{s_type}_se"]
            ax_flux.errorbar([xx + i*2], [val_y], yerr=[2*se_y], fmt='none', linewidth=0.75,
                             ecolor='black', capsize=2, zorder=2)
            ax_flux.scatter([xx + i*2], [val_y], color=color, s=sz, edgecolors='black',
                            linewidth=0.75, marker=markers_comp["year"], alpha=alpha_comp, zorder=3)

# -----------------------
# Plot Albedo
for xx, s_type in zip([x_full, x_mid], ["full","mid"]):
    for regime, color in zip(["shallow_clear","clear"], [colors_shc, colors_clear]):
        # Dry season
        val = points["dry"]["albedo_toa"][regime][f"{s_type}_mean"]
        se  = points["dry"]["albedo_toa"][regime][f"{s_type}_se"]
        ax_alb.errorbar([xx + x_shift], [val], yerr=[2*se], fmt='none', linewidth=0.75,
                        ecolor='black', capsize=2, zorder=2)
        ax_alb.scatter([xx + x_shift], [val], color=color, s=sz, edgecolors='black',
                        linewidth=0.75, marker=marker_main, zorder=3)

        # Full year
        val_y = points["year"]["albedo_toa"][regime][f"{s_type}_mean"]
        se_y  = points["year"]["albedo_toa"][regime][f"{s_type}_se"]
        ax_alb.errorbar([xx + x_shift], [val_y], yerr=[2*se_y], fmt='none', linewidth=0.75,
                        ecolor='black', capsize=2, zorder=2)
        ax_alb.scatter([xx + x_shift], [val_y], color=color, s=sz, edgecolors='black',
                        linewidth=0.75, marker=markers_comp["year"], alpha=alpha_comp, zorder=3)

# -----------------------
# Labels & aesthetics
ax_flux.set_ylabel(r"$\Delta \rm{F_{SW,fcc}^\uparrow}$, $\Delta \rm{F_{LW,fcc}^\uparrow}$ (W m$^{-2}$)", fontsize=9)
ax_alb.set_ylabel(r"$\Delta \alpha_{fcc}$", fontsize=9)
ax_flux.axhline(0, color='black', linestyle='--', linewidth=1)
ax_alb.axhline(0, color='black', linestyle='--', linewidth=1)
ax_flux.grid(axis='y', linestyle='--', linewidth=0.7)
ax_flux.tick_params(axis='both', labelsize=7)
ax_alb.tick_params(axis='y', labelsize=7)

# -----------------------
# Legend
handles = [
    mlines.Line2D([], [], color=colors_shc, marker='o', linestyle='None', markersize=5, label='All-sky'),
    mlines.Line2D([], [], color=colors_clear, marker='o', linestyle='None', markersize=5, label='Surface'),
    mlines.Line2D([], [], color='black', marker='o', linestyle='None', markersize=5, label='Dry season'),
    mlines.Line2D([], [], color='black', marker='^', linestyle='None', markersize=5, alpha=alpha_comp, label='Full year')
]
leg = ax_flux.legend(handles=handles, loc='upper right', fontsize=6, frameon=True)
leg.get_frame().set_linewidth(0.5)

# -----------------------
# X-axis setup
xticks = [0.75, 1.25, 2.75, 3.25, 4.75, 5.25]
ax_flux.set_xticks(xticks)
ax_flux.set_xticklabels([r"$f_{loss}>0$", r"$f_{loss}\geq0.5$"]*3, fontsize=6, rotation=15)

pair_centers = [(0.75+1.25)/2, (2.75+3.25)/2, (4.75+5.25)/2]
for x_center, label in zip(pair_centers,
                           [r"$\Delta \rm{F_{SW,fcc}^\uparrow}$",
                            r"$\Delta \rm{F_{LW,fcc}^\uparrow}$",
                            r"$\Delta \alpha_{fcc}$"]):
    ax_flux.text(x_center, -0.11, label, ha='center', va='top', fontsize=9,
                 fontweight='bold', transform=ax_flux.get_xaxis_transform())

ax_flux.set_xlim(0, 6)
ax_flux.set_title("All-sky biophysical feedback", fontsize=9, pad=8)

plt.tight_layout()
plt.show()