#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)

#from itertools import product
#from cftime import DatetimeNoLeap
import os,os.path,sys
import xarray as xr


#ifile = 'HadISST_sst.187001-201910.wy.nc'
#ifile = 'HadISST_sst.187001-202001.wy.nc'
#ifile = 'HadISST_sst.187001-202012.wy.nc'
#ifile = 'HadISST_sst.187001-202105.wy.nc' #2021-08-05
#ifile = 'HadISST_sst.187001-202202.wy.nc' #2021-08-05
ifile = 'LMR2019_HadISST_sst_0800-2000.nc' #WY 2023-08-29
ofile = ifile.replace('.nc', '.forHIRAM.nc')
if os.path.exists(ofile): print('[exists]:', ofile); sys.exit()

data_name = 'sst'
data_name_new = 'sst'
xdim,ydim,tdim = 'lon','lat','time'
#nc_encoding = {tdim:{'dtype': 'double', '_FillValue': False},
#            xdim:{'_FillValue': False},
#            ydim:{'_FillValue': False},
#            data_name_new:{'_FillValue': False},
#            'yr': {'dtype': 'int32', '_FillValue': False},
#           'mo': {'dtype': 'int32', '_FillValue': False},
#            'dy': {'dtype': 'int32', '_FillValue': False},
#            'nrecords': {'dtype': 'int32', '_FillValue': False},
#            }
#{'dtype': 'double'} seems necessary for HiRAM
nc_encoding = {tdim:{'_FillValue': None, 'dtype': 'double'},
            xdim:{'_FillValue': None},
            ydim:{'_FillValue': None},
            data_name_new:{'_FillValue': None},
            'yr': {'dtype': 'int32', '_FillValue': None},
            'mo': {'dtype': 'int32', '_FillValue': None},
            'dy': {'dtype': 'int32', '_FillValue': None},
            'nrecords': {'dtype': 'int32', '_FillValue': None},
            }


# read ifile
ds = xr.open_dataset(ifile)
da = ds[data_name]


# copy the attributes from the old dataset
da.attrs = ds[data_name].attrs

yr = xr.DataArray(da.time.dt.year.values, dims=['idim'])
yr.attrs = dict(long_name='year')

mo = xr.DataArray(da.time.dt.month.values, dims=['idim'])
mo.attrs = dict(long_name='month')

dy = xr.DataArray(da.time.dt.day.values, dims=['idim'])
dy.attrs = dict(long_name='day')

nrecords = da.time.size

ds_new = xr.Dataset({data_name_new: da,
    'yr': yr,
    'mo': mo,
    'dy': dy,
    'nrecords': nrecords})
#ds_new = da.to_dataset(name=data_name_new)
ds_new.attrs = ds.attrs

# save to ofile
nc_encoding[tdim]['calendar'] = ds.time.encoding['calendar']
nc_encoding[tdim]['units'] = ds.time.encoding['units']
ds_new.to_netcdf(ofile, unlimited_dims=tdim, encoding=nc_encoding)
print('[saved]:', ofile)