#!/usr/bin/env python

#;DIAG_NORCPM; INPUTFILENAME: input_file_not_set.nc
#;DIAG_NORCPM; OUTPUTPATTERN: 

## For DJF, 1st data will be mean 1st year Dec and 2nd year Jan Feb
##          last data will be Dec only

import xarray as xr
from os.path import exists
from sys import exit

infile  = "TBI_CoEx_PACE_P_ANOM_PRECT_ensmean_1980-2021.nc"
outpat  = "TBI_CoEx_PACE_P_ANOM_PRECT_ensmean_SEASON_1980-2021.nc" ## 'SEASAON' will be replaced as JJA...etc.

if not outpat: outpat = infile.replace('.nc','')+'_SEASON.nc'

JJA=('JJA',[6,7,8])
SON=('SON',[9,10,11])
DJF=('DJF',[12,1,2])
MAM=('MAM',[3,4,5])
seasons = [JJA,SON,DJF,MAM]

ds = xr.open_dataset(infile)
ds = ds.resample(time='QS-DEC').mean('time') ## without day weighting

for s in seasons:
    name,mons = s[0], s[1]
    ofn = outpat.replace('SEASON',name)
    if ofn == outpat: 
        print(f"ERROR: outpat = {outpat}, need contain string 'SEASON'")
        exit()
    ds_season = ds.sel(time=ds.time.dt.month == mons[0])
    ds_season.to_netcdf(ofn)
    print("Seasonal file: "+ofn)