import warnings
warnings.filterwarnings("ignore", category=DeprecationWarning)

import xarray as xr
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import cartopy.feature as cfeature
import numpy as np
import fiona
import geopandas as gpd

fiona.drvsupport.supported_drivers['libkml'] = 'rw' # enable KML support which is disabled by default
fiona.drvsupport.supported_drivers['LIBKML'] = 'rw' # enable KML support which is disabled by default

from cartopy.mpl.gridliner import LONGITUDE_FORMATTER, LATITUDE_FORMATTER
from cartopy.io import shapereader

ds = xr.open_mfdataset('/gpfs/wolf2/arm/atm124/proj-shared/bnf/bnfcsapr2squireS3.c1/*.nc')

ds

def gc_latlon_bear_dist(lat1, lon1, bear, dist):
    """
    Input lat1/lon1 as decimal degrees, as well as bearing and distance from
    the coordinate. Returns lat2/lon2 of final destination. Cannot be
    vectorized due to atan2.
    """
    re = 6371.1  # km
    lat1r = np.deg2rad(lat1)
    lon1r = np.deg2rad(lon1)
    bearr = np.deg2rad(bear)
    lat2r = np.arcsin((np.sin(lat1r) * np.cos(dist/re)) +
                      (np.cos(lat1r) * np.sin(dist/re) * np.cos(bearr)))
    lon2r = lon1r + np.atan2(np.sin(bearr) * np.sin(dist/re) *
                          np.cos(lat1r), np.cos(dist/re) - np.sin(lat1r) *
                          np.sin(lat2r))
    return np.rad2deg(lat2r), np.rad2deg(lon2r)
 
    
def add_scale_line(scale, ax, projection, color='k',
                  linewidth=None, fontsize=None, fontweight=None):
    """
    Adds a line that shows the map scale in km. The line will automatically
    scale based on zoom level of the map. Works with cartopy.
 
    Parameters
    ----------
    scale : scalar
        Length of line to draw, in km.
    ax : matplotlib.pyplot.Axes instance
        Axes instance to draw line on. It is assumed that this was created
        with a map projection.
    projection : cartopy.crs projection
        Cartopy projection being used in the plot.
 
    Other Parameters
    ----------------
    color : str
        Color of line and text to draw. Default is black.
    """
    frac_lat = 0.1  # distance fraction from bottom of plot
    frac_lon = 0.5  # distance fraction from left of plot
    e1 = ax.get_extent()
    center_lon = e1[0] + frac_lon * (e1[1] - e1[0])
    center_lat = e1[2] + frac_lat * (e1[3] - e1[2])
    # Main line
    lat1, lon1 = gc_latlon_bear_dist(
        center_lat, center_lon, -90, scale / 2.0)  # left point
    lat2, lon2 = gc_latlon_bear_dist(
        center_lat, center_lon, 90, scale / 2.0)  # right point
    if lon1 <= e1[0] or lon2 >= e1[1]:
        warnings.warn('Scale line longer than extent of plot! ' +
                      'Try shortening for best effect.')
    ax.plot([lon1, lon2], [lat1, lat2], linestyle='-',
            color=color, transform=projection, 
            linewidth=linewidth)
    # Draw a vertical hash on the left edge
    lat1a, lon1a = gc_latlon_bear_dist(
        lat1, lon1, 180, frac_lon * scale / 20.0)  # bottom left hash
    lat1b, lon1b = gc_latlon_bear_dist(
        lat1, lon1, 0, frac_lon * scale / 20.0)  # top left hash
    ax.plot([lon1a, lon1b], [lat1a, lat1b], linestyle='-',
            color=color, transform=projection, linewidth=linewidth)
    # Draw a vertical hash on the right edge
    lat2a, lon2a = gc_latlon_bear_dist(
        lat2, lon2, 180, frac_lon * scale / 20.0)  # bottom right hash
    lat2b, lon2b = gc_latlon_bear_dist(
        lat2, lon2, 0, frac_lon * scale / 20.0)  # top right hash
    ax.plot([lon2a, lon2b], [lat2a, lat2b], linestyle='-',
            color=color, transform=projection, linewidth=linewidth)
    # Draw scale label
    ax.text(center_lon, center_lat - frac_lat * (e1[3] - e1[2]) / 3.0,
            str(int(scale)) + ' km', horizontalalignment='center',
            verticalalignment='center', color=color, fontweight=fontweight,
            fontsize=fontsize)

# note: the KMZ file provided contains multiple geometry columns. 
in_layers = []
for layer in fiona.listlayers("locations/BNF.kmz"):
    print(layer)
    s = gpd.read_file("locations/BNF.kmz", layer=layer, engine='fiona')
    in_layers.append(s)

# Specify your start and end times here

start_time = '2025-05-01'
end_time = '2025-06-01'

# CSAPR2 PPIs are in 10 min timesteps, so the total accumulation is the sum of the rain rates divided by 6
dt = ds['time'].diff(dim='time').astype('timedelta64[s]')
total_accumulation = ds['rain_rate_combined'].where(ds['lowest_height'] < 2500.)
total = total_accumulation.sel(time=slice(start_time, end_time)).sum(dim='time', skipna=True)/6.

# Make the figure
fig, ax = plt.subplots(1, 1, subplot_kw=dict(projection=ccrs.PlateCarree()))

# Plot the mesh of accumulations
lat_grid = ds['lat'].values
lon_grid = ds['lon'].values
c = ax.pcolormesh(lon_grid, lat_grid, total.where(total > 10), vmin=0,
    vmax=500, cmap='Spectral_r')
ax.set_xlim([lon_grid.min(), lon_grid.max()])
ax.set_ylim([lat_grid.min(), lat_grid.max()])
plt.colorbar(c, label='Total accumulation [mm]')

# Add Grid lines
gl = ax.gridlines(draw_labels=True, linewidth=0.5, color='gray', alpha=0.7, linestyle='--')
gl.top_labels = False
gl.right_labels = False
gl.xformatter = LONGITUDE_FORMATTER
gl.yformatter = LATITUDE_FORMATTER            
# add in kmz file layers

# BNF Forest Preserve Land
in_layers[11].plot(transform=ccrs.PlateCarree(),
                   facecolor="none", 
                   edgecolor="green",
                   linewidth=2.0,
                   ax=ax,
                   label="BNF Forest Preserve",
                   zorder=2)
# CSAPR2 location
in_layers[8].plot(transform=ccrs.PlateCarree(),     
                  ax=ax, 
                  label="S3: CSAPR-2",
                  zorder=2,
                  marker="d",
                  markersize=65)
# X-SAPR location
in_layers[9].plot(transform=ccrs.PlateCarree(),     
                  ax=ax, 
                  label="S4: Ka/X-SACR",
                  zorder=2,
                  marker="d",
                  markersize=65)
# M1 location
in_layers[1].loc[[4], 'geometry'].plot(transform=ccrs.PlateCarree(),     
                                       ax=ax, 
                                       label="M1: WB, MET Tower, LD",
                                       zorder=2,
                                       marker="p",
                                       color="black",
                                       markersize=65)

# S20 location - MET, AERI, DL
in_layers[5].plot(transform=ccrs.PlateCarree(),     
                          ax=ax, 
                          label="S20: MET, AERI, DL",
                          zorder=2,
                          marker="x",
                          color="green",
                          markersize=65)

# S30 location - MET, RWP, LDIS
in_layers[6].plot(transform=ccrs.PlateCarree(),     
                          ax=ax, 
                          label="S30: MET, RWP, LDIS",
                          zorder=2,
                          marker="x",
                          color="red",
                          markersize=65)

# S40 location - MET, AERI, DL
in_layers[7].plot(transform=ccrs.PlateCarree(),     
                          ax=ax, 
                          label="S40: MET, DL",
                          zorder=2,
                          marker="x",
                          color="orange",
                          markersize=65)

# Plot the scale bar
add_scale_line(10, ax, projection=ccrs.PlateCarree())

                                                        
ax.set_title(f'Total accumulation {start_time} to {end_time}')
ax.legend()
fig.tight_layout()
plt.savefig('squire_May2025.png')