# import libraries
import numpy as np 
from netCDF4 import Dataset
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap
import warnings ; warnings.filterwarnings("ignore")
import os

# define wrapper function for interpolation 
def interp_senorge(fldIn, lonIn, latIn, 
                resOut=(1.25,180/191), lonBounds=(-1.25/2,360-1.25/2), latBounds=(-90-90/191,90+90/191), 
                method='conservative', maskOption='best'):
    """
Regrids se-norge output to a regular lon/lat grid (default=CESM's f09). 

Parameters:
-----------
Mandatory 
    
     fldIn : se-norge input field (2d or 3d) 
        
     lonIn : 2d longitude  

     latIn : 2d latitude          
        
              
Optional 

    resOut : resolution in deg; if skalar is specified the value is used for 
             both longitude and latitude; if tuple is specified the first value 
             is used for longitude and the second for latitude

 lonBounds : tuple that specifies the western and eastern domain bounds 

 latBounds : tuple that specifies the southern and northern domain bounds 

    method : interpolation method; must be one of 'conservative', 'bilinear',
             'patch', 'nearest_s2d' and 'nearest_d2s' 
                 
maskOption : must be one of 'conservative', 'best' and 'maximized';
             'conservative' sets points to nan if they contain any land,
             'best' if more than 50% land and 'maximized' if 100% land

Returns:
--------
    fldOut, lonOut, latOut : output field, longitude matrix, latitude matrix        

History: 
--------
2025.05.20 created (Ingo.Bethke@uib.no)
    """
    import xesmf as xe    
    global regridder, gridOut, maskOutInv
    #
    # make sure that input is a masked array; convert missing to mask
    fldIn = np.ma.masked_invalid(fldIn,copy=False)
    #
    # prepare weights if necessary
    if (not 'regridder' in globals()):   
        #
        # prepare se-norge coordinate information
        lonExt = np.concatenate((lonIn[:,0:1]-(lonIn[:,1:2]-lonIn[:,0:1])/2,lonIn,lonIn[:,-1:]+(lonIn[:,-1:]-lonIn[:,-2:-1])/2),axis=1)
        lonExt = np.concatenate((lonExt[0:1,:]-(lonExt[1:2,:]-lonExt[0:1,:])/2,lonExt,lonExt[-1:,:]+(lonExt[-1:,:]-lonExt[-2:-1,:])/2),axis=0)
        latExt = np.concatenate((latIn[:,0:1]-(latIn[:,1:2]-latIn[:,0:1])/2,latIn,latIn[:,-1:]+(latIn[:,-1:]-latIn[:,-2:-1])/2),axis=1)
        latExt = np.concatenate((latExt[0:1,:]-(latExt[1:2,:]-latExt[0:1,:])/2,latExt,latExt[-1:,:]+(latExt[-1:,:]-latExt[-2:-1,:])/2),axis=0)
        lon_b = 0.25 * (lonExt[0:-1,0:-1]+lonExt[0:-1,1:]+lonExt[1:,0:-1]+lonExt[1:,1:])
        lat_b = 0.25 * (latExt[0:-1,0:-1]+latExt[0:-1,1:]+latExt[1:,0:-1]+latExt[1:,1:])                                 
        gridIn = {'lon': lonIn, 'lon_b': lon_b, 'lat': latIn, 'lat_b': lat_b}
        # 
        # prepare regular grid coordinate information 
        if type(resOut) == int or type(resOut) == float:   
            resLon, resLat = resOut, resOut 
        else: 
            resLon, resLat = resOut[0], resOut[1]        
        periodic=True 
        if lonBounds[0] % 360 != lonBounds[1] % 360: 
            periodic=False 
        gridOut = xe.util.grid_2d(lonBounds[0],lonBounds[1],resLon,
                                  latBounds[0],latBounds[1],resLat)
        gridOut.lat_b[:] = np.maximum(-90,np.minimum(90,gridOut.lat_b[:]))        
        #
        # compute weights 
        weightsFile = '{:s}_{:d}x{:d}_{:d}x{:d}.nc'.format(method,fldIn.shape[1],fldIn.shape[0], \
                      gridOut.y.size,gridOut.x.size) 
        if os.path.isfile(weightsFile):
            regridder = xe.Regridder(gridIn, gridOut, method, periodic=periodic,
                                     reuse_weights=True,filename=weightsFile)
        else:
            regridder = xe.Regridder(gridIn, gridOut, method, periodic=periodic,
                                     reuse_weights=False)
        #
        # create ocean mask from data and interpolate to output grid
        maskIn = np.where(fldIn.mask, np.zeros(fldIn.shape), 
                          np.ones(fldIn.shape))
        maskOut = regridder(maskIn)        
        if maskOption == 'conservative':
            maskOut = np.where(maskOut>1.-1e-10, maskOut, np.nan)
        elif maskOption == 'maximized':
            maskOut = np.where(maskOut>0.+1e-10, maskOut, np.nan) 
        else:
            if maskOption != 'best':
                sys.exit('Unknown value for oceanMaskType. Will use "best".')
            maskOut = np.where(maskOut>0.5, maskOut, np.nan)
        maskOutInv = 1. / maskOut
    #
    # interpolate, normalize and mask
    # NOTE: To allow interpolation if one or more of contributing input cells 
    #       contain missing values, all missing values are set to zero in input. 
    #       This has the same effect as setting the grid cell area to zero, so 
    #       that cell effectively does not contribute. To ensure conservation, 
    #       the same is done in the interpolation of the ocean mask which is  
    #       used to normalize the output field.
    fldOut = regridder(np.where(fldIn.mask,0.,fldIn)) * maskOutInv
    return fldOut, np.array(gridOut.lon), np.array(gridOut.lat)

# read input data
IPATH='/nird/projects/NS9039K/users/filip/data/senorge_data/rr/rr_2025.nc'
nc = Dataset(IPATH)
lonIn = nc.variables['lon'][:]
latIn = nc.variables['lat'][:]
fldIn = nc.variables['rr'][:]
nc.close()   

# interpolate data 
fldOut, lonOut, latOut = interp_senorge(fldIn,lonIn,latIn)

# plot map from first time slice 
m = Basemap(projection='cyl',llcrnrlat=55,urcrnrlat=75,\
            llcrnrlon=0,urcrnrlon=35,resolution='l')
x, y = m(lonOut, latOut)
m.pcolormesh(x,y,fldOut[0,:,:])  
m.drawcoastlines()
plt.savefig('interp_senorge.png',format='png',dpi=200,bbox_inches='tight')