import numpy as np
import netCDF4 as nc
from netCDF4 import Dataset
from glob import glob
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap

data_path = '/Users/vicki/Dropbox/xch4_project/data/tropomi/'
figure_path = '/Users/vicki/Dropbox/xch4_project/figure/tropomi/texas'

filename_list = sorted(glob(data_path+'TROPOMI*nc'))
# Site coordinates
# site_lat, site_lon = 36.607322, -97.487643  #ARM
site_lat, site_lon = 31.1316741975655, -99.33607944174358 # Heart of Texas

farm_lat,farm_lon = 36.641168, -97.538272
okc_lat,okc_lon = 35.481918, -97.508469
austin_lat, austin_lon = 30.266666, -97.733330
dallas_lat, dallas_lon = 32.779167, -96.808891
houston_lat, houston_lon = 29.749907, -95.358421

# Define boundaries for a smaller region (e.g., +/- 2.5 degrees)
# lat_range = 38
# lon_range = 30
lat_range = 50
lon_range = 50

for filename in filename_list:
    year_month = filename[-9:-5]+'-'+filename[-5:-3]
    # Open the NetCDF file
    ds = nc.Dataset(filename, 'r')

    # Read data
    lon = ds.variables['LON'][:]
    lat = ds.variables['LAT'][:]
    ch4_monthly_mean = ds.variables['CH4_monthly_mean'][:]

    # Close the dataset
    ds.close()

    # Mask for 2D latitude and longitude arrays
    site_index = np.where((lon-site_lon)**2+(lat-site_lat)**2==np.min((lon-site_lon)**2+(lat-site_lat)**2))

    # Apply the mask to your 2D CH4 data array as well as longitude and latitude
    lon_subset = lon[site_index[0][0]-lat_range:site_index[0][0]+int(lat_range*1.2),site_index[1][0]-lon_range:site_index[1][0]+lon_range]
    lat_subset = lat[site_index[0][0]-lat_range:site_index[0][0]+int(lat_range*1.2),site_index[1][0]-lon_range:site_index[1][0]+lon_range]
    ch4_mean_subset = ch4_monthly_mean[site_index[0][0]-lat_range:site_index[0][0]+int(lat_range*1.2),site_index[1][0]-lon_range:site_index[1][0]+lon_range]

    fig, ax = plt.subplots(figsize=(10, 7))

    # Create a Basemap centered around the site with specified bounds
    m = Basemap(projection='merc', llcrnrlat=np.min(lat_subset), urcrnrlat=np.max(lat_subset),
            llcrnrlon=np.min(lon_subset), urcrnrlon=np.max(lon_subset), resolution='h')

    # Transform the longitude and latitude into the map projection coordinates
    lon_map, lat_map = m(lon_subset, lat_subset)


    # Ensure ch4_mean_subset is reshaped to match the transformed coordinates if necessary
    ch4_plot = m.pcolormesh(lon_map, lat_map, ch4_mean_subset*1e9, cmap='jet',vmin=1800,vmax=1900)


    # Add features and colorbar
    m.drawcoastlines()
    m.drawcountries()
    m.drawcounties()
    m.drawstates()

    # Add grid lines for latitudes and longitudes with labels
    m.drawmeridians(np.arange(round(np.min(lon_subset)), round(np.max(lon_subset)), 3), labels=[0,0,0,1], fontsize=10, color='gray')
    m.drawparallels(np.arange(round(np.min(lat_subset)), round(np.max(lat_subset)), 3), labels=[1,0,0,0], fontsize=10, color='gray')

    plt.colorbar(ch4_plot, label='CH4 Monthly Mean (ppb)',fraction=0.025)

#     Add a red star at the site location
#     x, y = m(site_lon, site_lat)
#     star1 = m.scatter(x, y, marker='*', c='r', s=30,label='ARM')

#     x, y = m(farm_lon, farm_lat)
#     star2 = m.scatter(x, y, marker='d', c='purple', s=30,label='AFO',alpha=0.8)

    x, y = m(okc_lon, okc_lat)
    star_okc = m.scatter(x, y, marker='o', c='k', s=50,label='OKC')
    
    x, y = m(austin_lon, austin_lat)
    star1 = m.scatter(x, y, marker='*', c='tab:pink', s=60,label='Austin')

    x, y = m(dallas_lon, dallas_lat)
    star2 = m.scatter(x, y, marker='d', c='purple', s=50,label='Dallas',alpha=0.8)

    x, y = m(houston_lon, houston_lat)
    star3 = m.scatter(x, y, marker='^', c='k', s=50,label='Houston')

    plt.legend(handles=[star1,star2,star3,star_okc],loc='lower right')

    plt.title('CH4 Monthly Mean '+year_month)
    # plt.show()
    fig.savefig(figure_path+'/TROPOMI_CH4_monthly_mean_'+year_month+'.png')