from netCDF4 import Dataset
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import matplotlib.colors as Normalize
import cartopy.crs as ccrs
import cartopy.feature as cfeature
from cartopy.mpl.ticker import LongitudeFormatter, LatitudeFormatter
import matplotlib.ticker as mticker
import matplotlib
import xarray as xr
import netCDF4
import numpy as np
import pandas as pd
import glob
import dask
import os
from matplotlib.colors import LinearSegmentedColormap, TwoSlopeNorm

import wrf
from wrf import (getvar, vinterp, interplevel, to_np, latlon_coords, get_cartopy,
                 cartopy_xlim, cartopy_ylim)

monthlist=['04','05','06']


for month in monthlist:
    #filename = f'/pscratch/sd/y/yuwei/Climate_Impact/long-term/data/RAP/current/rap_130_2017{month}30_0000_000.grb2'
    filename = f'/pscratch/sd/y/yuwei/Climate_Impact/long-term/data/RAP/current/rap_130_2017{month}*.grb2'

    # Use glob to find all matching files for the month
    file_list = sorted(glob.glob(filename))

    array_list=[]
    for file in file_list:

        # Load GRIB file
        ds = xr.open_dataset(file, engine="cfgrib", backend_kwargs={'filter_by_keys':{'stepType': 'instant', 'typeOfLevel': 'surface'}})
        array_list.append(ds)
        print(file)

    combined_ds = xr.concat(array_list, dim='time')
    combined_ds.to_netcdf(f'/pscratch/sd/d/dbrooks/RAP_analysis/rap_data_current_sfc_month{month}.nc')
    
'''

# For the 'future' files
for month in monthlist:

    var_list = ['3D', # w
                'psl', # slp in Pa
                'RHS', # relative humidity
                'sfc', # ?
                'tas', # temperature (2m and 80m)
                'uas', # u wind (10m and 80m)
                'vas'] # v wind (10m and 80m)

    for var in var_list:
        filename = f'/pscratch/sd/y/yuwei/Climate_Impact/long-term/data/RAP/future/rap_{var}.2017{month}*.grb'

        # Use glob to find all matching files for the month
        file_list = sorted(glob.glob(filename))

        array_list=[]
        for file in file_list:

            # Load GRIB file
            if var == '3D':
                ds = xr.open_dataset(file, engine="cfgrib", backend_kwargs={"filter_by_keys": {"typeOfLevel": "isobaricInhPa"}})
                ds[var] = ds['unknown']
                ds = ds.drop_vars(['unknown'])
            else:
                ds = xr.open_dataset(file, engine="cfgrib")
                ds[var] = ds['unknown']
                ds = ds.drop_vars(['unknown'])
            array_list.append(ds)
            print(file)

        combined_ds = xr.concat(array_list, dim='time')
        combined_ds.to_netcdf(f'/pscratch/sd/d/dbrooks/RAP_analysis/future/rap_{var}_future_month{month}.nc')

'''