# -*- coding: utf-8 -*-
"""
Plot 3 figures (one per timescale) from a NetCDF file with:
  cha(timescale, dr_char, stn), lon(stn), lat(stn)

Each figure has 6 subplots (2x3) for the 6 drought characteristics:
  frq, dm, dx, sev, svx, nday

Compatible with environments lacking advanced Matplotlib styles.
"""

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

# ---------------------- USER SETTINGS ----------------------
input_file = "/home/tanpv/Progs/Dr_4Books/Res_new/Chars_Dr_Clim_spei_mil.nc"   # Path to your NetCDF file
output_pref  = "map_timescale"            # output file prefix
char_names   = ["frq", "dm", "dx", "sev", "svx", "nday"]
# colormap per characteristic (tweak if you like)
colormaps    = ["viridis", "viridis", "viridis", "RdBu_r", "RdBu_r", "viridis"]
marker_size  = 80
dpi_value    = 200
# -----------------------------------------------------------

# (safe) style setup
try:
    plt.style.use("seaborn-whitegrid")
except Exception:
    pass  # fall back to default

# open dataset (do not decode non-CF time)
ds = xr.open_dataset(input_file, decode_times=False)

# dimensions and variables
char_dim = "dr_char" if "dr_char" in ds.dims else ("char" if "char" in ds.dims else None)
time_dim = "timescale" if "timescale" in ds.dims else ("time" if "time" in ds.dims else None)
stn_dim  = "stn" if "stn" in ds.dims else None
if char_dim is None or time_dim is None or stn_dim is None:
    raise ValueError("Missing required dims: dr_char/char, timescale/time, or stn.")

if "cha" not in ds:
    raise ValueError("Variable 'cha' not found.")
if "lon" not in ds or "lat" not in ds:
    raise ValueError("Variables 'lon' and 'lat' are required.")

cha = ds["cha"]
lon = ds["lon"].values
lat = ds["lat"].values

# mask missing values (-99 or _FillValue)
arr = cha.values.astype(float)
fill_val = cha.attrs.get("_FillValue", None)
miss_val = cha.attrs.get("missing_value", None)
if fill_val is not None:
    arr[arr == fill_val] = np.nan
if miss_val is not None:
    arr[arr == miss_val] = np.nan
arr[arr == -99] = np.nan
cha = xr.DataArray(arr, coords=cha.coords, dims=cha.dims, attrs=cha.attrs)

n_char = cha.sizes[char_dim]
n_time = cha.sizes[time_dim]

# robust color limits per characteristic (ignore outliers)
vmin_list, vmax_list = [], []
for ic in range(n_char):
    data_c = cha.isel({char_dim: ic}).values
    vmin_list.append(np.nanpercentile(data_c, 2))
    vmax_list.append(np.nanpercentile(data_c, 98))

# map extent from station bounds (pad a bit)
lon_min, lon_max = float(np.nanmin(lon)), float(np.nanmax(lon))
lat_min, lat_max = float(np.nanmin(lat)), float(np.nanmax(lat))
pad_x = max(0.3, 0.05 * (lon_max - lon_min))
pad_y = max(0.3, 0.05 * (lat_max - lat_min))
extent = [lon_min - pad_x, lon_max + pad_x, lat_min - pad_y, lat_max + pad_y]

def plot_timescale(t_index: int) -> None:
    """Plot one figure (6 subplots) for a given timescale index."""
    fig, axes = plt.subplots(2, 3, figsize=(14, 8), dpi=dpi_value)
    fig.suptitle(f"Drought characteristics (timescale index = {t_index})",
                 fontsize=15, fontweight="bold", y=0.97)

    for ic in range(n_char):
        ax = axes.flat[ic]
        vals = cha.isel({time_dim: t_index, char_dim: ic}).values
        cmap = colormaps[ic % len(colormaps)]

        sc = ax.scatter(
            lon, lat, c=vals, s=marker_size,
            edgecolor="k", linewidths=0.3, alpha=0.9,
            cmap=cmap, vmin=vmin_list[ic], vmax=vmax_list[ic]
        )
        ax.set_title(char_names[ic] if ic < len(char_names) else f"char_{ic}", fontsize=12, weight="bold")
        ax.set_xlim(extent[0], extent[1])
        ax.set_ylim(extent[2], extent[3])
        ax.set_xlabel("Longitude", fontsize=10)
        ax.set_ylabel("Latitude", fontsize=10)
        ax.grid(True, ls="--", alpha=0.3)

        cbar = plt.colorbar(sc, ax=ax, shrink=0.85, pad=0.03)
        if "units" in cha.attrs and cha.attrs["units"]:
            cbar.set_label(cha.attrs["units"], fontsize=9)
        cbar.ax.tick_params(labelsize=9)

    # hide unused panels if n_char < 6
    for j in range(n_char, len(axes.flat)):
        axes.flat[j].axis("off")

    plt.tight_layout(rect=[0, 0, 1, 0.95])
    out_path = f"{output_pref}_{t_index}.png"
    plt.savefig(out_path, dpi=dpi_value, bbox_inches="tight")
    plt.close(fig)
    print(f"Saved: {out_path}")

# plot all timescales
for t in range(n_time):
    plot_timescale(t)