#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import xarray as xr
import pandas as pd
import numpy as np
import os

# ==========================================================
# CONFIGURAÇÕES
# ==========================================================
DATA_DIR = "/share/ioper/tempo/MERGE/GPM/CLIMATOLOGY/MONTHLY_ACCUMULATED_YEARLY"
OUT_DIR  = "/share/ioper/tempo/MERGE/GPM/CLIMATOLOGY/ANOMALY_ACCUMULATED_YEARLY"

os.makedirs(OUT_DIR, exist_ok=True)

MONTHS = {
    "jan": 1, "feb": 2, "mar": 3, "apr": 4,
    "may": 5, "jun": 6, "jul": 7, "aug": 8,
    "sep": 9, "oct": 10, "nov": 11, "dec": 12
}

YEARS = range(1998, 2026)

VAR = "pacum"
FILL_VALUE = -999000000.0

# ==========================================================
# LEITURA DOS DADOS (tempo REAL)
# ==========================================================
datasets = []

for year in YEARS:
    for mon_str, mon_num in MONTHS.items():
        fname = f"{DATA_DIR}/MERGE_CPTEC_acum_{mon_str}_{year}.nc"
        if not os.path.exists(fname):
            continue

        # data representativa do mês
        time_coord = pd.Timestamp(year=year, month=mon_num, day=15)

        ds = xr.open_dataset(fname)

        # força time correto
        ds = ds.assign_coords(time=[time_coord])

        # trata FillValue
        if "_FillValue" in ds[VAR].attrs:
            fv = ds[VAR].attrs["_FillValue"]
        else:
            fv = FILL_VALUE

        ds[VAR] = ds[VAR].where(ds[VAR] != fv)

        datasets.append(ds)

# concatena tudo
ds_all = xr.concat(datasets, dim="time")

# ==========================================================
# CLIMATOLOGIA MENSAL
# ==========================================================
climatology = ds_all.groupby("time.month").mean("time", skipna=True)

# ==========================================================
# ANOMALIA MENSAL
# ==========================================================
anomaly = ds_all.groupby("time.month") - climatology

# ==========================================================
# GRAVA OS ARQUIVOS (estrutura idêntica)
# ==========================================================
for year in YEARS:
    for mon_str, mon_num in MONTHS.items():
        fname_in = f"{DATA_DIR}/MERGE_CPTEC_acum_{mon_str}_{year}.nc"
        if not os.path.exists(fname_in):
            continue

        ds_orig = xr.open_dataset(fname_in)

        time_sel = pd.Timestamp(year=year, month=mon_num, day=15)

        anom_sel = anomaly.sel(time=time_sel)

        ds_out = ds_orig.copy()

        # escreve dados (mantém time, lat, lon)
        ds_out[VAR][0, :, :] = anom_sel[VAR].values

        # atributos
        ds_out[VAR].attrs["long_name"] = "Monthly precipitation anomaly"

        # evita conflito de FillValue
        ds_out[VAR].attrs.pop("_FillValue", None)
        ds_out[VAR].encoding["_FillValue"] = FILL_VALUE

        out_name = f"{OUT_DIR}/MERGE_CPTEC_anom_{mon_str}_{year}.nc"
        ds_out.to_netcdf(out_name)

print("✅ Anomalias mensais calculadas e gravadas com sucesso.")