import os
import numpy as np
import xarray as xr
import matplotlib.pyplot as plt
from matplotlib import ticker
import cartopy.feature as cf
from datetime import datetime
#===============================================================================
# Python utilities

lm_ID:   str = '_CH_1100m'
wrk_dir: str = '/scratch/snx3000/jcanton/lmp/wd'
out_dir: str = 'lm_fine/1h_2D'
Lfname = lambda date: 'lffd{0:s}0000.nc'.format(date.strftime('%Y%m%d%H'))
Lfpath = lambda date: os.path.join(wrk_dir, date.strftime('%y%m0100')+lm_ID, out_dir, Lfname(date))

#-------------------------------------------------------------------------------
# Load point values
#
def loadCHpointsInTime(time_dates: list[datetime], \
        vname: str, grid_points: list[list[int]]) -> list[list[float]]:

    values = []
    for time_date in time_dates:
        # assemble fpath
        run_dir = time_date.strftime('%y%m0100') + lm_ID
        run_path = os.path.join(wrk_dir, run_dir)
        fname = Lfname(time_date)
        fpath = os.path.join(run_path, out_dir, fname)
        values.append( loadCHpoints(fpath, vname, grid_points) )

    return values

def loadCHpoints(fpath: str, vname: str, grid_points: list[list[int]]) -> list[float]:

    field = xr.load_dataset(fpath)[vname].isel(time=0)

    values: list[float] = []
    for gpoint in grid_points:
        values.append( float(field[gpoint[1], gpoint[0]].values) )

    return values

#-------------------------------------------------------------------------------
# Load entire fields
#
def loadCHfield(lm_ID, wrk_dir, out_dir, Lfname, years, months, days, hours, variables):
    nyears   = len(years)
    nmonths  = len(months)
    ndays    = len(days)
    nhours   = len(hours)
    ntstamps = nyears*nmonths*ndays*nhours
    itime = -1
    for year in years:
        for month in months:

            # assemble run dir and path
            run_dir = datetime(year, month, 1).strftime('%y%m%d%H') + lm_ID
            run_path = os.path.join(wrk_dir, run_dir)

            # print info
            print('\tWorking on: ' + run_dir)

            for day in days:
                for hour in hours:

                    # update counter
                    itime += 1

                    # assemble date
                    date = datetime(year, month, day, hour)

                    # get file name
                    fname = Lfname(date)
                    fpath = os.path.join(run_path, out_dir, fname)

                    # print info
                    print('\t\tLoading file: ' + fname)

                    # load file
                    ds = xr.load_dataset(fpath) #[vname].isel(time=0)

                    for v in variables:
                        if v['data'] is None:
                            v['lat'] = ds.lat.values
                            v['lon'] = ds.lon.values
                            nlat = ds.lat.shape[0]
                            nlon = ds.lon.shape[1]
                            v['data'] = np.zeros( (ntstamps, nlat, nlon) )
                            v['tstamp'] = ntstamps*[None]

                        v['data'][itime, :,:] = ds[v['name']].isel(time=0).values
                        v['tstamp'][itime] = date


    print('Done loading\n')
    return variables

#-------------------------------------------------------------------------------
def minDst(lon, lat, loc):
    #
    # Find the lon/lat point with minimum distance from a location
    dst2 = ((lon - loc[0])**2 + (lat - loc[1])**2)
    idx_min = np.unravel_index(np.argmin(dst2, axis=None), dst2.shape)
    return (idx_min[1], idx_min[0]), dst2[idx_min]**0.5

#-------------------------------------------------------------------------------
def addCbar(fig, ax, im, label, lonMM=[5.96, 10.50], latMM=[45.82, 47.81], extend='neither', borders=True, ticks=None, dy=0):
    ax.axis('equal')
    ax.set_xlim([lonMM[0], lonMM[1]])
    ax.set_ylim([latMM[0], latMM[1]])
    ax.axis('off')
    if borders:
        ax.coastlines()
        ax.add_feature(cf.BORDERS)
    axS = ax.get_position().x0
    axB = ax.get_position().y0
    axW = ax.get_position().x1 - ax.get_position().x0
    axS += axW/10
    axW  = axW/10*8
    axH = 0.01
    cbaxes1 = fig.add_axes([axS, axB+dy*axH, axW, axH])
    cb = plt.colorbar(im, cax=cbaxes1, orientation="horizontal", extend=extend) #, format=ticker.FuncFormatter(tickFmt))
    cb.set_label(label, fontsize='small')
    cb.ax.tick_params(labelsize='x-small')
    if ticks is None:
        tick_locator = ticker.MaxNLocator(nbins=5)
        cb.locator = tick_locator
    else:
        cb.set_ticks(ticks)
    cb.update_ticks()
    plt.draw()

#-------------------------------------------------------------------------------
def borderMask(mask: np.ndarray, nr_layers: int = 1, diags: bool = False) -> np.ndarray:
    # add layers to a logical mask
    import copy


    if nr_layers > 1:
        mask = borderMask(mask, nr_layers-1)

    I, J = mask.shape
    mask2 = copy.copy(mask)
    for i in range(I):
        for j in range(J):
            if mask[i,j]:
                mask2[i,j] = mask2[i+1,j] = mask2[i-1,j] = mask2[i,j+1] = mask2[i,j-1] = True
                if diags:
                    mask2[i+1,j+1] = mask2[i+1,j-1] = mask2[i-1,j+1] = mask2[i-1,j-1] = True

    return mask2