import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import pickle
import datetime as dt
import sys,os,glob
import cartopy.crs as ccrs
import cartopy.feature as cfeature
import logging
logging.basicConfig(level=logging.WARNING)
#logging.basicConfig(level=logging.INFO)
from netCDF4 import Dataset
# pull from https://github.com/Kang-Sun-CfA/Oversampling_matlab
sys.path.append('/user/zolalaya/Oversampling_matlab/')
from CAREER.tempo import TEMPO
from popy import Level3_Data, Level3_List
%matplotlib inline

Process Level 4

start_dt = dt.datetime(2024,5,1,0) #str(sys.argv[1])
end_dt = dt.datetime(2024,8,31,23) #str(sys.argv[2])
altitude = 100 #int(str(sys.argv[3]))

grid_size = 0.05
flux_grid_size = 0.05
interp_met_kw = {'which_met':'ERA5',
                 'met_dir':'/projects/academic/kangsun/data/ERA5/Y%Y/M%m/D%d/CONUS_3D_%Y%m%d.nc',
                 'interp_fields':['u10','v10'],
                 'altitudes':[altitude]}
calculate_gradient_kw = {'write_diagnostic':True}
gradient_kw = {'x_wind_field':f'era5_u{altitude}',
               'y_wind_field':f'era5_v{altitude}',
               'x_wind_field_sfc':'era5_u10',
               'y_wind_field_sfc':'era5_v10',
               'interp_met_kw':interp_met_kw,
               'calculate_gradient_kw':calculate_gradient_kw
              }

# Cornbelt: subwest=-105;subeast=-80;subsouth=36;subnorth=50
wesn_dict = dict(west=-105,east=-80,south=36,north=50)

# ### Only when producing L4
# # l2_path_pattern='/projects/academic/kangsun/data/TEMPONO2/TEMPO_NO2_L2_V03/%Y/%m/%d/'
# l2_path_pattern='/projects/academic/kangsun/data/TEMPONO2/TEMPO_NO2_L2_V03/%Y/%m/%d/TEMPO_NO2_L2_V03_%Y%m%d*.nc4'
# l4_path_pattern='/projects/academic/kangsun/data/TEMPONO2/cornbelt_l4_0.05/%Y/%m/%d/{}_TEMPO_NO2_L4_V03_%Y%m%d.nc'.format(
#     gradient_kw['x_wind_field'])
# # mons = pd.period_range('2024-2','2024-8',freq='1m')
# mons = pd.period_range(start_dt,end_dt,freq='1D')
# for mon in mons:
#     logging.warning('running {}'.format(mon.strftime('%Y%m%d')))
#     t = TEMPO(product='NO2',**wesn_dict,
#               start_dt=mon.start_time,
#               end_dt=mon.end_time,
#               grid_size=grid_size,flux_grid_size=flux_grid_size)

#     t.regrid_from_l2(l2_path_pattern=l2_path_pattern,
#                      l4_path_pattern=l4_path_pattern,
#                      gradient_kw=gradient_kw,
#                      maxsza=70,maxcf=0.2,
#                      ncores=20,block_length=200,
#                      use_TEMPOL2=False,
#                      l4_save_fields = ['column_amount','local_hour','terrain_height',\
#                                        'wind_topo','wind_column','wind_column_xy','wind_column_rs',\
#                                        'wind_topo_xy','wind_topo_rs'])

# Load L4 scans and also calculate scans in local hours between 9-11, 11-3, and 13.15
l4_path_pattern=\
'/projects/academic/kangsun/data/TEMPONO2/cornbelt_l4_0.05/%Y/%m/%d/era5_u100_TEMPO_NO2_L4_V03_%Y%m%d*.nc'
tendency=['c','b','f']
fields_name = ['column_amount','local_hour','terrain_height',\
               'wind_topo', 'wind_column','wind_column_xy','wind_column_rs',\
               'wind_topo_xy','wind_topo_rs']

tend = TEMPO(product='NO2',**wesn_dict,
             start_dt=start_dt,end_dt=end_dt)
tend.load_scans(l4_path_pattern,fields_name=fields_name,tendency=tendency,local_hour_centers=[10,12,14],local_hour_spans=[2,2,2])

/user/zolalaya/Oversampling_matlab/popy.py:3730: RuntimeWarning: Mean of empty slice
  storage = np.nanmean([

# Tune add_storage (wind_column_storage)
tendency=['c','b','f']
tend.l3s.add_storage(vcd_field='column_amount',
                     emission_field='wind_column_xy',tendency=tendency,
                     mode='priority',nan_tendency_as_zero=True)

tend.l3s[0].keys()

dict_keys(['column_amount', 'local_hour', 'terrain_height', 'wind_topo', 'wind_column', 'wind_column_xy', 'wind_column_rs', 'wind_topo_xy', 'wind_topo_rs', 'xgrid', 'ygrid', 'num_samples', 'total_sample_weight', 'storage_c_column_amount', 'storage_b_column_amount', 'storage_f_column_amount', 'wind_column_storage', 'wind_column_xy_storage'])

# Integrate (average) l3s for 9-11, 11-13, and 13-15 for the whole time between start_dt and end_dt
hours = [10, 12, 14]
l3_merge = {}
for hour in hours:
    mask = tend.l3_lhs.df.index.hour == hour
    l3_hour = None
    idxs = np.where(mask)[0]   
    for i in idxs:
        l3_idx = tend.l3_lhs[i]
        if l3_hour is None:
            l3_hour = l3_idx
        else:
            l3_hour = l3_hour.merge(l3_idx)   
    l3_merge[hour] = l3_hour

# Plot different fields, averaged over the whole time between start_dt and end_dt
fields_name = ['column_amount', 'wind_column', 'wind_column_xy', 'wind_column_storage']
hours = [10, 12, 14]

fig,axs = plt.subplots(len(fields_name),len(hours),figsize=(15,15),
                       constrained_layout=True,subplot_kw={"projection": ccrs.PlateCarree()})

for i, field in enumerate(fields_name):
    for j, hour in enumerate(hours):
        ax = axs[i, j]
        if field == 'column_amount':
            im = l3_merge[hour].plot(
                field,
                ax=ax,
                draw_colorbar=False
            )
        else:
            im = l3_merge[hour].plot(
                field,
                func=lambda x:1e9*x,
                ax=ax,
                draw_colorbar=False,
                cmap='viridis',
                vmin=-3,vmax=2
            )
        if i == 0:
            ax.set_title(f'{hour-1}:00 to {hour+1}:00', fontsize=10)
        if j == 2:
            cbar = plt.colorbar(
                im['pc'],
                ax=ax,
                orientation='vertical',
                shrink=0.7,  # height of colorbar
                pad=0.04         # space between plot and colorbar
            )
y_positions = [0.84, 0.59, 0.35, 0.08]
for y, field in zip(y_positions, fields_name):
    fig.text(-0.02, y, field, rotation='vertical', fontsize=12)# plt.tight_layout()
plt.show()

for hour in hours:
    figout = l3_merge[hour].plot(
        'column_amount',
        vmin=0,
        vmax=1.5e-4,
    )
    ax = plt.gca()
    ax.set_title(f'hours {hour-1}-{hour+1}')
#     figout['fig'].savefig(f'test_{hour}.png')

for hour in hours:
    figout = l3_merge[hour].plot(
        'wind_column',
        func=lambda x:1e9*x,
        vmin=-5,
        vmax=5,
        cmap='viridis'
    )
    ax = plt.gca()
    ax.set_title(f'hours {hour-1}-{hour+1}')
#     figout['fig'].savefig(f'test_{hour}.png')

for hour in hours:
    figout = l3_merge[hour].plot(
        'storage_c_column_amount',
        func=lambda x:1e9*x,
        vmin=-4,
        vmax=2,
        cmap='viridis'
    )
    ax = plt.gca()
    ax.set_title(f'hours {hour-1}-{hour+1}')
#     figout['fig'].savefig(f'test_{hour}.png')

Combinations with land cover CDL

# Load land cover type pickle
with open('/user/zolalaya/TEMPO_AG/tempo_cdl_land_cover_fractions_2024_10m_0.05.pkl', 'rb') as f:
    lc = pickle.load(f)

# Plot land cover
indices = [1, 82] # 1:Corn; 82:Developed
fig, axs = plt.subplots(len(indices), 1, figsize=(10, 10),
                        subplot_kw={'projection': ccrs.PlateCarree()})

for ax, idx in zip(axs, indices):
    im = ax.pcolormesh(lc['lons'], lc['lats'], lc['fractions'][:, :, idx], 
                       cmap='viridis', transform=ccrs.PlateCarree())
    ax.coastlines(linewidth=0.5,color='gray')
    ax.add_feature(cfeature.BORDERS, linewidth=0.5)
    ax.add_feature(cfeature.STATES, linewidth=0.3)
    plt.colorbar(im, ax=ax, label='Fraction', shrink=0.7)
    ax.set_xlim(-105, -80)
    ax.set_ylim(36, 50)
    ax.set_xticks(np.linspace(-105, -80, 6))
    ax.set_yticks(np.linspace(36, 50, 6))
    ax.set_xlabel('Longitude')
    ax.set_ylabel('Latitude')
    ttl = lc['land_cover'][idx]
    ax.set_title(f'Fraction of Land Cover: {ttl}')
plt.tight_layout()
plt.show()

/cvmfs/soft.ccr.buffalo.edu/versions/2023.01/easybuild/software/avx512/MPI/gcc/11.2.0/openmpi/4.1.1/cartopy/0.20.3/lib/python3.9/site-packages/cartopy/mpl/geoaxes.py:1797: MatplotlibDeprecationWarning: shading='flat' when X and Y have the same dimensions as C is deprecated since 3.3.  Either specify the corners of the quadrilaterals with X and Y, or pass shading='auto', 'nearest' or 'gouraud', or set rcParams['pcolor.shading'].  This will become an error two minor releases later.
  result = matplotlib.axes.Axes.pcolormesh(self, *args, **kwargs)
/cvmfs/soft.ccr.buffalo.edu/versions/2023.01/easybuild/software/avx512/MPI/gcc/11.2.0/openmpi/4.1.1/cartopy/0.20.3/lib/python3.9/site-packages/cartopy/mpl/geoaxes.py:1797: MatplotlibDeprecationWarning: shading='flat' when X and Y have the same dimensions as C is deprecated since 3.3.  Either specify the corners of the quadrilaterals with X and Y, or pass shading='auto', 'nearest' or 'gouraud', or set rcParams['pcolor.shading'].  This will become an error two minor releases later.
  result = matplotlib.axes.Axes.pcolormesh(self, *args, **kwargs)

# developed = lc['fractions'][:, :, 82] + lc['fractions'][:, :, 121] + lc['fractions'][:, :, 122] + lc['fractions'][:, :, 123] + lc['fractions'][:, :, 124] > 0.2
# plt.figure(figsize=(10, 6))
# ax = plt.axes(projection=ccrs.PlateCarree())
# im = ax.pcolormesh(lc['lons'], lc['lats'], developed, 
#                    cmap='viridis', transform=ccrs.PlateCarree())
# ax.coastlines(linewidth=0.5,color='gray')
# ax.add_feature(cfeature.BORDERS, linewidth=0.5)
# ax.add_feature(cfeature.STATES, linewidth=0.3)
# plt.colorbar(im, ax=ax, label='Fraction', shrink=0.7)
# ax.set_xlim(-105, -80)
# ax.set_ylim(36, 50)
# x_ticks = np.linspace(-105, -80, 6)
# y_ticks = np.linspace(36, 50, 6)
# ax.set_xticks(x_ticks)
# ax.set_yticks(y_ticks)
# plt.title(f'Fraction of Land Cover: {class_name}')
# plt.xlabel('Longitude')
# plt.ylabel('Latitude')

# Land cover masks threshhold
corn_mask = lc['fractions'][:,:,1] > 0.5
devel_mask = lc['fractions'][:, :, 82] + lc['fractions'][:, :, 121] +\
             lc['fractions'][:, :, 122] + lc['fractions'][:, :, 123] + lc['fractions'][:, :, 124] > 0.5
print(np.count_nonzero(corn_mask))
print(np.count_nonzero(devel_mask))
# np.count_nonzero((np.nansum(lc['fractions'], axis=2) >= 0.999999) &
#                  (np.nansum(lc['fractions'], axis=2) <= 1.000001))

3556
158

# Integrate emission flux (mol/m2/s) to emission rates (mol/s) 
# over the pixels with fraction of land type above a threshold (mask)

# Define masks
corn_mask = lc['fractions'][:,:,1] > 0.5
devel_mask = lc['fractions'][:, :, 82] + lc['fractions'][:, :, 121] +\
             lc['fractions'][:, :, 122] + lc['fractions'][:, :, 123] + lc['fractions'][:, :, 124] > 0.5
masks = [corn_mask, devel_mask]
mask_names = ['Corn', 'Developed']

# Integrate emission flux over thosed masked pixels for different emission estimators
fields_name = ['wind_column_xy', 'wind_topo_xy', 'wind_column_storage']
emission_df = []
for mask in masks:
    emission_masked = np.full([len(tend.l3_lhs), len(fields_name)], np.nan, dtype='float32')
    for i in range(len(tend.l3_lhs)):
        emission_masked_i = tend.l3_lhs[i].sum_by_mask(mask=mask, fields_to_sum=fields_name)
        for j, field in enumerate(fields_name):
            emission_masked[i, j] = emission_masked_i[field]
    
    df = pd.DataFrame(emission_masked, columns=fields_name, index=tend.l3_lhs.df.index)
    emission_df.append(df)

# Plot corn and deleveloped
fig, axes = plt.subplots(2, 3, figsize=(18, 8), sharex=True)
hours = [10, 12, 14]

for row, (df, mask_label) in enumerate(zip(emission_df, mask_names)):
    for col, field in enumerate(fields_name):
        ax = axes[row, col]
        for hour in hours:
            hourly_series = df[df.index.hour == hour][field]
            ax.plot(hourly_series.index, hourly_series.values, label=f'{hour}:00')
        
        if row == 0:
            ax.set_title(field)
        if col == 0:
            ax.set_ylabel(f'{mask_label}\nEmission Rate (mol/s)')
        if row == 0 and col == 2:
            ax.legend(loc='upper right')
        ax.grid(True)

for ax in axes[1, :]:
    ax.set_xlabel('Date')
    ax.tick_params(axis='x', rotation=45)

fig.text(0.53, 1.05, 'NO${_2}$ Emission over the US Corn Belt \n (for pixels with corn/developed fraction>0.5',
         ha='center', va='top', fontsize=16)
plt.tight_layout()
plt.show()

# Define masks
corn_mask = lc['fractions'][:,:,1] > 0.5

# Plot only corn
emission_masked = np.full([len(tend.l3_lhs), len(fields_name)], np.nan, dtype='float32')
for i in range(len(tend.l3_lhs)):
    emission_masked_i = tend.l3_lhs[i].sum_by_mask(mask=corn_mask, fields_to_sum=fields_name)
    for j in range(len(fields_name)):
        emission_masked[i,j] = emission_masked_i[fields_name[j]]
emission_df = pd.DataFrame(emission_masked, columns=fields_name, index=tend.l3_lhs.df.index)

fig, axes = plt.subplots(3, 1, figsize=(11, 9), sharex=True)
hours = [10, 12, 14]
for ax, field in zip(axes, fields_name):
    for hour in hours:
        hourly_series = emission_df[emission_df.index.hour == hour][field]
        ax.plot(hourly_series.index, hourly_series.values, label=f'{hour}:00')
    ax.set_xlim(dt.datetime(2024,6,1,0), dt.datetime(2024,7,1,0))
    ax.tick_params(axis='x', rotation=45)
    ax.set_title(field)
    ax.set_ylabel('Emission Rate (mol/s)')
    ax.grid(True)
axes[2].set_xlabel('Date')
axes[0].legend()
fig.text(0.53, 1.04, 'Corn NO${_2}$ Emission over the US Corn Belt \n (for pixels with corn fraction>0.5)',
         ha='center', va='top', fontsize=16)
plt.tight_layout()
plt.show()

# corn_mask = corn_fraction > 0.3
# nebraska_mask = get_nebraska_mask()

# combined_mask = corn_mask & nebraska_mask

# Level3_List.sum_by_mask(mask=combined_mask, fields_to_sum=['wind_column_topo_chem'])