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



yr = int(sys.argv[1])
mn = int(sys.argv[2])
output_dir = "/home/disk/rocinante/DATA/temp/TNC/imerg/data/"


urllist = 'imerg.v6B_urls.csv'
udf = pd.read_csv(urllist, parse_dates=['time'])

df = udf[(udf.time.dt.year == yr) & (udf.time.dt.month == mn)]
days = np.unique(df.time.dt.day)


#urls = list(df.url.values)
#for dy in days:    
#
#    data = list(df[df.time.dt.day == dy].url.values)
#    print(len(data))
#    dxs = []
#    for d in data:
#        #print(d)
#        dx = xr.open_dataset(d)
#        dxs = dxs + [dx]
#        
#    out = "{}/IMERG_{}_{:02}_{:02}.nc".format(odir, yr, mn, dy)        
#    print('\t', out)
#    ds = xr.concat(dxs, dim='time')
#    ds.to_netcdf(out)
#    dx.close()
#    ds.close()
#

for dy in days:
    df = df[df.time.dt.day == dy]

    for i, row in df.iterrows():
        odir = "{}/{}_{:02}/{:02}/".format(output_dir, yr, mn, dy)

        hr = row.time.hour
        mm = row.time.minute
        ds = xr.open_dataset(row.url)
        out = "{}/IMERG_{}-{:02}-{:02}_{:02}_{:02}.nc".format(odir, yr, mn, dy, hr, mm)
        os.makedirs(odir, exist_ok=True)

        print(out)
        ds.to_netcdf(out)




# Process url list
#df['time'] = df.apply(lambda x: x.url.split('3IMERG.')[1][0:16], axis=1)
#df.index = pd.to_datetime(df.time, format='%Y%m%d-s%H%M%S')