#!/usr/bin/env python3
"""Download groups for ECMWF S2S reforecasts at native 1.5 deg (ECDS).

Variable set reproduces the 40-variable table of HW_HPC_1_cla.docx (Bang 1),
plus a large-scale domain used to build EOF/PC predictors (feature set F2/F3
of the previous HW session).

Grouping is dictated by the dataset constraints (verified against
catalogue/v1/collections/s2s-reforecasts/constraints.json on 2026-09-03):
  - "Daily averaged" variables use 24h WINDOW leadtimes (0_24 ... 1080_1104), 46 values
  - Instantaneous/accumulated variables use POINT leadtimes
  - mx2t6/mn2t6 exist only on the 6-hourly point grid (6 ... 1104), 184 values
  - pressure fields: key is `level_value` with "<p>_hpa" values, `level_type` = "pressure"
"""

DATASET = "s2s-reforecasts"
API_URL = "https://ecds.ecmwf.int/api"

MODEL_VERSION_YEAR = "2023"
HYEARS = [str(y) for y in range(2003, 2023)]      # 20 hindcast years in ONE request
GRID = "1.5/1.5"

# cdsapi area order: [North, West, South, East]
AREA_VN = [25, 100, 6, 114]          # domain A - Vietnam
AREA_WIDE = [45, 80, -15, 150]       # domain B - large scale (EOF predictors)

PLEVS = ["1000_hpa", "925_hpa", "850_hpa", "700_hpa", "500_hpa", "300_hpa", "200_hpa"]
PLEVS_WIDE = ["850_hpa", "500_hpa", "200_hpa"]

LT_WINDOW_24 = [f"{i*24}_{(i+1)*24}" for i in range(46)]   # daily-averaged, 0_24..1080_1104
LT_POINT_24 = [str(i * 24) for i in range(47)]             # 0 .. 1104
LT_POINT_06 = [str(i * 6) for i in range(1, 185)]          # 6 .. 1104 (mx2t6/mn2t6)

# short name each requested variable maps to in the resulting GRIB/NetCDF
SHORTNAME = {
    "2_m_temperature": "t2m", "2_m_dewpoint_temperature": "d2m",
    "skin_temperature": "skt", "convective_available_potential_energy": "cape",
    "total_column_water": "tcw", "total_cloud_cover": "tcc",
    "soil_moisture_top_20_cm": "sm20", "soil_moisture_top_100_cm": "sm100",
    "soil_temperature_top_20_cm": "st20", "soil_temperature_top_100_cm": "st100",
    "surface_sensible_heat_flux": "sshf", "surface_latent_heat_flux": "slhf",
    "mean_sea_level_pressure": "msl", "surface_pressure": "sp",
    "10_m_u_component_of_wind": "u10", "10_m_v_component_of_wind": "v10",
    "surface_solar_radiation_downwards": "ssrd",
    "surface_thermal_radiation_downwards": "strd",
    "surface_net_solar_radiation": "ssr", "surface_net_thermal_radiation": "str",
    "top_net_thermal_radiation": "ttr",
    "convective_precipitation": "cp", "total_precipitation": "tp",
    "maximum_2_m_temperature_in_the_last_6_hours": "mx2t6",
    "minimum_2_m_temperature_in_the_last_6_hours": "mn2t6",
    "geopotential_height": "gh", "temperature": "t", "specific_humidity": "q",
    "u_component_of_wind": "u", "v_component_of_wind": "v",
    "vertical_velocity": "w",
}

GROUPS = {
    "srf_daily": dict(
        level_type="single_level",
        variable=["2_m_temperature", "2_m_dewpoint_temperature", "skin_temperature",
                  "convective_available_potential_energy", "total_column_water",
                  "total_cloud_cover", "soil_moisture_top_20_cm",
                  "soil_moisture_top_100_cm", "soil_temperature_top_20_cm",
                  "soil_temperature_top_100_cm"],
        leadtime_hour=LT_WINDOW_24, area=AREA_VN,
    ),
    "srf_24h": dict(
        level_type="single_level",
        variable=["surface_sensible_heat_flux", "surface_latent_heat_flux",
                  "mean_sea_level_pressure", "surface_pressure",
                  "10_m_u_component_of_wind", "10_m_v_component_of_wind",
                  "surface_solar_radiation_downwards",
                  "surface_thermal_radiation_downwards",
                  "surface_net_solar_radiation", "surface_net_thermal_radiation",
                  "top_net_thermal_radiation", "convective_precipitation",
                  "total_precipitation"],
        leadtime_hour=LT_POINT_24, area=AREA_VN,
    ),
    "txtm": dict(
        level_type="single_level",
        variable=["maximum_2_m_temperature_in_the_last_6_hours",
                  "minimum_2_m_temperature_in_the_last_6_hours"],
        leadtime_hour=LT_POINT_06, area=AREA_VN,
    ),
    "plev": dict(
        level_type="pressure",
        variable=["geopotential_height", "temperature", "specific_humidity",
                  "u_component_of_wind", "v_component_of_wind", "vertical_velocity"],
        level_value=PLEVS,
        leadtime_hour=LT_POINT_24, area=AREA_VN,
    ),
    "wide_plev": dict(
        level_type="pressure",
        variable=["geopotential_height", "temperature",
                  "u_component_of_wind", "v_component_of_wind"],
        level_value=PLEVS_WIDE,
        leadtime_hour=LT_POINT_24, area=AREA_WIDE,
    ),
    "wide_olr": dict(
        level_type="single_level",
        variable=["top_net_thermal_radiation"],
        leadtime_hour=LT_POINT_24, area=AREA_WIDE,
    ),
}


def build_request(group, mv_month, mv_day):
    """One ECDS request: one model-version date x all 20 hindcast years."""
    g = GROUPS[group]
    req = {
        "origin": "ecmwf",
        "year": MODEL_VERSION_YEAR,
        "month": mv_month,
        "day": mv_day,
        "time": "00:00",
        "hyear": list(HYEARS),
        "hmonth": [mv_month],
        "hday": [mv_day],
        "forecast_type": "control_forecast",
        "data_format": "grib",
        "grid": GRID,
    }
    req.update(g)
    return req


def load_mvdates(path="/home/tanpv/Data/QG.25.196/hdate_2023.txt"):
    """Model-version dates (mm-dd) of the 2023 ECMWF cycle: 2 runs/week -> 104 dates."""
    with open(path) as fh:
        toks = fh.read().split()
    return sorted({t for t in toks if len(t) == 5 and t[2] == "-"})