#!/usr/bin/env python

'''
Script to regrid the MODIS BNU LAI data from 0.1 x 0.1 degree resolution.
to 0.25 x 0.25 degree resolution.  This will allow us to split the LAI
into 73 different variables (one per land type) so that HEMCO can read
and regrid the data.
'''

# imports
import os
import numpy as np
import xarray as xr
import xesmf as xe
from datetime import datetime

# Today's date
today = datetime.now().strftime('%Y-%m-%d')

# Main data directory
maindatadir = '/n/holyscratch01/jacob_lab/msulprizio/GC/data/MODIS_LAI'

# Define the lon and lat arrays for the output grid: 0.25 x 0.25
lat_out = np.arange(-89.875, 90.0, 0.25)
lon_out = np.arange(-179.875, 180.0, 0.25)

# Loop over years
for yyyy in range(2001,2021):

    # -----------------------
    # Get input data
    # -----------------------
    
    # Input file
    infile = os.path.join(maindatadir,
                          'MODIS.LAI.BNUv6.generic.01x01.{}.nc'.format(yyyy))


    # Read data from input file
    print('Reading {}'.format(infile))
    ds = xr.open_dataset(infile)

    # --------------------
    # Regrid data
    # --------------------  
          
    # Attributes for new Dataset
    attrs = {'title': ds.attrs['title'],
             'institution': ds.attrs['institution'],
             'citation': ds.attrs['citation'],
             'creation_date': ds.attrs['creation_date'],
             'contact': ds.attrs['contact'],
             'history': today}
                 
    # Coordinates for new Dataset
    coords = {'time': ds.coords['time'],
              'lat': lat_out,
              'lon': lon_out}
          
    # Define a new Dataset and ovewrite dims, attrs, coords
    ds_out = xr.Dataset(attrs=attrs, coords=coords)
          
    # Define the regridder object
    regridder = xe.Regridder(ds, ds_out, 'bilinear', reuse_weights=False)

    # Regrid the data
    dr = regridder(ds['MODIS'])
    print(dr)
    
    # Copy the same variable attributes to the regridded data
    dr.attrs = ds['MODIS'].attrs

    # Save to the new dataset
    ds_out['MODIS'] = dr

    # -----------------------
    # Save output to disk
    # -----------------------  

    # Define output file
    outfile = os.path.join(maindatadir,
                          'MODIS.LAI.BNUv6.generic.025x025.{}.nc'.format(yyyy))
    # Write to netCDF
    print('Writing {}'.format(outfile))
    ds_out.to_netcdf(outfile)