# -*-coding:utf-8 -*-
'''
@Author: Haoran Yu
@Desc: 极端气候指标计算
'''

import ETCCDI
import xarray as xr
import os
import numpy as np
from tqdm import tqdm

# 参数
CMIP = 'HiCPC_pr'
MODE = 'ACCESS-ESM1-5'
SSP = ['ssp126', "ssp245", 'ssp370', 'ssp585']
BASELINE_YEARS = (1979, 2014)

# 指标函数
INDICATORS = {
    'PRCPTOT': ETCCDI.PRCPTOT,
    'SDII': ETCCDI.SDII,
    'RX1DAY': ETCCDI.RX1DAY,
    'RX5DAY': ETCCDI.RX5DAY,
    'CDD': ETCCDI.CDD,
    'CWD': ETCCDI.CWD,
    'R1mm_days': ETCCDI.R1mm_days,
    'R10mm_precipitation': ETCCDI.R10mm_precipitation,
    'R10mm_days': ETCCDI.R10mm_days,
    'R10mm_intensity': ETCCDI.R10mm_intensity,
    'R10mm_ratio': ETCCDI.R10mm_ratio,
    'R20mm_precipitation': ETCCDI.R20mm_precipitation,
    'R20mm_days': ETCCDI.R20mm_days,
    'R20mm_intensity': ETCCDI.R20mm_intensity,
    'R20mm_ratio': ETCCDI.R20mm_ratio,
    'R25mm_precipitation': ETCCDI.R25mm_precipitation,
    'R25mm_days': ETCCDI.R25mm_days,
    'R25mm_intensity': ETCCDI.R25mm_intensity,
    'R25mm_ratio': ETCCDI.R25mm_ratio,
    'R50mm_precipitation': ETCCDI.R50mm_precipitation,
    'R50mm_days': ETCCDI.R50mm_days,
    'R50mm_intensity': ETCCDI.R50mm_intensity,
    'R50mm_ratio': ETCCDI.R50mm_ratio,
    'R90mm_precipitation': lambda pr, b: ETCCDI.R90mm_precipitation(pr, b, baseline_years=BASELINE_YEARS),
    'R90mm_ratio': lambda pr, b: ETCCDI.R90mm_ratio(pr, b, baseline_years=BASELINE_YEARS),
    'R90mm_days': lambda pr, b: ETCCDI.R90mm_days(pr, b, baseline_years=BASELINE_YEARS),
    'R90mm_intensity': lambda pr, b: ETCCDI.R90mm_intensity(pr, b, baseline_years=BASELINE_YEARS),
    'R95mm_precipitation': lambda pr, b: ETCCDI.R95mm_precipitation(pr, b, baseline_years=BASELINE_YEARS),
    'R95mm_ratio': lambda pr, b: ETCCDI.R95mm_ratio(pr, b, baseline_years=BASELINE_YEARS),
    'R95mm_days': lambda pr, b: ETCCDI.R95mm_days(pr, b, baseline_years=BASELINE_YEARS),
    'R95mm_intensity': lambda pr, b: ETCCDI.R95mm_intensity(pr, b, baseline_years=BASELINE_YEARS),
    'R99mm_precipitation': lambda pr, b: ETCCDI.R99mm_precipitation(pr, b, baseline_years=BASELINE_YEARS),
    'R99mm_ratio': lambda pr, b: ETCCDI.R99mm_ratio(pr, b, baseline_years=BASELINE_YEARS),
    'R99mm_days': lambda pr, b: ETCCDI.R99mm_days(pr, b, baseline_years=BASELINE_YEARS),
    'R99mm_intensity': lambda pr, b: ETCCDI.R99mm_intensity(pr, b, baseline_years=BASELINE_YEARS),
}

BASELINE_REQUIRED_INDICATORS = {
    'R90mm_precipitation', 'R90mm_ratio', 'R90mm_days', 'R90mm_intensity',
    'R95mm_precipitation', 'R95mm_ratio', 'R95mm_days', 'R95mm_intensity',
    'R99mm_precipitation', 'R99mm_ratio', 'R99mm_days', 'R99mm_intensity'
}

for s in SSP:
    print(f"\n处理: {s}")

    f = f"./Data/{MODE}/{CMIP}_{MODE}_{s}.nc"
    hist = xr.open_dataset(f).sel(time=slice(str(BASELINE_YEARS[0]), str(BASELINE_YEARS[1]))).sel(lat=slice(35,43), lon=slice(112, 120))
    fut = xr.open_dataset(f).sel(time=slice('2015', '2100')).sel(lat=slice(35,43), lon=slice(112, 120))
    baseline = None

    for name, data in [('historical', hist), ('future', fut)]:
        print(f"  {name}")

        for var in data.data_vars:
            if str(var).lower() in ['pr', 'pre', 'precipitation']:
                pr = data[var].load()
                break
        else:
            continue

        rename = {}
        for d in pr.dims:
            dim_name = str(d).lower()
            if 'lat' in dim_name:
                rename[d] = 'lat'
            elif 'lon' in dim_name:
                rename[d] = 'lon'
        if rename:
            pr = pr.rename(rename)

        if name == 'historical' and baseline is None:
            y = pr.time.dt.year
            baseline = pr.where((y >= BASELINE_YEARS[0]) & (y <= BASELINE_YEARS[1]), drop=True)

        for ind, func in tqdm(INDICATORS.items(), desc="指标"):
            if ind in BASELINE_REQUIRED_INDICATORS:
                res = func(pr, baseline)
            else:
                res = func(pr)
            if res is None:
                continue

            if 'year' in res.coords:
                years = sorted(np.asarray(res['year'].values).astype(int).tolist())
            elif 'time' in res.dims:
                years = sorted(np.unique(res.time.dt.year.values).astype(int).tolist())
            else:
                years = []

            years_label = "unknown"
            if years:
                years_label = str(years[0]) if len(years) == 1 else f"{years[0]}-{years[-1]}"
            p = f"./Output/Indices data/{MODE}/{s}/{MODE}-{s}-{ind}-{years_label}.nc"
            os.makedirs(os.path.dirname(p), exist_ok=True)
            res.to_dataset(name=ind).to_netcdf(p)


