"""
gulf_stream_section.py – Contour and scatter plots for Gulf Stream glider
(glider 25B20901sat, profiles 41-91).
"""
import numpy as np
import scipy.io as sio
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import matplotlib.dates as mdates
from pathlib import Path

# ── Settings ──────────────────────────────────────────────────────────────────
DATA_FILE  = r"C:\Users\bwerb\Documents\ProposalFigures\data\25B20901sat.mat"
OUT_DIR    = Path(r"C:\Users\bwerb\Documents\ProposalFigures\gulf_stream_section")
OUT_DIR.mkdir(exist_ok=True)

PROF_IDX   = slice(40, 91)   # profiles 41-91 (0-indexed: 40–90 inclusive)
BIN_SIZE   = 50              # m
GAP_THRESH = 2               # days
DEPTH_MAX  = 1000            # m
FIG_SIZE   = (13.33, 7.5)
DPI        = 300

# ── MATLAB datenum → matplotlib date number ───────────────────────────────────
def mat2mpl(dn):
    return np.asarray(dn, float) - 366.0

# ── Load ──────────────────────────────────────────────────────────────────────
raw = sio.loadmat(DATA_FILE, squeeze_me=True, struct_as_record=False)
s   = raw['s']

sdn = s.sdn[PROF_IDX]
dep = s.depth[:, 0] if s.depth.ndim == 2 else s.depth  # depth vector (m)

FIELDS = ['pHin', 'doxy', 'tc', 'psal', 'pHin_esper', 'pHin_canb', 'pco2in', 'dic_canb']
raw_data = {}
for f in FIELDS:
    arr = getattr(s, f)
    raw_data[f] = arr[:, PROF_IDX] if arr.ndim == 2 else arr[PROF_IDX]

n_profs = len(sdn)

# ── Depth binning (50 m) ──────────────────────────────────────────────────────
bin_edges = np.arange(0, DEPTH_MAX + BIN_SIZE, BIN_SIZE)
bin_ctrs  = bin_edges[:-1] + BIN_SIZE / 2
n_bins    = len(bin_ctrs)

binned = {}
for f, arr in raw_data.items():
    if arr.ndim != 2:
        continue
    out = np.full((n_bins, n_profs), np.nan)
    for b in range(n_bins):
        mask = (dep >= bin_edges[b]) & (dep < bin_edges[b + 1])
        if mask.any():
            with np.errstate(all='ignore'):
                out[b] = np.nanmean(arr[mask], axis=0)
    binned[f] = out

sdn_mpl = mat2mpl(sdn)

# ── Insert NaN columns at time gaps ──────────────────────────────────────────
def add_gaps(t, Z, thresh=GAP_THRESH):
    t, Z = t.copy(), Z.copy()
    gaps = np.flatnonzero(np.diff(t) > thresh)
    for i in reversed(gaps):
        t_mid = (t[i] + t[i + 1]) / 2
        t = np.insert(t, i + 1, t_mid)
        Z = np.insert(Z, i + 1, np.nan, axis=1)
    return t, Z

# ── Contour plot helper ───────────────────────────────────────────────────────
def save_contour(t, Z, dep_ctrs, cmap, title, cbar_label, fname,
                 clim=None, norm=None):
    t_p, Z_p = add_gaps(t, Z)
    kw = dict(cmap=cmap)
    if norm is not None:
        kw['norm'] = norm
    elif clim is not None:
        kw['vmin'], kw['vmax'] = clim

    fig, ax = plt.subplots(figsize=FIG_SIZE)
    cf = ax.contourf(t_p, dep_ctrs, Z_p, 50, **kw)
    ax.set_ylim(DEPTH_MAX, 0)
    ax.xaxis_date()
    ax.xaxis.set_major_formatter(mdates.DateFormatter('%b %d %Y'))
    ax.xaxis.set_major_locator(mdates.AutoDateLocator())
    plt.setp(ax.get_xticklabels(), rotation=20, ha='right')
    ax.set_ylabel('Depth (m)', fontsize=12)
    ax.set_title(title, fontsize=14)
    ax.grid(True, alpha=0.3)

    pos = ax.get_position()
    cb  = fig.colorbar(cf, ax=ax, pad=0.01)
    cb.set_label(cbar_label, fontsize=10)
    ax.set_position(pos)

    fig.tight_layout()
    path = OUT_DIR / fname
    fig.savefig(path, dpi=DPI, bbox_inches='tight')
    plt.close(fig)
    print(f'Saved  {path.name}')

# ── Diverging (red-blue) contour helper ───────────────────────────────────────
def save_diverging(t, Z, dep_ctrs, title, cbar_label, fname):
    finite = Z[np.isfinite(Z)]
    vmax   = np.percentile(np.abs(finite), 98) if finite.size else 0.05
    vmax   = max(vmax, 1e-3)
    norm   = mcolors.TwoSlopeNorm(vmin=-vmax, vcenter=0.0, vmax=vmax)
    save_contour(t, Z, dep_ctrs, 'RdBu_r', title, cbar_label, fname, norm=norm)

# ── Scatter plot helper ───────────────────────────────────────────────────────
def save_scatter(Z_raw, dep_vec, xlabel, title, fname):
    dep_flat = np.tile(dep_vec[:, None], (1, Z_raw.shape[1])).ravel()
    z_flat   = Z_raw.ravel()
    mask     = np.isfinite(dep_flat) & np.isfinite(z_flat) & (dep_flat <= DEPTH_MAX)

    fig, ax = plt.subplots(figsize=(6, 8))
    ax.scatter(z_flat[mask], dep_flat[mask], s=3, alpha=0.4,
               color='steelblue', rasterized=True)
    ax.axvline(0, color='k', lw=0.8, ls='--')
    ax.set_xlabel(xlabel, fontsize=12)
    ax.set_ylabel('Depth (m)', fontsize=12)
    ax.set_ylim(DEPTH_MAX, 0)
    ax.set_title(title, fontsize=13)
    ax.grid(True, alpha=0.3)
    fig.tight_layout()
    path = OUT_DIR / fname
    fig.savefig(path, dpi=DPI, bbox_inches='tight')
    plt.close(fig)
    print(f'Saved  {path.name}')

# ═══════════════════════════════════════════════════════════════════════════════
# Contour plots
# ═══════════════════════════════════════════════════════════════════════════════

save_contour(sdn_mpl, binned['pHin'], bin_ctrs, 'turbo',
             r'Gulf Stream – pH$_{in\ situ}$',
             r'pH$_{in\ situ}$ [total]', 'pHin_section.png',
             clim=[7.85, 8.05])

save_contour(sdn_mpl, binned['doxy'], bin_ctrs, 'viridis',
             r'Gulf Stream – Dissolved Oxygen',
             r'DO ($\mu$mol kg$^{-1}$)', 'doxy_section.png',
             clim=[0, 300])

save_contour(sdn_mpl, binned['tc'], bin_ctrs, 'RdYlBu_r',
             r'Gulf Stream – Temperature',
             r'Temperature ($^{\circ}$C)', 'tc_section.png',
             clim=[4, 30])

save_contour(sdn_mpl, binned['psal'], bin_ctrs, 'viridis',
             r'Gulf Stream – Salinity',
             r'Salinity (PSS)', 'psal_section.png',
             clim=[33, 37])

save_contour(sdn_mpl, binned['pco2in'], bin_ctrs, 'plasma',
             r'Gulf Stream – p$CO_2$',
             r'p$CO_2$ ($\mu$atm)', 'pco2_section.png',
             clim=[200, 700])

save_contour(sdn_mpl, binned['dic_canb'], bin_ctrs, 'viridis',
             r'Gulf Stream – DIC (CANYON-B)',
             r'DIC ($\mu$mol kg$^{-1}$)', 'dic_section.png',
             clim=[1800, 2200])

# ΔpH diverging contour
delta_esper_bin = binned['pHin'] - binned['pHin_esper']
save_diverging(sdn_mpl, delta_esper_bin, bin_ctrs,
               r'Gulf Stream – $\Delta$pH (measured $-$ ESPER)',
               r'$\Delta$pH', 'deltapH_esper_section.png')

# ═══════════════════════════════════════════════════════════════════════════════
# Scatter plots: depth vs ΔpH (raw, un-binned)
# ═══════════════════════════════════════════════════════════════════════════════

delta_esper_raw = raw_data['pHin'] - raw_data['pHin_esper']
delta_canb_raw  = raw_data['pHin'] - raw_data['pHin_canb']

save_scatter(delta_esper_raw, dep,
             r'pH$_{in\ situ}$ $-$ pH$_{ESPER}$',
             r'Gulf Stream – $\Delta$pH vs Depth (ESPER)',
             'scatter_deltapH_esper.png')

save_scatter(delta_canb_raw, dep,
             r'pH$_{in\ situ}$ $-$ pH$_{CANYON-B}$',
             r'Gulf Stream – $\Delta$pH vs Depth (CANYON-B)',
             'scatter_deltapH_canb.png')

print(f'\nAll figures saved to: {OUT_DIR}')