#!/usr/bin/env python
#wy2024-03-11: modified from wydo_era5_t2m_subregion.py
if __name__ == '__main__':
    from misc.timer import Timer
    wyt = Timer(f'start {__file__}') 
import xarray as xr
import os, os.path, sys, time
import  salem
import pandas as pd
import numpy as np
import geoxarray

cntry_name = 'Japan'
cntry_sname = cntry_name.replace(' ', '')
data_name = 'd2m' #'RH2m' #'t2m'
daname_long = '2m_dewpoint_temperature' # '2m_temperature'
xname = 'longitude'
yname = 'latitude'
shdf = salem.read_shapefile('gadm36_JPN_shp/gadm36_JPN_1.shp')
name_col = 'NAME_1'
#t = pd.date_range('2013-01', '2020-02', freq='MS')
t = pd.date_range('1982-01', '2020-12', freq='MS')
print(data_name)
for year,month in zip(t.year, t.month):
    print(f'{year:04d}-{month:02d}')
    ifile = f'/tigress/wenchang/data/era5/analysis/{daname_long}/daily/era5.{daname_long}.daily.{year:04d}-{month:02d}.nc'
    ofile = f'tmp/era5.{data_name}.daily.{cntry_sname}.subregion.{year:04d}-{month:02d}.nc'
    if os.path.exists(ofile):
        print('[exists]:', ofile)
        continue
    ds = salem.open_xr_dataset(ifile)
    # transform lon [0,360) -> [-180,180)
    ds = ds.roll(longitude=ds.longitude.size//2, roll_coords=True)
    lon_new = lon = ds.longitude.where(ds.longitude<180, other=ds.longitude-360).values
    ds = ds.assign_coords(longitude=lon_new)
    
    das = []
    #region_IDs = []
    region_names = []
    N = len(shdf)
    for i in range(N):
        r = shdf.iloc[i, :]
        lon_slice = slice(r.min_x, r.max_x)
        lat_slice = slice(r.min_y, r.max_y)
        g = r.geometry
     
        try:
            da = ds[data_name].salem.roi(geometry=g).geo.fldmean()
            das.append(da)
            #region_IDs.append(int(r.GEOID))
            region_names.append(r[name_col])
            print('[OK]:', i+1,'of', N, r[name_col])
        except:
            print('\t[Failed]', i+1,'of', N, r[name_col])
            pass
    
    i_record = range(len(das))
    da = xr.concat(das, pd.Index(region_names, name='region_name')).transpose()
    da.attrs = ds[data_name].attrs
    ds_new = da.to_dataset(name=data_name)
    ds_new.attrs = ds.attrs
    ds_new.attrs['region_name'] = '; '.join(region_names)
    
    #scodes = xr.DataArray(scodes, dims='i_county', coords=[i_county,])
    #ccodes = xr.DataArray(ccodes, dims='i_county', coords=[i_county,])
    #xr.Dataset(dict(precip=da, scode=scodes, ccode=ccodes)).to_netcdf(f'pr{year}.nc', unlimited_dims='time')
    ds_new.to_netcdf(ofile, unlimited_dims='time', encoding={data_name: {'_FillValue': None}})
    
    print('[saved]:', ofile)

if __name__ == '__main__':
    wyt.check('**done**')