# Read online NINO3.4 anomaly data and plot as a shaded red/blue time
# series. Created with chatGPT. Eli, 202506

import pandas as pd
import matplotlib.pyplot as plt
from urllib.request import urlopen
import matplotlib.ticker as mticker   # only once, at top

URL = "https://psl.noaa.gov/data/correlation/nina34.anom.data"

# ---------------------------------------------------------------------
# 1. Load and tidy the Niño-3.4 data
# ---------------------------------------------------------------------
def load_nino34(url=URL) -> pd.Series:
    text   = urlopen(url).read().decode("utf-8")
    toks   = text.split()
    first  = int(toks[0])            # 1948
    last   = int(toks[1])            # 2025
    data   = toks[2:]

    recs, i = [], 0
    for yr in range(first, last + 1):
        i += 1                       # skip the year token
        for month in range(1, 13):
            val = float(data[i]); i += 1
            if val < -90:            # −99.99 ⇒ missing
                continue
            recs.append({"Date": pd.Timestamp(yr, month, 15), "Anom": val})

    return (pd.DataFrame(recs)
              .set_index("Date")
              .sort_index()["Anom"])

s = load_nino34()

# ---------------------------------------------------------------------
# 2. Two equal halves: 1948-1986 and 1987-2025 (39 yr each)
# ---------------------------------------------------------------------
periods = [(1948, 1986), (1987, 2025)]

def decimal_year(idx):
    return idx.year + (idx.month - 1) / 12.0

# ---------------------------------------------------------------------
# 3. Plot: one figure, two stacked panels
# ---------------------------------------------------------------------
fig, axes = plt.subplots(
    nrows=2, ncols=1, figsize=(8,5), sharey=True, constrained_layout=True
)

for ax, (start, end) in zip(axes, periods):
    seg = s[str(start):str(end)]
    x   = decimal_year(seg.index)

    ax.plot(x, seg, lw=1)
    ax.axhline(0, color="black", lw=0.8)

    ax.fill_between(x, 0, seg.where(seg > 0, 0),
                    where=seg > 0, alpha=0.8, color="red")
    ax.fill_between(x, 0, seg.where(seg < 0, 0),
                    where=seg < 0, alpha=0.8, color="blue")

    # ----------------------------------------------------------------
    # 1) Major ticks & grid  (every 5 years in this example)
    # ----------------------------------------------------------------
    major_locator = mticker.MultipleLocator(5)
    ax.xaxis.set_major_locator(major_locator)
    ax.grid(True, which="major", linestyle="-", linewidth=0.25, alpha=1.0)

    # ----------------------------------------------------------------
    # 2) Minor ticks (every single year, no labels, no grid)
    # ----------------------------------------------------------------
    ax.set_xticks(range(start, end + 1), minor=True)
    ax.tick_params(axis="x", which="minor", length=4, labelbottom=False)

    ax.set_xlim(start, end + 1)
    ax.set_xticks(range(start, end + 1, 3 if (end - start) > 20 else 1))
    #ax.set_title(f"{start} – {end}", fontsize=11, pad=6)

axes[-1].set_xlabel("Year")
for ax in axes:
    ax.set_ylabel("Niño 3.4 (°C)")



## save as pdf:
f = plt.gcf()  # f = figure(n) if you know the figure number
f.set_size_inches(8, 5)
f.savefig("Output/Figure-nino3.4.pdf",format='pdf');

plt.show()