import glob
import os
import re

import numpy as np
import pandas as pd
import xarray as xr
import matplotlib.pyplot as plt

indir = "/global/cfs/cdirs/e3sm/inputdata/atm/datm7/WIEMIP/UKESM/TPHWL6Hrly"
start_year, end_year = 1980, 2000

files = sorted(glob.glob(os.path.join(indir, "UKESM.TPQWL.*.nc")))

dates, values = [], []
for f in files:
    m = re.search(r"TPQWL\.(\d{4})-(\d{2})\.nc$", f)
    if m is None:
        continue
    yr, mon = int(m.group(1)), int(m.group(2))
    if not (start_year <= yr <= end_year):
        continue

    # Take the date from the filename, so time decoding is not needed
    with xr.open_dataset(f, decode_times=False) as ds:
        tbot = ds["TBOT"]                    # (time, lat, lon), K
        tbot = tbot.where(tbot < 1e30)       # mask fill values

        # Latitude: 1D "lat", or 2D "LATIXY" for some DATM files
        lat = ds["lat"] if "lat" in ds else ds["LATIXY"]
        w = np.cos(np.deg2rad(lat))

        space_dims = [d for d in tbot.dims if d != "time"]
        gm = tbot.weighted(w).mean(dim=space_dims)   # global mean at each time step
        monthly = float(gm.mean("time"))             # monthly mean of the global mean

    dates.append(pd.Timestamp(yr, mon, 15))
    values.append(monthly - 273.15)                  # K to °C

    if len(values) % 120 == 0:
        print(f"Processed through {yr}-{mon:02d}")

ts = pd.Series(values, index=pd.DatetimeIndex(dates), name="TBOT_global_C").sort_index()
ts.to_csv("tbot_global_mean_monthly_1900_2000.csv")

# Annual mean, weighted by days in each month
days = ts.index.days_in_month
annual = (ts * days).groupby(ts.index.year).sum() / days.to_series(index=ts.index).groupby(ts.index.year).sum()

# Plot
fig, ax = plt.subplots(figsize=(10, 4))
ax.plot(ts.index, ts.values, lw=0.6, alpha=0.5, label="Monthly")
ax.plot(pd.to_datetime(annual.index.astype(str) + "-07-01"), annual.values,
        lw=2, label="Annual mean")
ax.set_xlabel("Year")
ax.set_ylabel("Global mean TBOT (°C)")
ax.set_title("UKESM forcing: global mean air temperature (TBOT), 1900-2000")
ax.grid(alpha=0.3)
ax.legend()
plt.tight_layout()
plt.savefig("tbot_global_mean_1900_2000.png", dpi=200)
plt.show()