BRCA_Deep_Dive.Rmd

---
title: "BRCA_Deep_Dive_Processing"
author: "KEC"
date: "2024-04-04"
output: html_document
editor_options: 
  chunk_output_type: console
---

## CCVA Data Processing

This R markdown prepares data inputs for a Climate Change Vulnerability 
Assessment (CCVA) for water supplies at BRCA

### 1. Setup Workspace

First, install and/or load required packages and functions.

```{r setup, include=FALSE}

knitr::opts_chunk$set(echo = TRUE)

source('setup.R')
library("readxl")
library("trend")
library("elevatr")
library("viridis")
library("SPEI")

 # Optional preview water source locations
mapviewOptions(fgb = FALSE, 
               georaster = FALSE, 
               basemaps = c("Esri.WorldTopoMap",
                            "Esri.WorldImagery"))

```


Define park using 4-digit NPS Unit Code. Codes for all parks can be found at: 
<https://www.nps.gov/aboutus/foia/upload/NPS-Unit-List.xlsx>.

Download the park boundary from the NPS IRMA DataStore as an sf object.

```{r park_dat, eval = TRUE, echo = FALSE, message = FALSE, warning = FALSE}

park <- "BRCA"

# Get park geometry and metadata
park_boundary <- getParkBoundary(park = park)
state <- park_boundary$STATE
park_name <- park_boundary$UNIT_NAME
park_name_short <- gsub(" National Park", "", park_name)

# Define a color palette to use throughout - note can change for regions? feels?
pal = c("#DFCB34","#0067A2", "#CB7223", "#289A84", "#7FA4C2", "#AF7E56",
        "#8C2B0E", "#FEB359", "#132F5B", "#435F90", "#68434E", "#B47E83",
        "#444E7E", "red","hotpink", "#B7ABBC","#FD8700", "#D8511D")

# Define report name
report_name <- 
  paste0("Climate Change Vulnerability Assessment for Water Supplies at ",
         park_name)

```

Data import option

Read in data if you've already downloaded it
```{r read_data}

# Commented out for now because this is old data.

## path to data folder (from project directory)
#path <- "data/park/"

# read in data
#try(
#load(paste0(path, "/", park, "/", park, "_report_data_v2.RData")))
## terra object issue when saving as .RData, so saved as separate .tif
```


### 2. Water Supply Database

A water supply system database was created for this project using various NPS
and public databases.  The database contains a "supply" table and a "source"
table.  The "supply" table includes a row for each water supply system at
the park. The source table contains a row for each source of water (e.g.,
individual wells, diversions, and springs) associated with each supply. The
two tables are related by a shared column -> wsd_system_id.

```{r, eval = TRUE, echo = FALSE, message = FALSE, warning = FALSE}

table_path <- "data/Water_Supply_Systems/NPS_Water_Systems_Database.xlsx"

# Supply table (Currently filtered to Utah)
supply_table <- read_excel(table_path, sheet = 2, na = "NA") %>%
  janitor::clean_names() %>%
  dplyr::filter(park_name == park_boundary$UNIT_NAME,
                in_use == "Active")
  
# Source table
source_table <- read_excel(table_path, sheet = 3, na = "NA") %>%
  janitor::clean_names() %>%
  #dplyr::filter(state == "UT")
  dplyr::filter(wsd_system_id %in% supply_table$wsd_system_id) %>%
  dplyr::mutate(well_depth = as.numeric(well_depth))

# Copy of source table as sf_object - some columns without location data are
# dropped. Circle back to this later.
source_table_locs <- source_table %>%
  drop_na(c("source_longitude", "source_longitude")) %>%
  st_as_sf(., 
           coords = (c("source_longitude","source_latitude")), 
           crs = 4326,
           remove = FALSE) 


mapview(source_table_locs,
        zcol = "water_system_name",
        layer.name = "Source:",
        col.regions = pal) + 
  mapview(park_boundary,
          col.regions = "forestgreen",
          alpha.regions = 0.2,
          legend = FALSE) #%>%

# mapshot(., file = "/Users/kcognac/Desktop/BRCA_Deep_Dive/Site_map2.jpg", dpi = 600)

```


### 3. Delineate Source Watershed and/or Aquifer

```{r source_water}

# Get the aquifer and watershed extent
source_table_locs <- get_aoi_source(source_table_locs)

# Extract the watershed as a separate sf_object
watersupply_watershed <- source_table_locs$aoi %>%
  bind_rows() %>%
  dplyr::filter(type == "Watershed") %>%
  distinct() %>%
  st_transform(., crs = st_crs(park_boundary))

# Load NHD flowlines and filter for specific reaches within the watersupply
# watershed that represent the "main stem" of the alluvial aquifer
main_stem_flowlines <- 
  dplyr::summarize(watersupply_watershed) %>%
  mapNHDPlusHR() %>%
  dplyr::filter(GNIS_Name == "East Creek",
                ReachCode != 16030002000853) %>%# | 	ReachCode == 16030002000859) %>%
  dplyr::summarize()

# Similarly, load NHD flowlines but filter for a slightly broader extent of the
# alluvial aquifer.
all_alluvium_flowlines <-   
  dplyr::summarize(watersupply_watershed) %>%
  mapNHDPlusHR() %>%
  dplyr::filter(GNIS_Name == "East Creek" | 	ReachCode == 16030002000859) %>%
  dplyr::summarize()

# Now, create a buffer using each of those flowline selections
main_stem_buffer <- 
  main_stem_flowlines %>% 
  st_transform(., crs = st_crs(watersupply_watershed)) %>%
  dplyr::summarize() %>% 
  st_buffer(., .005)

all_alluvium_buffer <- 
  all_alluvium_flowlines %>% 
  st_transform(., crs = st_crs(watersupply_watershed)) %>%
  dplyr::summarize() %>% 
  st_buffer(., .005)

# View
#mapview(main_stem_buffer) + mapview(all_alluvium_buffer, col.regions = "pink")

# Now, crop the alluvium formation using each of the buffers to get a select
# extent of the alluvial mainstem
main_stem <- source_table_locs$aoi %>%
  bind_rows() %>%
  dplyr::filter(type == "Formation") %>%
  distinct() %>%
  st_transform(., crs = st_crs(watersupply_watershed))  %>%
    st_intersection(., main_stem_buffer) %>%
  st_intersection(., watersupply_watershed)

all_alluvium <- source_table_locs$aoi %>%
  bind_rows() %>%
  dplyr::filter(type == "Formation") %>%
  distinct() %>%
  st_transform(., crs = st_crs(watersupply_watershed))  %>%
    st_intersection(., all_alluvium_buffer) %>%
  st_intersection(., watersupply_watershed)
  

# View all
mapview(source_table_locs,
        zcol = "water_system_name",
        layer.name = FALSE,
        col.regions = c("tomato","dodgerblue"),
        cex = 4,
        homebutton = FALSE) + 
  mapview(park_boundary,
          col.regions = "forestgreen",
          alpha.regions = 0.2,
          homebutton = FALSE,
          legend = FALSE) +
  mapview(source_table_locs$aoi %>% bind_rows() %>%
            dplyr::mutate(name = ifelse(type == "Formation", "East Creek Alluvium", 
                                        "East Creek Watershed")),
          zcol = "name",
          layer.name = FALSE,
          col.regions = pal,
          homebutton = FALSE,
          alpha.regions = 0.2)  +
  
  mapview(all_alluvium,
          col.regions = "orange",
          alpha.regions = 0.4,
          layer.name = "Alluvium broad AOI") +
  mapview(main_stem,
          col.regions = pal[3],
          layer.name = "Alluvium Main Stem AOI") 

```


Various climate, hydrology, and geography data are required to generate a CCVA 
for the selected National Park. The following code chunks download and describe 
this data.


### 4. Park PODs

Park water supplies may be sourced from within or beyond the park boundary. This
chunk pulls in state-reported water supply locations, or points of diversion (PODS)
that occur within a buffer distance of the park boundary. Then, PODs are filtered
to identify specific PODs associated with park water supplies. For some parks, 
the water supply system ID is linked using the water supply database. For 
others, we assume the water supply is any POD owned by NPS within the select buffer.

\*\*\*\*KEC: This assumes POD_state has "OWNER" column and that NPS is 
identified by string "NATIONAL PARK". Should probably update match strings in 
future as other POD databases are brought in.

```{r, eval = TRUE, echo = FALSE, message = FALSE, warning = FALSE}

# Now, get state specific points of diversion (POD) (i.e., water supply points)
# EDIT NOTE - Eventually nest into a getPODall function with aoi and dist
# args which then applies the state-specific function.

# Set distance and AOI for POD search
buffer_dist <- .05 # in decimal degrees lat/long

aoi <- 
  park_boundary 

# Dictionary to map POD functions to states
POD_dict <- c(
  "CA" = getPODCalifornia,
  "CO" = getPODColorado,
  "MT" = getWaterRightsMontana,
  "NV" = getPODNevada,
  "UT" = getPODUtah
)

# Assign POD_state using the condition map
if (state %in% names(POD_dict)) {
  
  POD_state <- POD_dict[[state]](aoi, buffer_dist)  # Pass additional arguments

  } else {
    
  print("No PODs returned.")

  }


# Filter using water supply database to identify PODs associated park water 
# supplies. If table does not have adequate information, select by POD
# owner and location metadata

if(supply_table %>% drop_na(water_rights_id) %>% nrow() > 0) {
  
  POD_supply <- POD_state %>%
    #dplyr::filter(WRNUM %in% c("61-893", "2061001M00")) %>%
    dplyr::filter(WRNUM %in% supply_table$water_rights_id) %>%
    dplyr::distinct(LOCATION, .keep_all = TRUE) %>%
    dplyr::mutate(
      name = supply_table %>% 
        pull (water_system_name)) 
  
} else {
    
  POD_supply <- POD_state %>%
    dplyr::filter(
      OWNER %like% "NATIONAL PARK",
      str_detect(SUMMARY_ST, "A|P")) %>% # approved / perfected applications
      #str_detect(USES,"D|M|O|I")) %>%  # filter by POD use category
    dplyr::distinct(WRNUM,
                    .keep_all = TRUE) %>%
    mutate(name = source)
      
      #if (nrow(Suppliers) >= 1) {
      #  POD_supply <- POD_supply %>%
      #    left_join(., Suppliers, by = WRID)
      #}
}

# All PODS within park
POD_all <- 
  POD_state %>% 
  dplyr::filter(park_right == "Park") %>% 
  dplyr::distinct(WRNUM, .keep_all = TRUE)


mapview(POD_state,
        #zcol = "CFS",
         layer.name = "State Reported PODs",
         col.regions = pal[4:6],
        homebutton = FALSE,
        cex = 3) +
mapview(source_table_locs,
        zcol = "water_system_name",
        layer.name = "Source:",
        homebutton = FALSE,
        col.regions = c("tomato","dodgerblue"),
        cex = 4) + 
  mapview(park_boundary,
          col.regions = "forestgreen",
          alpha.regions = 0.2,
          homebutton = FALSE,
          legend = FALSE) +
  mapview(source_table_locs$aoi %>% bind_rows(),
          zcol = "type",
          layer.name = "AOIs",
          col.regions = pal,
          homebutton = FALSE,
          alpha.regions = 0.2)
  


```

### 5. Nearest NWIS Sites

For many parks, the nearest USGS stream gage is far away. Therefore, here we
pull in all NWIS gages within 100 km of the park boundary. Of those, we only 
select stream gages that are considered "reference" gages in the 
[GAGES-II database (Falcone, 2011)](https://pubs.usgs.gov/publication/70046617).
We then select the gage that is closest to the park with the most complete
period of record and delineate associated watershed using the `get_nldi_basin()` 
function from the {nhdplusTools} package.

```{r, eval = TRUE, message = FALSE, warning = FALSE}

# Get all NWIS sites
nwis <- listNWIS(aoi = watersupply_watershed %>% 
                   dplyr::summarize(), dist = .3) 

# Get NWIS Stream Gages
ref_gages <- get_gagesII(id = nwis$site_no) %>%
             dplyr::filter(class == "Ref")


# Select by maximum overlapping POR (1980-2023) 
nwis_select_stream_gage <- nwis %>%
  dplyr::filter(site_no %in% ref_gages$staid,
                data_type_cd == "dv") %>%
  dplyr::left_join(st_drop_geometry(ref_gages), 
                   by = c("site_no" ="staid")) %>%
  dplyr::mutate(distance = sf::st_distance(geometry, 
                                           watersupply_watershed),
                overlap_por = year(end_date) - ifelse(year(begin_date) < 1980, 1980, year(begin_date))) %>%
  dplyr::slice_max(., overlap_por) # grab nearest only


# Get watersheds associated with nearest stream gage
nwis_select_watershed <- 
  nwis_select_stream_gage$site_no %>%
  purrr::map_dfr(~nldi_finder(site_no = .)) %>%
  dplyr::mutate(data = map(site_no, ~nldi_meta(site_no = .))) %>%
  unnest(cols = c(data)) %>%
  dplyr::left_join(st_drop_geometry(nwis_select_stream_gage), by = "site_no")


# Download data from reference stream sites
nwis_select_discharge <- 
  dataRetrieval::readNWISdv(siteNumbers = nwis_select_stream_gage$site_no,
                            parameterCd = c('00060','00065')) %>%
    dplyr::rename(c("discharge" = "X_00060_00003",
                    "date" = "Date")) %>%
    dplyr::filter(year(date) >= 1980) %>%
    dplyr::group_by(site_no) %>%
    dplyr::mutate(dev_mean = discharge - mean(discharge, na.rm = TRUE)) %>% 
    dplyr::select(date, site_no, discharge, dev_mean) %>%
  dplyr::mutate(
    discharge_in_nwis = 86400 * 12 * discharge /
      (nwis_select_watershed$drain_sqkm * 1.076e+7))



# Now Get Gw sites
#62611 (Groundwater level above NAVD 1988, feet)
#72019 (Depth to water level, feet below land surface)
nwis_groundwater <- nwis %>%
  dplyr::filter(begin_date != end_date,
                year(end_date) > 1985,
                n_obs > 50,
                code == 62611) %>%
  dplyr::mutate(dist = st_distance(geometry,park_boundary) %>%
                as.numeric()) %>%
  #dplyr::filter(dist <= 1600*3) %>%
  add_gw_meta()


# pull those sites groundwater level data and convert to monthly mean
  nwis_groundwater_levels <- 
    dataRetrieval::readNWISgwl(nwis_groundwater$site_no) %>%
    dplyr::filter(parameter_cd == 62611,
                  year(as.Date(lev_dt)) >= 1980) %>% # 72019 =Depths, 62611=elevation
    dplyr::mutate(ym = lubridate::ym(substr(lev_dt, 1, 7))) %>% 
    dplyr::group_by(ym, site_no) %>%
    dplyr::summarize(mean_lev_va = mean(sl_lev_va, na.rm. = TRUE), # elevation
                     #mean_lev_va = mean(lev_va, na.rm. = TRUE), # depths
                     .groups = "keep") %>%
    dplyr::select(ym, site_no, mean_lev_va) %>%
    dplyr::group_by(site_no) %>%
    dplyr::mutate(dev_mean = mean_lev_va - mean(mean_lev_va, na.rm = TRUE))
    #tidyr::pivot_wider(names_from = site_no, names_prefix = "well_", 
                       #values_from = mean_lev_va) #%>%

ggarrange(  
  ggplot(nwis_select_discharge, 
         aes(x = date, y = discharge, color = site_no)) +
    geom_line() +
    xlim(as.Date("1980-01-01"), as.Date("2025-01-01")) +
    scale_color_manual("NWIS Site #", values = pal[2:6]) +
    theme_bw() +
    labs(x = "", y = "Discharge (cfs)"),
  ggplot(nwis_groundwater_levels, 
       aes(x = ym, y = dev_mean, color = site_no)) + 
    geom_line() + 
    scale_color_manual("NWIS Site #", values = pal) + 
    theme_bw() +
    labs(x = "", y = "GW Lev (ft) deviation from mean"),
  ncol = 1,
  align = "hv"
)

  ggplot(nwis_groundwater_levels, 
       aes(x = ym, y = dev_mean, color = site_no)) + 
    geom_line() + 
    scale_color_manual("NWIS Site #", values = pal) + 
    theme_bw() +
    #labs(x = "", y = "GW Lev (ft) deviation from mean") %>% 
    facet_wrap(~site_no)


 
mapview(nwis_groundwater, 
        #zcol = "site_no", 
        col.regions = "yellow",
        cex = 4,
        homebutton = FALSE,
       layer.name = "NWIS GW Sites") +
mapview(nwis_select_stream_gage,
        #zcol = "site_no",
        layer.name = "Mammoth Creek NWIS Gage",
        col.regions = "hotpink",
        homebutton = FALSE,
        cex = 4) +
  mapview(nwis_select_watershed,
          legend = FALSE,
          homebutton = FALSE,
          col.regions = "hotpink",
          alpha.regions = 0.3) +
mapview(source_table_locs,
        zcol = "water_system_name",
        layer.name = "Source:",
        homebutton = FALSE,
        col.regions = c("tomato","dodgerblue"),
        cex = 4) + 
  mapview(park_boundary,
          col.regions = "forestgreen",
          alpha.regions = 0.2,
          homebutton = FALSE,
          legend = FALSE) +
  mapview(source_table_locs$aoi %>% bind_rows(),
          zcol = "type",
          layer.name = "AOIs",
          col.regions = pal,
          homebutton = FALSE,
          alpha.regions = 0.2)
  
#path <- "/Users/kcognac/Desktop/BRCA_Deep_Dive/Site_map.jpg"

#mapshot(
#  a,
#  file = path,
#  remove_controls = c("zoomControl", "layersControl", "homeButton", 
#    "drawToolbar", "easyButton"))


```

### 6. Park supplied well data

Here, we pull in the daily well level data for both supply wells associated
with the East Creek Water system.  This data was digitized from manual field
notes that were made by park staff.

```{r well_dat, eval = TRUE, message = FALSE, warning = FALSE}

well_data <- read_csv('data/park/BRCA/manual/BRCA_Well_Data.csv', na = c("NaN", "NA", "")) %>%
  clean_names() %>%
  dplyr::mutate(static_in = ifelse(static_in == "-", NA,
                              ifelse(static_in == "NaN", NA, 
                                as.numeric(static_in))),
                date = mdy(date),
                static_ft = -static_in/12) %>%
  # Remove some outliers / bad data (next-day level change > 1.5 ft)
   mutate(diff = abs(static_ft - lag(static_ft, default = NA)),
         static_ft = ifelse(diff > 1.5, NA, static_ft)) %>%
  dplyr::select(date,static_ft,well) %>%
  complete(date = seq.Date(min(date),max(date), by = "day")) %>% 
  arrange(date) %>%
  dplyr::mutate(static_ft_c = static_ft) %>%
  group_by(well) %>%
  mutate(static_ft_c = na.approx(static_ft, na.rm = FALSE)) %>%
  dplyr::mutate(aq_volume_ft3 = 
                  map_dbl(static_ft_c, ~aquifer_volume(dtw_ft = .)))
 
well_data_an <- well_data %>%
  dplyr::mutate(year = year(date)) %>%
  dplyr::group_by(year, well) %>%
  dplyr::summarize(static_ft_c = mean(static_ft_c)) %>%
  dplyr::mutate(aq_volume_ft3 = 
                  map_dbl(static_ft_c, ~aquifer_volume(dtw_ft = ., full_thickness_ft = 30)))

# Plot Well data
ggplot(well_data) +
  geom_point(aes(x = date, y = static_ft, color = well)) + 
  geom_line(aes(x = date, y = static_ft_c, color = well), linetype = "dashed") +
  theme_bw() +
  labs(y = "Groundwater Depth (ft)", x = "") +
  scale_color_manual("", values = c("dodgerblue","tomato")) 





#ggsave("/Users/kcognac/Desktop/BRCA_Deep_Dive/GW_Levels.jpg", dpi = 600, width = 6,
#  height = 3,
#  units = c("in"))


```


### 7. Selected Climate Futures

Climate futures were previously compiled for park and Koppen centroids. This
data was used to select which of the CMIP5 climate models to use to represent
"hot dry" and "warm wet" scenarios. 

Source for Koppen-Geiger climate classification maps:

Beck, H. E., Zimmermann, N. E., McVicar, T. R., Vergopolan, N., Berg, A., & 
Wood, E. F. (2018). Present and future Köppen-Geiger climate classification maps
at 1-km resolution. Scientific data, 5(1), 1-12.

<https://figshare.com/articles/dataset/Present_and_future_K_ppen-Geiger_climate_classification_maps_at_1-km_resolution/6396959/2>


```{r, eval = TRUE, echo = FALSE, message = FALSE, warning = FALSE}


select_cfs <- data.table::fread('data/parkwide_gcms_wbm_filtered.csv') %>%
  dplyr::filter(park %in% {{park}},
                  CF %in% c("Warm Wet", "Hot Dry")) %>%
  distinct(GCM, .keep_all = TRUE)

# If there is more than one model selection for a CF scenario, grab the more
# divergent.
if (nrow(select_cfs) > 2) {
  if (nrow(select_cfs %>% dplyr::filter(CF == "Hot Dry")) > 1) {
        select_cfs <- select_cfs %>%
                group_by(CF) %>%
                slice_min(delta_pr)
  }
  if ((nrow(select_cfs %>% dplyr::filter(CF == "Warm Wet")) > 1) ) {
          select_cfs <- select_cfs %>%
            group_by(CF) %>%
            slice_max(delta_pr)
  }
}



```


### 8. Climate data 

Get historical (GridMET) and future (MACA GCMs) climate data for the watersupply
watershed (source area).

```{r, eval = TRUE, echo = FALSE, message = FALSE, warning = FALSE}

# Get historical climate data for water source area
clim_source_hist <- 
  get_climate_historic(
    sf = watersupply_watershed,
    col_name = "name",
    start = "1979-01-01",
    end = "2023-12-31"
    ) %>%
  # convert daily values for each grid cell to daily mean for all cells
  dplyr::mutate(CF = "Historical")
                                        

# Get future climate data for water source area

tictoc::tic()
clim_source_fut <- 
  get_climate_future(
    sf = watersupply_watershed,
    col_name = "name",
    start = "2006-01-01",
    end = "2070-01-01",
    GCM = select_cfs$GCM
    ) %>%
  dplyr::mutate(GCM = paste0(GCM,'.',RCP)) %>%
  left_join(select_cfs %>% dplyr::select(c("GCM","CF")), by = c("GCM")) %>%
  dplyr::select(-c("GCM","Ensemble","RCP")) 

tictoc::toc()



# Get historical climate data for nearest NWIS stream gage
clim_nwis <- 
  get_climate_historic(
    sf = nwis_select_watershed,
    col_name = "site_no",
    start = nwis_select_discharge$date %>% min(),
    end = "2023-01-01") %>%
  dplyr::mutate(CF = "Historical") %>%
  # Join wbm results with discharge
  left_join(., nwis_select_discharge , by = c("date", "site_no")) 

# Get unique points from each climate dataset
source_fut_pts <- clim_source_fut %>% 
  ungroup() %>% 
  distinct(x,y, .keep_all = TRUE) %>% 
  st_as_sf(., coords = c("x","y"), crs = st_crs(4326), remove = FALSE) %>%
  dplyr::select(c(date,CF,name,x,y))

source_hist_pts <- clim_source_hist %>% 
  ungroup() %>% 
  distinct(x,y, .keep_all = TRUE) %>% 
  st_as_sf(., coords = c("x","y"), crs = st_crs(4326), remove = FALSE) %>%
  dplyr::select(c(date,CF,name,x,y))

nwis_pts <- clim_nwis %>% 
  ungroup() %>% 
  distinct(x,y, .keep_all = TRUE) %>% 
  st_as_sf(., coords = c("x","y"), crs = st_crs(4326), remove = FALSE) %>%
  dplyr::select(c(date,CF,site_no,x,y))

# Plot
mapview(source_fut_pts, 
        col.regions = "black",
        layer.name = "MACA") + 
  mapview(watersupply_watershed, 
          col.regions = "dodgerblue", 
          alpha.regions = 0.2,
          layer.name = "Source Watershed") + 
  mapview(source_hist_pts, 
          col.regions = "red",
          layer.name = "GridMET") +
mapview(source_hist_pts %>% 
          st_transform(., 32612) %>% 
          st_buffer(, dist = 2000),
        col.regions = "red",
        alpha.regions = 0.2,
        legend = FALSE) +
  mapview(nwis_pts,
          col.regions = "red",
          legend = FALSE) +
  mapview(nwis_select_watershed,
          col.regions = "seagreen",
          alpha.regions = 0.3,
          layer.name = "NWIS Watershed")

```


Exploration of bias correction

```{r}

distances <- st_distance(source_fut_pts,source_hist_pts)
closest_matches <- apply(distances,1,which.min)
source_fut_pts$id <- paste0("pt_",1:nrow(source_fut_pts))
source_hist_pts$id <- paste0("pt_",closest_matches)

# Regression using daily values for each grid cell
csf <- left_join(clim_source_fut, 
                 source_fut_pts %>% 
                   st_drop_geometry() %>% 
                   dplyr::select(x,y,id), 
                 by = c("x","y")) %>% 
  dplyr::filter(year(date) < 2023) %>% 
  ungroup() %>% 
  dplyr::select(-c(x,y)) %>% 
  rename_with(~paste0(., "_csf"),-c("date", "id","name")) %>%
  dplyr::mutate(id = as.factor(id))
  #dplyr::mutate(ym = ym(paste0(year(date), "-", month(date)))) %>%
  #dplyr::select(-date) %>%
  #dplyr::group_by(ym,name,CF_csf,id) %>% 
  #dplyr::summarize_all(sum) 

csh <-left_join(clim_source_hist, 
                source_hist_pts %>% 
                  st_drop_geometry() %>% 
                  dplyr::select(x,y,id), by = c("x","y")) %>% 
  dplyr::filter(year(date) > 2005) %>% 
  ungroup() %>% 
  dplyr::select(-c(x,y)) %>% 
  rename_with(~paste0(., "_csh"),-c("date", "id","name")) %>%
  dplyr::mutate(id = as.factor(id)) #%>%
 # dplyr::mutate(ym = ym(paste0(year(date), "-", month(date)))) %>%
  #dplyr::select(-date) %>%
  #dplyr::group_by(ym,name,CF_csh,id) %>% 
  #dplyr::summarize_all(sum)

bias <- left_join(csf,csh, by = c("date", "id","name")) %>% dplyr::arrange(ppt_mm_csf)

bias_ww <- bias %>% 
  dplyr::filter(CF_csf == "Warm Wet") %>%
  dplyr::arrange(ppt_mm_csf) %>%
  dplyr::mutate(ppt_csh_sort = sort(ppt_mm_csh))

bias_hd <- bias %>% 
  dplyr::filter(CF_csf == "Hot Dry") %>%
  dplyr::arrange(ppt_mm_csf) %>%
  dplyr::mutate(ppt_csh_sort = sort(ppt_mm_csh))

ggplot() +
  geom_point(data = bias_hd,
             aes(x = ppt_csh_sort, y = ppt_mm_csf, color = "Hot Dry")) +
  stat_smooth(data = bias_hd, 
              aes(x = ppt_csh_sort, y = ppt_mm_csf, color = "Hot Dry"),
              method = "lm", se = FALSE, formula = y~x, geom = "line", alpha = 0.5) +
  geom_point(data = bias_ww,
             aes(x = ppt_csh_sort, y = ppt_mm_csf, color = "Warm Wet")) +
  stat_smooth(data = bias_ww, 
              aes(x = ppt_csh_sort, y = ppt_mm_csf, color = "Warm Wet"),
              method = "lm", se = FALSE, formula = y~x, geom = "line", alpha = 0.5) +
  theme_bw() + 
  scale_color_manual("", values = c("tomato", "dodgerblue")) +
  geom_abline() +
  facet_wrap(~factor(id)) +
  labs(x = "GridMET Precip (mm)", y = "MACA Precip (mm)")
  


 # Bias correction using daily mean of all grid cells

csf <- clim_source_fut %>%
  dplyr::filter(year(date) < 2023) %>% 
  dplyr::ungroup() %>%
  dplyr::select(-c(x,y)) %>%
  group_by(date, name, CF) %>%
  dplyr::summarize_all(mean) %>%
  rename_with(~paste0(., "_csf"),-c("date","name")) 

csh <- clim_source_hist %>%
  dplyr::filter(year(date) < 2023) %>% 
  dplyr::ungroup() %>%
  dplyr::select(-c(x,y)) %>%
  group_by(date, name, CF) %>%
  dplyr::summarize_all(mean) %>%
  rename_with(~paste0(., "_csh"),-c("date","name")) 


bias <- left_join(csf,csh, by = c("date","name")) %>% ungroup()
bias_ww <- bias %>% 
  dplyr::filter(CF_csf == "Warm Wet") %>%
  dplyr::arrange(ppt_mm_csf) %>%
  dplyr::mutate(ppt_csh_sort = sort(ppt_mm_csh))
bias_hd <- bias %>% 
  dplyr::filter(CF_csf == "Hot Dry") %>%
  dplyr::arrange(ppt_mm_csf) %>%
  dplyr::mutate(ppt_csf_sort = sort(ppt_mm_csf),
                ppt_csh_sort = sort(ppt_mm_csh))


ggplot() +
  geom_point(data = bias_hd, aes(x = ppt_csh_sort, y = ppt_mm_csf, color = "Hot Dry")) +
  geom_point(data = bias_ww, aes(x = ppt_csh_sort, y = ppt_mm_csf, color = "Warm Wet")) +
  scale_color_manual("", values = c("tomato", "dodgerblue","black")) +
  geom_abline() +
  theme_bw() +
  labs(x = "GridMET Precip (mm)", y = "MACA Precip (mm)")

```

### 9. WBM 

#### 9.1 Define Run_WBM Fun

This will be external eventually. Leaving here now because it's constantly
being modified for testing at this time.

```{r run_wbm, eval = TRUE, echo = FALSE, message = FALSE, warning = FALSE}

# Load DEM and Rasters
  
# DEM (from NPS WBM group) -> in meters
dem <- terra::rast(here::here('data/all/elevation_cropped.tif')) %>%
  terra::project(crs(nwis_pts))

# storage properties raster (from NPS WBM group)

soil <- terra::rast(here::here('data/all/water_storage.tif')) %>%
  terra::project(crs(nwis_pts)) * 10# cm to mm

# Calculate slope & aspect
# For slope, 4 better for "smooth" surfaces, 8 better for rough See:
# https://www.rdocumentation.org/packages/raster/versions/3.1-5/topics/terrain
slope <- terra::terrain(dem, v = "slope", 
                        unit = "degrees", 
                        neighbors = 8) 

aspect <- terra::terrain(dem, 
                         v = "aspect", 
                         unit = "degrees")

# Hock, 2003 (https://doi.org/10.1016/S0022-1694(03)00257-9)
# Melt factor for Goosebury Creek in Utah is 2.5

# KEC Description of CSU additions to run_NPS_WBM()
#' @param Direct_Frac - Fraction of rain that gets directly routed to runoff
#' @param Return_Rate - Rate that cumulative storage term returns to runoff
#' @param PET_mult - Multiplier to alter PET (e.g., if known bias)
#' @param Soil_mult - Multiplier to alter Soil storage capacity 
#' 
#' When set to default values (below), the model should run identical to the
#' unmodified version: 
#' Direct_Frac = 0,
#' Return_Rate = 1, 
#' PET_mult = 1, and 
#' Soil_mult = 1

run_NPS_wbm <- function(points, 
                        climate_data, 
                        col_name,
                        Direct_Frac = 0, # CSU ADD
                        Return_Rate = 1, # CSU ADD
                        PET_mult = 1,  # CSU ADD
                        Soil_mult = 1, # CSU ADD
                        PET_Method = c("Oudin"), 
                        hock_coef = 4, 
                        Snowpack.Init = 0,   # according to Mike's code
                        Soil.Init = 0, # Init full from Mike's code, empty in Ambers
                        Shade.Coeff = 1,
                        T.Base = 0) {

  # Code for testing:
  #points <- source_hist_pts
  #climate_data <- clim_source_hist
  #col_name <- "name"
  #PET_Method <- "Oudin" 
  #hock_coef <-  4 
  #Snowpack.Init <- 0   # according to Mike's code
  #Soil.Init <- 0
  #Shade.Coeff <- 1
  #T.Base <- 0
  #Soil_mult <- 1
  #Direct_Frac <- 0.3
  #Return_Rate <- 0.2
  
#' Run NPS water balance model for a given set of points and associated climate
#' data.   

# Extract params for points
print("Extracting params at points...")
  points2 <- 
    points %>%
    data.table() %>%
    dplyr::mutate(pt_num = row_number(),
                  Elev = terra::extract(dem, points)$elevation_cropped,
                  Aspect = terra::extract(aspect, points)$aspect,
                  Slope = terra::extract(slope, points)$slope,
                  SWC.Max =  Soil_mult * terra::extract(soil, points)$water_storage,
                  Snowpack.Init = Snowpack.Init,
                  Soil.Init = Soil.Init, 
                  Shade.Coeff = Shade.Coeff)

  climate_data <- climate_data %>%
    data.table() %>%
    left_join(., points2 %>% dplyr::select(x,y,pt_num) %>% st_drop_geometry(), by = c("x","y"))
    
  CFs <- climate_data[["CF"]] %>% unique()
  
# Now, run WBM for each point & GCM
  DailyWB_ret <- vector(mode = "list", 
                        length = nrow(points2)*length(CFs))
  ct <- 0
  
print("Running wbm...")
for (j in 1:length(CFs)) {

  for (i in 1:nrow(points2)) {
    ct <- ct + 1
    #print("Running WBM for grid cell:")
    #print(i)
    #i <- 1
    point <- points2[i,]  
    DailyWB <- climate_data %>% 
      dplyr::filter(pt_num == point$pt_num,
                    CF == CFs[j]) %>%
      data.table()
    
    DailyWB$doy <- yday(DailyWB$date)
    DailyWB$daylength <- get_daylength(DailyWB$date, point$y)
    DailyWB$jtemp = as.numeric(get_jtemp(point$y, point$x))
    DailyWB$F = get_freeze(DailyWB$jtemp, DailyWB$tmean_C)
    DailyWB$RAIN = get_rain(DailyWB$ppt_mm, DailyWB$F)
    DailyWB$SNOW = get_snow(DailyWB$ppt_mm, DailyWB$F)
    DailyWB$MELT = get_melt(DailyWB$tmean_C, DailyWB$jtemp, hock=hock_coef, DailyWB$SNOW, point$Snowpack.Init)
    DailyWB$PACK = get_snowpack(DailyWB$jtemp, DailyWB$SNOW, DailyWB$MELT)
    DailyWB$DIRECT = DailyWB$RAIN*Direct_Frac # CSU ADD
    
    #DailyWB$W = DailyWB$MELT + DailyWB$RAIN
    DailyWB$W = DailyWB$MELT + DailyWB$RAIN - DailyWB$DIRECT # CSU ADD
    
    if(PET_Method == "Hamon"){
      DailyWB$PET = ET_Hamon_daily(DailyWB)
    } else {
      if(PET_Method == "Penman-Monteith"){
        DailyWB$PET = ET_PenmanMonteith_daily(DailyWB)
      } else {
        if(PET_Method == "Oudin"){
          DailyWB$PET = get_OudinPET(DailyWB$doy, point$y, DailyWB$PACK, 
                                     DailyWB$tmean_C, point$Slope, 
                                     point$Aspect, point$Shade.Coeff)
        } else {
          print("Error - PET method not found")
        }
      }
    }
    DailyWB$PET_mod <- DailyWB$PET * PET_mult # CSU ADD
    #DailyWB$PET = modify_PET(DailyWB$PET, Slope, Aspect, Lat, Shade.Coeff)
    DailyWB$W_PET = DailyWB$W - DailyWB$PET
    DailyWB$SOIL = get_soil(DailyWB$W, point$Soil.Init, DailyWB$PET_mod, DailyWB$W_PET, point$SWC.Max)
    DailyWB$DSOIL = diff(c(point$Soil.Init, DailyWB$SOIL))
    DailyWB$AET = get_AET(DailyWB$W, DailyWB$PET, DailyWB$SOIL, point$Soil.Init)
    
    #DailyWB$W_ET_DSOIL = DailyWB$W - DailyWB$AET - DailyWB$DSOIL
    DailyWB$STORAGE_ADD = DailyWB$W - DailyWB$AET - DailyWB$DSOIL # CSU ADD
    DailyWB$STORAGE_RELEASE = get_storage(DailyWB$STORAGE_ADD,
                                          Return_Rate)["storage_release"] # CSU ADD
    DailyWB$W_ET_DSOIL =  DailyWB$STORAGE_RELEASE + DailyWB$DIRECT
    DailyWB$STORAGE_REMAIN = get_storage(DailyWB$STORAGE_ADD,
                                         Return_Rate)["storage_remain"] # CSU ADD
    
    DailyWB$D = DailyWB$PET - DailyWB$AET
    DailyWB$GDD = get_GDD(DailyWB$tmean_C, T.Base)
    DailyWB_ret[[ct]] <- DailyWB %>%
      ungroup() %>%
      dplyr::select(date, x, y, CF, pt_num, ppt_mm, 
                    "RUNOFF" = W_ET_DSOIL, RAIN, SNOW, MELT,AET,PET_mod, tmean_C,"excess_wat" = STORAGE_ADD) %>%
      pivot_longer(-c(date,x,y,pt_num,CF), 
                   values_to = "vals", 
                   names_to = "vars") %>%
      mutate(temp_name = point[[{{col_name}}]]) %>%
      dplyr::rename(!!col_name := "temp_name")
  }
}
   
# Group and summarize all runs
    DailyWB_ret2 <- DailyWB_ret %>% 
      bind_rows() %>%
      mutate(vals = ifelse(vars == "tmean_C", vals, vals /25.4)) %>% # mm to in
      group_by(date,vars, !!sym(col_name),CF) %>%
      dplyr::select(-c(x,y,pt_num)) %>%
      dplyr::summarize(vals = mean(vals, na.rm = TRUE),
                       .groups = "keep") %>%
       pivot_wider(names_from = vars, values_from = vals) %>%
    rename(., c("RUNOFF" = "runoff_in_wbm", 
                "RAIN" = "rain_in_wbm",
                "SNOW" = "snow_in_wbm",
                "MELT" = "melt_in_wbm",
                "AET" = "aet_in_wbm",
                "ppt_mm" = "precip_in_wbm",
                "PET_mod" = "pet_in_wbm",
                "excess_wat" = "excess_water_in_wbm")) 
  }


```

#### 9.2 Source Area WBM

Run the NPS water balance model for each gridmet point in the AOI and then
aggregate to area mean to return spatially averaged daily value.   Run for
historical and future datasets separately for now because point locations
are slightly different.  KEC --- perhaps rather than running the model 10
times for each grid point (which we've talked about changing to), we could
grab the average for each parameter for 10 sampled points around the select
location... Thoughts?

```{r wbm_run}

# run wbm - results returned in in/day
source_hist_wbm <- run_NPS_wbm(points = source_hist_pts,
                        climate_data = clim_source_hist,
                        col_name = "name",
                        hock_coef = 2,
                        Direct_Frac = 0.1, # CSU ADD
                        Return_Rate = 0.05) # CSU ADD)

source_fut_wbm <- run_NPS_wbm(points = source_fut_pts,
                              climate_data = clim_source_fut,
                              col_name = "name",
                              hock_coef = 2,
                              Direct_Frac = 0.1,
                              Return_Rate = 0.05)

source_wbm <- bind_rows(source_hist_wbm, source_fut_wbm %>% dplyr::filter(year(date) > 2023)) %>%
  left_join(., select_cfs %>% 
              dplyr::select(GCM, CF) %>%
              dplyr::mutate(GCM = str_extract(GCM, "^[^.]+")), 
            by = "CF") %>%
  dplyr::mutate(GCM = ifelse(is.na(GCM),"GridMET",GCM))

#rm("source_hist_wbm","source_fut_wbm")

source_wbm_annual <- source_wbm %>%
  ungroup %>%
  dplyr::mutate(year = year(date)) %>%
  dplyr::select(-date) %>%
  group_by(year,name,GCM,CF) %>%
  dplyr::summarise(melt_in_wbm = sum(melt_in_wbm),
                   rain_in_wbm = sum(rain_in_wbm),
                   runoff_in_wbm = sum(runoff_in_wbm),
                   snow_in_wbm = sum(snow_in_wbm),
                   precip_in_wbm = sum(precip_in_wbm),
                   aet_in_wbm = sum(aet_in_wbm),
                   excess_water_in_wbm = sum(excess_water_in_wbm),
                   tmean_C = mean(tmean_C, na.rm = TRUE),
                   .groups = "keep") %>%
  mutate(run_per = runoff_in_wbm/precip_in_wbm)

source_wbm_monthly <- source_wbm %>%
  ungroup %>%
  dplyr::mutate(ym = as.Date(paste0(year(date),'-',month(date),'-15'))) %>%
  dplyr::select(-date) %>%
  group_by(ym,name,GCM,CF) %>%
  dplyr::summarise(melt_in_wbm = sum(melt_in_wbm),
                   rain_in_wbm = sum(rain_in_wbm),
                   runoff_in_wbm = sum(runoff_in_wbm),
                   snow_in_wbm = sum(snow_in_wbm),
                   precip_in_wbm = sum(precip_in_wbm),
                   aet_in_wbm = sum(aet_in_wbm),
                   excess_water_in_wbm = sum(excess_water_in_wbm),
                   tmean_C = mean(tmean_C, na.rm = TRUE),
                   .groups = "keep") %>%
  mutate(run_per = runoff_in_wbm/precip_in_wbm)

# Example annual plot
ggplot(source_wbm_annual %>%
         dplyr::ungroup() %>%
         dplyr::select(c(year,
                         GCM,
                         CF,
                         "Precipitation" = "precip_in_wbm",
                         "Temperature" = "tmean_C",
                         "AET" = "aet_in_wbm",
                         "Snowmelt" = "melt_in_wbm",
                         "Rain" = "rain_in_wbm",
                         "Runoff" = "runoff_in_wbm")) %>%
         pivot_longer(., -c(year,GCM,CF), names_to = "vars", values_to = "vals"), 
       aes(x = year, y = vals, color = CF, fill = CF)) + 
   geom_line(lwd = .5, alpha = 0.5) +
  geom_smooth(method = "lm") +
  theme_bw() +
  scale_color_manual("",values = c("black","red3","dodgerblue2")) +
  scale_fill_manual("",values = c("black","red3","dodgerblue2")) +
  labs(y = "", x = "") +
  facet_wrap(~ factor(vars, levels = c("Precipitation","Rain","Snowmelt","Temperature","AET","Runoff")),
             scales = "free") + 
  theme(legend.position = "bottom")
  #theme(legend.position=c(.07, .9), legend.title=element_blank())


  ggsave("/Users/kcognac/Desktop/BRCA_Deep_Dive/WBM.jpg", dpi = 600, width = 6,
  height = 3.5,
  units = c("in"))
  

# Examply daily plot
ggplot(source_wbm %>% dplyr::filter(year(date) > 2010, year(date) < 2012), 
       aes(x = date, y = runoff_in_wbm, color = CF, fill = CF)) + 
   geom_line(lwd = 1) +
  theme_bw() +
  scale_color_manual("",values = c("black","red3","dodgerblue2")) +
  scale_fill_manual("",values = c("black","red3","dodgerblue2")) +
  labs(y = "Runoff (in/d)", x = "", 
       title = "Example year when GridMET overlaps MACA. \nNote how monsoon runoff increases while spring runoff decreases\n for Warm Wet future.")

```



#### 9.3 NWIS area WBM

For the NWIS area, we have discharge data that we can use to calibrate the
NPS wbm.  Here, we attempt calibration using two runoff routing methods.
The first is a simple routing performed within the WBM through the "Direct_Frac"
and "Return_Rate" parameters. The next is an IHACREs type runoff routing
function that routes "excess precipitation" using unit-hydrograph type methods
with a recursive filter function.

##### 9.3.1 Manual Calibration 

Use in-line parameters to manually calibrate calibrate to discharge at nearest 
NWIS gage.

```{r nwis_wbm}

# Run WBM with "Direct_Frac" and "Return_Rate" routing parameters.
nwis_wbm <- run_NPS_wbm(points = nwis_pts,
                        climate_data = clim_nwis,
                        col_name = "site_no",
                        hock_coef = 4,
                        PET_mult = 1,
                        Soil_mult = 5,
                        Direct_Frac = 0.05, 
                        Return_Rate = 0.01) %>%
  left_join(., nwis_nearest_discharge, by = c("date","site_no")) %>%
  ungroup()

ggplot(nwis_wbm %>% dplyr::filter(year(date) > 2001, year(date) < 2013)) + 
  geom_line(aes(x = date, y = excess_water_in_wbm, color = "Old Runoff"), alpha = 0.2) +
  geom_line(aes(x = date, y = runoff_in_wbm, color = "New Runoff")) + 
  geom_line(aes(x = date, y = discharge_in_nwis, color = "NWIS Discharge")) +
  scale_color_manual("",values = c("black","tomato","dodgerblue")) +
  theme_bw()
```


##### 9.3.2 Automated Calibration 

Use IHACRES style routing function to calibrate runoff
```{r}

# This section applies the IHACRES routing function AFTER the standard
# NPS WBM runoff is calculated. The runoff is routed to discharge using a 
# calibration routine with test and validation datasets.

# First, run wbm without Direct_Frac and Return_Rate effects
nwis_wbm <- run_NPS_wbm(points = nwis_pts,
                        climate_data = clim_nwis,
                        col_name = "site_no",
                        hock_coef = 2,
                        PET_mult = 1,
                        Soil_mult = 5,
                        Direct_Frac = 0, 
                        Return_Rate = 1) %>%
  left_join(., nwis_nearest_discharge, by = c("date","site_no")) %>%
  ungroup()

# Define test dataset (on which to calibrate)
test_dates <- interval(as.Date("1995-01-01"), as.Date("2005-01-01"))
test_dat <- nwis_wbm %>%
  dplyr::filter(date %within% test_dates)

# And validation dataset (data not used for calibration)
val_dates <- interval(as.Date("2005-01-01"), as.Date("2015-01-01"))
val_dat <- nwis_wbm %>%
  dplyr::filter(date %within% val_dates)

# Now, calibrate...

# First define fitting function for calibration
#' @param params input parameters for routing function (tau_q_rain, tau_s_rain,
#' vs_rain, tau_s_melt, and delay)
#' @param dat dataframe containing test data for which to calibrate
#' @param return_dat True/False on whether to return daily values or just 
#' objective function. For calibration routine, needs to return only objective
#' function.
#' 

fit_fun <- function(params, 
                    dat = test_dat, 
                    return_dat = FALSE) {
  
  # dataset for testing
  #params <- c(3,90,.5,50,1)
  #dat <- nwis_wbm
  
  Q_mod <-runoff_routing_snow_fun(df = dat, 
                             tau_q_rain = params[1], 
                             tau_s_rain = params[2], 
                             vs_rain = params[3],
                             tau_s_melt = params[4],
                              delay = params[5]
                            )
  
  df_out <- left_join(Q_mod, dat, by = "date") %>%
    drop_na(discharge_in_nwis)
  
  # objective function to minimize during calibration (i.e., error term)
  ret_er <- NSE(sim = df_out$Q_tot, obs = df_out$discharge_in_nwis)
  #ret_er <- SSE(sim = Q_mod$Q_tot, obs = discharge)
  #ret_er <- MAE(sim = Q_mod$Q_tot, obs = discharge)
  
  print(ret_er)
  
  if(return_dat==TRUE){
    Q_mod <- df_out %>% dplyr::select(date,Q_tot,discharge = discharge_in_nwis)
    return(Q_mod) 
  } else {
    return(ret_er)
  }
} 

# Set range of parameters for calibration - i.e., what are the upper and lower
# limits to test during calibration and what are the starting values to feed
# the algorithm.
  #               tau_qr, tau_sr, vsr, Tau_sm, delay
 lower_bounds <- c(  1,    30,   0.01,   10,   0)#,  0.01)
      initial <- c(  5,    80,   0.3,    90,   0)#,  0.5)
upper_bounds <- c(   30,  200,   0.9,   150,  31)#,   1)

# Now, run the optimization routine

optim_par <- 
  optim(
    par = initial, 
    fn = fit_fun,
    method = 'L-BFGS-B',
    lower = lower_bounds, 
    upper = upper_bounds,
    control = c("trace" = 0, "factr" = 1e-6, fnscale = -1))# function


# Now, use optimized parameters to calculate values for test and validation 
# datasets
test_dis <- fit_fun(dat = test_dat,
                   params = optim_par$par,
                   return_dat = TRUE)

val_dis <- fit_fun(dat = val_dat,
                   params = optim_par$par,
                   return_dat = TRUE) %>% 
  left_join(nwis_wbm %>% dplyr::select(-discharge), by = "date")

# Plot results
# https://www.hec.usace.army.mil/confluence/hmsdocs/hmstrm/calibration/calibration-summary-statistics
a <- ggplot(test_dis) +
  geom_line(aes(x = date, y = discharge, color = "NWIS")) +
  geom_line(aes(x = date, y = Q_tot, color = "Modeled")) +
  theme_bw() +
  scale_color_manual("",values = c("tomato","black")) +
  labs(title = paste0("Test Data\n    NSE: ", 
                      round(NSE(sim = test_dis$Q_tot, obs = test_dis$discharge),5),
                      "\n    RSR: ", 
                      round(RSR(sim = test_dis$Q_tot, obs = test_dis$discharge),5),
                      "\nP BIAS: ",
                      round(PBIAS(sim = test_dis$Q_tot, obs = test_dis$discharge),2),
                      "\n    COD: ",
                      round(COD(sim = test_dis$Q_tot, obs = test_dis$discharge),5)),
       x = "")



b <- ggplot(val_dis) +
  geom_line(aes(x = date, y = discharge, color = "NWIS")) +
  geom_line(aes(x = date, y = Q_tot, color = "Modeled")) +
  #geom_line(aes(x = date, y = runoff_in_wbm, color = "Runoff"),alpha = 0.5) +
  theme_bw() +
  scale_color_manual("",values = c("tomato","black", "dodgerblue")) +
  labs(title = paste0("Validation Data\n    NSE: ", 
                      round(NSE(sim = val_dis$Q_tot, obs = val_dis$discharge),5),
                      "\n    RSR: ", 
                      round(RSR(sim = val_dis$Q_tot, obs = val_dis$discharge),5),
                      "\nP BIAS: ",
                      round(PBIAS(sim = val_dis$Q_tot, obs = val_dis$discharge),2),
                      "\n    COD: ",
                      round(COD(sim = val_dis$Q_tot, obs = val_dis$discharge),5)),
       x = "", y = "Discharge (in/d)") 

ggarrange(a,b)


```



##### 9.3.4 GW Calibration

Here is a redundant exercise to above where I try to calibrate the routing 
function to groundwater levels. Heads up, it doesn't work. However, next up is 
to test the storage term..
```{r fit_to_gw_lev}

nwis_wbm <- run_NPS_wbm(points = nwis_pts,
                        climate_data = clim_nwis,
                        col_name = "site_no",
                        hock_coef = 2,
                        PET_mult = 1,
                        Soil_mult = 5,
                        Direct_Frac = 0.15, 
                        Return_Rate = 0.02) %>%
  left_join(., nwis_nearest_discharge, by = c("date","site_no")) %>%
  ungroup()


source_dates <- interval(as.Date("2004-01-01"), as.Date("2023-01-01"))
source_dat <- source_wbm %>%
  data.table() %>%
  dplyr::filter(date %within% source_dates) %>%
  left_join(., well_data %>% dplyr::filter(well == "Well 1"), by = "date") 


# Fitting function for calibration
fit_fun_gw <- function(params, 
                    dat = source_dat, 
                    return_dat = FALSE) {
  
  #params <- c(3,90,.5,50,1)
  #dat <- source_dat
  
  # Clean gw level data for calibration - assume linear
  # interpolation through time is okay 
  gw_lev <- dat[["static_ft"]] 
  if (is.na(gw_lev[1])) {
    gw_lev[1] <- first(gw_lev[!is.na(gw_lev)])
  }
  
  if (is.na(gw_lev[length(gw_lev)])) {
    gw_lev[length(gw_lev)] <- last(gw_lev[!is.na(gw_lev)])
  }
  
  gw_lev <- zoo::na.approx(gw_lev, na.rm = FALSE)
  gw_lev <- gw_lev - gw_lev[1]
  
  area_factor <- .01
  gw_lev <- (gw_lev)*.15*area_factor  
  
  
  dat$gw_clean <- gw_lev
  
  Q_mod <-runoff_routing_snow_fun(df = dat, 
                             tau_q_rain = params[1], 
                             tau_s_rain = params[2], 
                             vs_rain = params[3],
                             tau_s_melt = params[4],
                              delay = params[5])
  
  df_out <- left_join(Q_mod, dat, by = "date")
  
  ret_er <- NSE(sim = df_out$Q_tot, obs = df_out$gw_clean)
  #ret_er <- SSE(sim = Q_mod$Q_tot, obs = discharge)
  #ret_er <- MAE(sim = Q_mod$Q_tot, obs = discharge)
  
  print(ret_er)
  
  if(return_dat==TRUE){
    Q_mod <- df_out %>% dplyr::select(date,Q_tot, gw_clean)
    return(Q_mod) 
  } else {
    return(ret_er)
  }
} 


 # daily data
  #               tau_qr, tau_sr, vsr, Tau_sm, delay
 lower_bounds <- c(  30,    60,   0.1,     1,  0.00001)
      initial <- c(  50,    300,   0.5,     20,  0.000012)
upper_bounds <- c(   90,  360,   0.99,   360,   0.000013)

# daily data
  #               tau_qr, tau_sr, vsr,  delay
# lower_bounds <- c(  1,    10,   0.01,     0)#,  0.01)
#      initial <- c(  5,    80,   0.3,      1)#,  0.5)
#upper_bounds <- c(   10,  150,   0.9,    31)#,   1)



# tau_q, tau_s, vs, delay, loss
optim_par <- 
  optim(
    par = initial, 
    fn = fit_fun_gw,
    method = 'L-BFGS-B',
    lower = lower_bounds, 
    upper = upper_bounds,
    control = c("trace" = 0, "factr" = 1e-6, fnscale = -1))# function


source_lev <- fit_fun_gw(dat = source_dat,
                   params = optim_par$par,
                   return_dat = TRUE)


# https://www.hec.usace.army.mil/confluence/hmsdocs/hmstrm/calibration/calibration-summary-statistics
ggplot(source_lev) +
  geom_line(aes(x = date, y = gw_clean, color = "Observed GW Level")) +
  geom_line(aes(x = date, y = Q_tot, color = "Model")) +
  theme_bw() +
  scale_color_manual("",values = c("tomato","black")) +
  labs(title = paste0("    NSE: ", 
                      round(NSE(sim = source_lev$Q_tot, obs = source_lev$gw_clean),5),
                      "\n    RSR: ", 
                      round(RSR(sim = source_lev$Q_tot, obs = source_lev$gw_clean),5),
                      "\nP BIAS: ",
                      round(PBIAS(sim = source_lev$Q_tot, obs = source_lev$gw_clean),2),
                      "\n    COD: ",
                      round(COD(sim = source_lev$Q_tot, obs = source_lev$gw_clean),5)))




```


Explore sources of runoff
```{r water_runoff}

ggplot(val_dis %>% 
         dplyr::filter(year(date) > 2010, 
                       year(date) < 2012) %>%
         dplyr::mutate(melt_frac = (melt_in_wbm / (melt_in_wbm + rain_in_wbm)) * runoff_in_wbm,
                       rain_frac = (rain_in_wbm / (melt_in_wbm + rain_in_wbm)) * runoff_in_wbm)) +
  geom_col(aes(x = date, y = melt_frac+rain_frac, color = "Snowmelt Runoff"), 
           fill = "pink",lwd = 1, alpha = 0.5) +
  geom_col(aes(x = date, y = rain_frac, color = "Rain Runoff"), 
           fill = "dodgerblue",lwd = 1, alpha = 0.5) +
  geom_line(aes(x = date, y = discharge, color = "Discharge"), lwd = 1) +
  theme_bw() +
  scale_color_manual("",values = c("black", "dodgerblue", "pink")) + 
  scale_fill_manual("", values = c("dodgerblue","pink")) +
  labs(x = "", y = "Flow (in/d)", title = "Mammoth Creek") 



```

### 11. Park Visitation

```{r, eval = TRUE, echo = FALSE, message = FALSE, warning = FALSE}
# NPS monthly visitor information
visitors <- 
  getUnitVisitation(units = park, startYear = 1980, endYear = 2023) %>%
  dplyr::mutate(ym = ym(paste0(Year, "-", Month))) %>%
  dplyr::mutate(TotalVisitors = RecreationVisitors + NonRecreationVisitors) %>%
  dplyr::select(ym, RecreationVisitors, TotalVisitors) %>%
  mutate(ifelse(TotalVisitors == 0,NA,TotalVisitors),
         ifelse(RecreationVisitors ==0, NA, RecreationVisitors))

visitors_an <- visitors %>%
  dplyr::mutate(year = year(ym)) %>%
  dplyr::group_by(year) %>%
  dplyr::select(-ym) %>%
  dplyr::summarize_all(., .funs = sum)


```

### 12. Water use data

```{r, eval = TRUE, echo = FALSE, message = FALSE, warning = FALSE}

# Get state specific water use data --> in monthly format

if (state == "UT") {

  water_supply_id <- 
    getWaterSuppliersUtah(aoi = park_boundary) %>%
    filter(grepl("National", WRNAME, ignore.case=TRUE) | 
           grepl("National", WRENAME, ignore.case = TRUE)) %>%
    .$WRID

  if (!sum(water_supply_id) == 0) {
  
   water_use <- 
     map_dfr(water_supply_id, 
                     \(x) getWaterUseUtah(WRID = x)[[1]] %>%
      dplyr::filter(!is.na(suppressWarnings(as.numeric(Year)))) %>%
      tidyr::pivot_longer(-c("Year","Method of Measurement"), 
                          names_to = "month", 
                          values_to = "use_acre_feet") %>%
      dplyr::filter(month != "Annual inAcre Feet") %>%
      dplyr::mutate(ym = ym(paste0(Year, "-", month))) %>%
      drop_na(ym) %>%
      dplyr::group_by(ym) %>%
      dplyr::summarize(use_acre_feet = sum(as.numeric(use_acre_feet), 
                                           na.rm = TRUE)) %>%
        mutate(WRID=paste0("WRID_",x))) %>%
     pivot_wider(names_from = WRID, values_from = use_acre_feet) %>%
     mutate(use_acre_feet = rowSums(dplyr::select(., !starts_with("ym")), 
                                    na.rm = TRUE)) %>%
     dplyr::mutate(use_ft3 = use_acre_feet* 43559.9)
  
  } else if (water_supply_id == 0 ) {
    water_use = Visitors %>%
      mutate(water_use_gal = TotalVisitors*5)
    }

}

water_use <- water_use %>%
  left_join(., visitors, by = "ym")

# Calculate average monthly useage
use_mon <- water_use %>% 
         dplyr::filter(year(ym) >= 2000) %>%
         dplyr::mutate(month = month(ym)) %>% 
         group_by(month) %>% 
         dplyr::select(-ym) %>% 
         dplyr::summarize(use_acre_feet = mean(use_acre_feet, na.rm = TRUE),
                          TotalVisitors = mean(TotalVisitors, na.rm = TRUE), .groups = "keep") %>%
         dplyr::mutate(use_ft3 = use_acre_feet * 43559.9) %>%
         dplyr::ungroup() %>%
        dplyr::mutate(cum_use = cumsum(use_ft3)) 

use_an <- water_use %>%
  dplyr::mutate(year = year(ym)) %>%
  dplyr::group_by(year) %>%
  dplyr::summarize(use_acre_feet = sum(use_acre_feet)) %>%
  dplyr::mutate(
    use_acre_feet = ifelse(use_acre_feet == 0, NA, use_acre_feet),
    use_acre_feet = na.approx(use_acre_feet),
    use_ft3 = use_acre_feet * 43559.9) %>%
  left_join(., visitors_an, by = "year")



# Explore water use
# Annual water use
ggplot() +
  #geom_point(data = use_an, aes(x = as.Date(paste0(year,"-01-01")), y = use_ft3, color = "Water Use"), size = 3) +
   geom_line(data = use_an,aes(x = as.Date(paste0(year,"-01-01")), y = use_ft3, color = "Water Use"), lwd = 1) +
#geom_point(data = use_an,aes(x = as.Date(paste0(year,"-01-01")), y = TotalVisitors, color = "Visitors"), size = 3) +
  geom_line (data = use_an,aes(x = as.Date(paste0(year,"-01-01")), y = TotalVisitors, color = "Visitors"), size = 1) +
    theme_bw() +
   scale_y_continuous(labels = comma) +
  labs(x = "", y = "Annual") +
  scale_color_manual("", values = c(pal[1], pal[2], pal[2])) 

#ggsave("/Users/kcognac/Desktop/BRCA_Deep_Dive/Water_use.jpg", dpi = 600, width = 6,
#  height = 2.5,
#  units = c("in"))


# Annual Water Use + visitation
factor = 1
shift = 0
ggplot() +
  geom_point(data = use_an, 
             aes(x = as.Date(paste0(year,"-01-01")), 
                 y = use_ft3, 
                 color = "Water Use"), 
             size = 3) +
   geom_line(data = use_an,
             aes(x = as.Date(paste0(year,"-01-01")), 
                 y = use_ft3, 
                 color = "Water Use"), 
             lwd = 1) +
geom_point(data = use_an,
           aes(x = as.Date(paste0(year,"-01-01")), 
               y = TotalVisitors, 
               color = "Visitors"), 
           size = 3) +
  geom_line (data = use_an,
             aes(x = as.Date(paste0(year,"-01-01")), 
                 y = TotalVisitors, 
                 color = "Visitors"), 
             size = 1) +
  geom_line(data = well_data_an, 
            aes(x = as.Date(paste0(year,"-01-01")), 
                y = factor*(aq_volume_ft3 + shift), 
                color = "Aq Volume"), 
            lwd = 1, 
            alpha = 0.5) +
  geom_point(data = well_data_an, 
             aes(x = as.Date(paste0(year,"-01-01")),
                 y = factor*(aq_volume_ft3+ shift), 
                 color = "Aq Volume"), 
             size = 3,
             alpha = 0.5) +
    theme_bw() +
   scale_y_continuous(sec.axis = sec_axis(~./factor - shift, name = "Aquifer Volume (CF)"), labels = comma) +
  labs(x = "", y = "Annual water use (CF)") +
  scale_color_manual("", values = c(pal[3], pal[1], pal[2])) +
  xlim(as.Date("2000-01-01"), as.Date("2023-01-01"))



ggsave("/Users/kcognac/Desktop/BRCA_Deep_Dive/Water_use_Vis_Depth.jpg", dpi = 600, width = 6,
  height = 3,
  units = c("in"))

# Mean monthly water use
b <- ggplot(use_mon) + 
geom_line(aes(x = as.Date(paste0("2020-",month,"-15")), y = use_ft3, color = "Water Use"), lwd = 1) + 
  geom_line(aes(x = as.Date(paste0("2020-",month,"-15")), y = TotalVisitors, color = "Visitors"), lwd = 1) + 
  scale_y_continuous(labels = comma) +
  scale_x_date(labels = date_format("%B")) +
  theme_bw() + 
  labs(y = "Mean Monthly", x ="") +
  scale_color_manual("",values = pal)

a <- ggplot(water_use) +
  #geom_point(aes(x = ym, y = use_ft3, color = "Water Use"), size = 3) +
   geom_line(aes(x = ym, y = use_ft3, color = "Water Use"), lwd = 1) +
#geom_point(aes(x = ym, y = TotalVisitors, color = "Visitors"), size = 3) +
  geom_line (aes(x = ym, y = TotalVisitors, color = "Visitors"), lwd = 1, alpha = 0.8) +
    theme_bw() +
   scale_y_continuous(labels = comma) +
  labs(x = "", y = "", title = "Monthly") +
  scale_color_manual("", values = c(pal[1],pal[2])) +
  xlim(as.Date("2018-01-01"), as.Date("2024-01-01"))

# Annual water use
b <- ggplot() +
  #geom_point(data = use_an, aes(x = as.Date(paste0(year,"-01-01")), y = use_ft3, color = "Water Use"), size = 3) +
   geom_line(data = use_an,aes(x = as.Date(paste0(year,"-01-01")), y = use_ft3, color = "Water Use"), lwd = 1) +
#geom_point(data = use_an,aes(x = as.Date(paste0(year,"-01-01")), y = TotalVisitors, color = "Visitors"), size = 3) +
  geom_line (data = use_an,aes(x = as.Date(paste0(year,"-01-01")), y = TotalVisitors, color = "Visitors"), size = 1) +
    theme_bw() +
   scale_y_continuous(labels = comma) +
  labs(x = "", y = "", title = "Annual") +
  scale_color_manual("", values = c(pal[1], pal[2], pal[2])) 

plot <- ggarrange(a,b, nrow = 1, align = "v", common.legend = TRUE,legend="bottom")

annotate_figure(plot, top = text_grob("Water Use (CF) and Visitors (-)", 
                face = "bold", size = 12))

ggsave("/Users/kcognac/Desktop/BRCA_Deep_Dive/Water_use_m_y.jpg", dpi = 600, width = 6,
  height = 3,
  units = c("in"))



# Mean cumulative water use
 ggplot(use_mon) +
  geom_line(aes(x = as.Date(paste0("2020-",month,"-15")), y = cum_use), lwd = 1) + 
  #geom_line(aes(x = as.Date(paste0("2020-",month,"-15")), y = use_ft3, color = "Monthly"), lwd = 1) + 
  theme_bw() + 
  labs(y = "Water Use (acre-feet)", x = "") + 
  scale_x_date(labels = date_format("%B")) +
  scale_y_continuous(labels = comma) +
  labs(y = "Mean") +
  scale_color_manual("", values = c("black","#619CFF"))
```


### 13 CCVA Indices
#### 13.1 Exposure
##### Runoff Magnitude

Runoff % change =  (RF - RH / RH) where runoff (R) = mean annual runoff.

```{r}

runoff_change <- source_wbm_annual %>%
  mutate(stat_group = ifelse(year < 2010, "Hist",
                             ifelse(year > 2040, "Fut","Present"))) %>%
  dplyr::ungroup() %>%
  dplyr::select(stat_group, CF, runoff_in_wbm) %>%
  dplyr::group_by(stat_group, CF) %>%
  dplyr::summarize_all(., .funs = mean, .groups = "keep") %>%
  dplyr::filter(stat_group != "Present",
                !(stat_group == "Hist" & CF == "Hot Dry"),
                !(stat_group == "Hist" & CF == "Warm Wet")) %>%
  ungroup() %>%
  mutate(hist_runoff = runoff_in_wbm[stat_group == "Hist"],
         perchg = 100 * (runoff_in_wbm - hist_runoff) / hist_runoff)

ggplot(source_wbm_annual, aes(x = year, y = runoff_in_wbm, color = CF)) + 
  geom_line() + 
  geom_hline(data = runoff_change, aes(yintercept = runoff_in_wbm, color = CF), linetype = "dashed") +
  theme_bw() +
  scale_color_manual("", values = c("black","tomato","dodgerblue")) +
  labs(y = "Annual Runoff (in)", x = "")

```

##### Runoff Timing
##### Runoff Frequency
##### Runoff Low-Flow Duration
##### Climate Related Climate Exposure
#### 13.2 Sensitivity
##### Annual Sensitivity
##### Seasonal Sensitivity
##### Source Historic Trends
##### Source water right conflicts


##### Withdrawal Ratio
Demand - Supply Ratio
```{r}

Supply_Annual <- source_wbm_annual %>%
  dplyr::filter(year < 2021  |
                year > 2039) %>%
  ungroup() %>%
  group_by(CF) %>%
  dplyr::select(CF,year,runoff_in_wbm) %>%
  dplyr::mutate(runoff_cf = (runoff_in_wbm/12) * as.numeric(st_area(watersupply_watershed$geometry)) * 10.7639) %>% # m2 to ft2
  dplyr::summarize(mean_supply_cf3 = mean(runoff_cf),
                   min_supply_cf3 = min(runoff_cf))
  



# Data table of maximum use on Annual, Monthly, and Daily basis
Max_Historical_Use <- max(use_an$use_ft3)     # Annual
Mean_Historical_Use <- mean(use_an$use_ft3, na.rm = TRUE)

              
# Aquifer Storage



source_wbm_annual %>%
  dplyr::filter(year > )

kable(SD, escape = FALSE, booktabs = T,
     digits = 3) %>%
  kable_styling(bootstrap_options = c("striped", "hover", "condensed", "responsive")) %>%
kableExtra::landscape() %>%
  #kable_styling(full_width = F) %>%
    kable_paper(full_width = F)



#Supply_Monthly<- source_wbm_monthly %>%
#  dplyr::filter(year(ym) < 2021 |
#                  year(ym) > 2039) %>%
#  ungroup() %>%
#  group_by(CF) %>%
#  dplyr::select(CF, ym, runoff_in_wbm) %>%
#  dplyr::mutate(runoff_cf = (runoff_in_wbm/12) * as.numeric(st_area(watersupply_watershed$geometry)) * 10.7639) %>%
#  dplyr::summarize(mean_supply_cf3 = mean(runoff_cf),
#                   min_supply_cf3 = min(runoff_cf))

```

##### Storage Ratio


```{r storage_demand}

# Storage:Demand Ratio
TP <- c("Annual","Monthly","Daily")
Max_Use = c(max(use_an %>% 
                  dplyr::filter(year > 1999) %>% 
              pull("use_ft3")), 
            max(water_use$use_ft3),
            max(water_use$use_ft3/30))
SD <- data.frame(TP,Max_Use) %>%
  dplyr::mutate(Supply = supply_table$storage_capacity_gallons,
                S_D = Supply/Max_Use)

kable(SD, escape = FALSE, booktabs = T,
     digits = 3) %>%
  kable_styling(bootstrap_options = c("striped", "hover", "condensed", "responsive")) %>%
kableExtra::landscape() %>%
  #kable_styling(full_width = F) %>%
    kable_paper(full_width = F)

```

##### Fire sensitivity

### 15. Data Analysis

  tbd...

### 16.  Example Tables

```{r, eval = TRUE, echo = FALSE, message = FALSE, warning = FALSE}


kable(select_cfs, escape = FALSE, booktabs = T,
     digits = 3) %>%
  kable_styling(bootstrap_options = c("striped", "hover", "condensed", "responsive")) %>%
kableExtra::landscape() %>%
  #kable_styling(full_width = F) %>%
    kable_paper(full_width = F)
```


### 17. BRCA data exploration
```{r}

wbm_an_stats <- source_wbm %>% 
  dplyr::filter(year(date) < 2024) %>%
  dplyr::mutate(y = year(date)) %>%
  dplyr::ungroup() %>%
  dplyr::select(-c(date, CF, GCM, name)) %>%
  dplyr::group_by(y) %>%
  dplyr::summarize_all(., .funs = sd)

# GW_level variability as sensitivity indicator
well_data_stats <- well_data %>%
  dplyr::mutate(y = year(date)) %>%
  dplyr::group_by(well,y) %>% 
  dplyr::summarize(mean = mean(static_ft, na.rm = TRUE),
                   st_dev = sd(static_ft, na.rm = TRUE)) %>%
  left_join(., wbm_an_stats, by = "y", relationship = "many-to-many")

ggplot(well_data_stats, aes(x = runoff_in_wbm, y = st_dev, color = well)) +
  geom_point() +
  geom_smooth()

# Runoff ratio (runoff / precip) which
ggplot(source_wbm_annual, aes(x = year, y = runoff_in_wbm/precip_in_wbm)) +
  geom_point() +
  theme_bw() +
  labs(y = "Runoff ratio (runoff/precip)")



# Melt to runoff ratio through time
ggplot(source_wbm_annual %>%
         dplyr::filter(year < 2075), aes(x = year, 
                                         y = melt_in_wbm/runoff_in_wbm, color = GCM)) +
  geom_point() +
  geom_smooth(method = "lm") +
  theme_bw() +
  labs(y = "melt/runoff")


# Runoff Elasticity
ggplot(source_wbm_annual, aes(x = rain_in_wbm, y = runoff_in_wbm)) +
  geom_point() +
  geom_smooth(method = "lm", formula = y ~ x) +
  stat_cor(label.y = max(source_wbm_annual$runoff_in_wbm) * .9) +
  stat_regline_equation(label.y = max(source_wbm_annual$runoff_in_wbm) * .98) +
  theme_bw() +
  labs(title = "runoff elasticity (runoff/precip)")




```



Relate gw levels to wbm
```{r}
 
z_score <- function(x) {
  # if NA values, fill. If fist
  # val is NA, replace with first 
  # non-NA value.
  if (is.na(x[1])) {
    x[1] <- first(x[!is.na(x)])
  }
  x <- zoo::na.approx(x, na.rm = FALSE)
  mean_x <- mean(x, na.rm = TRUE)
  stdev_x <- sd(x, na.rm = TRUE)
  dep <- (x - mean_x)/stdev_x
}


# Calculate annual z-score
wbm_nwis_z <- source_wbm %>%
  ungroup() %>%
  left_join(., well_data %>% dplyr::filter(well == "Well 1"), by = "date") %>%
  dplyr::mutate(y = year(date),
                m = month(date),
                ym = as.Date(paste0(y,'-',m,'-15'))) %>%
  dplyr::filter(y > 2001,
                y < 2024) %>%
  dplyr::group_by(y) %>%
  dplyr::summarize(precip = sum(precip_in_wbm, na.rm = TRUE),
                   temp = mean(tmean_C, na.rm = TRUE),
                   gw = mean(static_ft_c),
                   melt = sum(melt_in_wbm, na.rm = TRUE),
                   rain = sum(rain_in_wbm, na.rm = TRUE),
                   runoff = sum(runoff_in_wbm, na.rm = TRUE)) %>%
  ungroup() %>%
  dplyr::mutate(melt_z = z_score(melt),
                precip_z = z_score(precip),
                temp_z = z_score(temp),
                gw_z = z_score(gw),
                melt_z = z_score(melt),
                rain_z = z_score(rain),
                runoff_z = z_score(runoff)) 

ggplot(wbm_nwis_z) + 
  geom_line(aes(x = y, y = gw_z, color = "GW"), lwd = 1) +
  geom_line(aes(x = y, y = melt_z, color = "Snowmelt"), lwd = 1) +
  geom_line(aes(x = y, y = rain_z, color = "Rain"), lwd = 1) +
  geom_line(aes(x = y, y = runoff_z, color = "Runoff"), lwd = 1) +
  #geom_line(aes(x = y, y = )) +
  theme_bw() +
  scale_color_manual("", values = c("black","dodgerblue", "tomato", "pink2")) 



ggplot(wbm_nwis_z) +
  #geom_line(aes(x = y, y = precip_z, color = "Precip"), lwd = 1) + 
  geom_line(aes(x = y, y = gw_z, color = "GW Depth"), lwd = 2) + 
  #eom_line(aes(x = y, y = temp_dep, color = "Temp C"), lwd = 1) +
  #geom_line(aes(x = y, y = melt_z, color = "Snowmwelt"), lwd = 1) +
  #geom_line(aes(x = y, y = rain_dep, color = "Rain"), lwd = 1) +
  #geom_line(aes(x = y, y = lag(runoff_z,0), color = "Runoff"), lwd = 1) +
  geom_line(aes(x = y, y = lag(runoff_z,1), color = "Runoff_lag1year"), lwd = 1, alpha = 0.5) +
  #geom_line(aes(x = y, y = lag(frollmean(runoff_z,2,align = "right"),1), color = "Runoff_rolling2"), lwd = 1, alpha = 0.5) +
  theme_bw() + 
  scale_color_manual("", values = c("black","dodgerblue","seagreen","tomato","pink2","orange")) +
  labs(y = "Z-Score")

#ggplotly(
ggplot(wbm_nwis_z, aes(x = lag(runoff_z,1), y = gw)) + 
  geom_point() + 
  geom_smooth(method = "lm", formula = y ~ x, color = "black") +
  #stat_cor(label.y = max(wbm_nwis_z$gw_z, na.rm = TRUE) * 1.05) +
  #stat_regline_equation(label.y = max(wbm_nwis_z$gw_z, na.rm = TRUE) * 1.2) +
   stat_regline_equation(
    aes(label =  paste(..eq.label.., ..adj.rr.label.., sep = "~~~~~~~~~")),
    formula = y ~x) +
  theme_bw() + 
  labs(x = "Runoff (lag_1)", y = "Minimum Annual GW Depth")#)

#ggarrange(a,b,ncol=1, align = "v")




```

SPEI

```{r}

# Calculate SPEI using spei packages

ag_n <- 1 # number of months to aggregate for time series, 1 = monthly
freq <- 12/ag_n # frequency 
agg_months <- ifelse(ag_n == 1,paste0(ag_n, " month"), paste0(ag_n, " months"))

# Aggregate into ag_n dataframe
spei_wbm <- source_wbm %>% 
  dplyr::ungroup() %>%
  dplyr::select(-name) %>%
  left_join(well_data %>% 
              dplyr::filter(well == "Well 1") %>%
              dplyr::ungroup() %>%
              dplyr::select(c("date","static_ft_c","static_ft_o" = "static_ft")), 
            by = "date") %>%
  dplyr::mutate(year = year(date),
                month = month(date),
                time_ag = floor_date(date, agg_months)) %>%
  dplyr::group_by(year,time_ag, CF) %>%
  dplyr::summarize(precip_in_wbm = sum(precip_in_wbm),
                   tmean_C = mean(tmean_C),
                   pet_in_wbm = sum(pet_in_wbm),
                   static_ft = min(static_ft_c),
                   runoff_in_wbm = sum(runoff_in_wbm),
                   static_ft_o = min(static_ft_o, na.rm = TRUE)) %>%
  dplyr::mutate(balance = precip_in_wbm - pet_in_wbm,
                month = month(time_ag))

# Grab last year
end_year = spei_wbm$year %>% tail(., n = 1)
end_month = spei_wbm$month %>% tail(., n = 1)

# define timesries for historical, hot dry, and warm wet periods. For 
# continuity in SPEI calc, hot dry and warm wet should also include the
# historical. 
#spei_ts <- ts(spei_wbm, end = c(end_year, end_month), frequency = 12)
spei_historical_ts <- ts(spei_wbm %>% 
                           dplyr::filter(CF == "Historical"),
                         end = c(2023,12), frequency = freq)

spei_hot_dry_ts <- ts(spei_wbm %>% 
                        dplyr::filter(CF %in% c("Historical","Hot Dry")),
                      end = c(end_year,end_month), frequency = freq)

spei_warm_wet_ts <- ts(spei_wbm %>% 
                         dplyr::filter(CF %in% c("Historical","Warm Wet")),
                      end = c(end_year,end_month), frequency = freq)

spei_historical <- 
  spei_wbm %>%
  dplyr::ungroup() %>%
      dplyr::filter(CF %in% c("Historical")) %>%
      dplyr::mutate(
        spei_historical = c(spei(spei_historical_ts[,'balance'], freq,
                                   ref.end = c(2023,1))[["fitted"]]))

spei_hot_dry <- 
  spei_wbm %>%
  dplyr::ungroup() %>%
      dplyr::filter(CF %in% c("Historical","Hot Dry")) %>%
      dplyr::mutate(
        spei_hot_dry = c(spei(spei_hot_dry_ts[,'balance'], freq,
                                   ref.end = c(2023,1))[["fitted"]]))

spei_warm_wet <- 
  spei_wbm %>%
  dplyr::ungroup() %>%
      dplyr::filter(CF %in% c("Historical","Warm Wet")) %>%
      dplyr::mutate(
        spei_warm_wet = c(spei(spei_warm_wet_ts[,'balance'], freq,
                                   ref.end = c(2023,1))[["fitted"]]))


## Plot all SPEI

ggplot() +
  geom_line(data = spei_hot_dry, aes(x = time_ag, y = spei_hot_dry, color = "Hot Dry"), lwd = 1) +
   geom_line(data = spei_warm_wet %>% dplyr::filter(year(time_ag) > 2024), aes(x = time_ag, y = spei_warm_wet, color = "Warm Wet"), lwd = 1) +
  geom_line(data = spei_historical, aes(x = time_ag, y = spei_historical, color = "Historical"), lwd = 1) +
  scale_color_manual("", values = c("black","tomato","dodgerblue")) + 
  theme_bw() +
  labs(x = "", y = "SPEI")


factor <- 2.5
shift <- 0
lag_num <- 4
roll_num <- 12

lag_roll <- function(dat,lag_num, roll_num) {
  # lag
  lag_vals <- rep(NA,abs(lag_num))
  if (lag_num > 0) {
    lag_dat <- c(lag_vals,dat[1:(length(dat)-lag_num)])
  } else if (lag_num < 0) {
    lag_dat <- c(dat[(abs(lag_num)+1):length(dat)], lag_vals) 
  } else {
    lag_dat <- dat
  }
  
  # roll
  roll_dat <- frollmean(lag_dat,roll_num,align=c("center")) 
  
}


a <- ggplot(spei_historical %>% dplyr::filter(year > 2004, year < 2024)) +
  geom_line(aes(x = time_ag, 
                y = static_ft, color = time_ag), alpha = 0.5) +
  geom_line(aes(x = time_ag, 
                y=lag_roll(static_ft,0,roll_num),
                color = time_ag), 
            lwd = 1) +
  geom_line(aes(x = time_ag, 
                y = factor*(spei_historical+shift), 
                alpha = 0.5), color = "dodgerblue") +
  geom_line(aes(x = time_ag, 
                y = factor*(lag_roll(spei_historical,0,roll_num) + shift)), color = "dodgerblue",lwd = 1) +

   annotate("rect",xmin=as.Date("2013-07-01"), 
            xmax=as.Date("2016-07-01"), 
            ymin=-Inf, 
            ymax=Inf, alpha = 0.1) +
  annotate("text",
           y = 4.5,
            x=as.Date("2015-01-01"), 
            label = "Deviation \nPeriod") +
  scale_y_continuous(sec.axis = sec_axis(~(.-shift)/factor , name = "SPEI (in)")) +
 theme_bw() +
  #scale_color_manual("",values=c("black", "dodgerblue","red")) +
  scale_color_viridis(discrete = FALSE) +
  #scale_color_gradient(low = pal[5], high = pal[13]) +
  #scale_fill_manual("",values=c("black", "dodgerblue","red")) +
  labs(y = "GW Depth", x = "") +
  theme(
    axis.title.y = element_text(color = "black"),
    axis.text.y = element_text(color = "black"),
    axis.title.y.right = element_text(color = "dodgerblue"),
    axis.text.y.right = element_text(color = "dodgerblue"),
    legend.position = "none"
  )

b <- ggplot(spei_historical %>%
         dplyr::filter(year > 2004, year < 2022),
       aes(x = lag_roll(spei_historical,lag_num,roll_num), 
           y = lag_roll(static_ft,0,roll_num),
           color = time_ag)) +
  geom_point() +
  scale_color_viridis(discrete = FALSE) +
  stat_smooth(method = "lm", formula = y ~ x, color = "black") +
 stat_cor(label.y = max(spei_historical$static_ft, na.rm = TRUE) * 1.05) +
 stat_regline_equation(
    aes(label =  paste(..eq.label.., ..adj.rr.label.., sep = "~~~~")),
   formula = y ~x) +
    labs(x = "SPEI", y = "Minimum Monthly GW Depth") +
  theme_bw() +
  theme(legend.position = "none")

ggarrange(a,b, ncol = 1, align = "v")


# Extract linear model for projections
reg_mod <- lm(lag_roll(static_ft, 0,roll_num) ~ 
              lag_roll(spei_historical,lag_num, roll_num),
  data = spei_historical %>%
         dplyr::filter(year > 2004, year < 2024))
reg_sum <- summary(reg_mod)

intercept <- as.numeric(reg_sum$coefficients[1])
slope <- as.numeric(reg_sum$coefficients[2])

spei_historical2 <- spei_historical %>%
  dplyr::filter(year > 2000) %>%
  dplyr::ungroup() %>%
  dplyr::mutate(y_predict = intercept + slope*(lag_roll(spei_historical,lag_num,roll_num)),
                static_ft_predict = lag_roll(y_predict,-lag_num,1),
                residuals = static_ft_predict - static_ft)

# Calculate estimated variance
sigma_hat_squared <- var(spei_historical2$residuals, na.rm = TRUE)
# Calculate standard errors
n <- length(spei_historical2$residuals)
x_bar <- mean(spei_historical2$static_ft, na.rm = TRUE)  # Assuming 'x' is the predictor variable

spei_hot_dry2 <- spei_hot_dry %>%
  dplyr::filter(year > 2000) %>%
  dplyr::ungroup() %>%
  dplyr::mutate(y_predict = intercept + slope*(lag_roll(spei_hot_dry,lag_num,roll_num)),
                static_ft_predict = lag_roll(y_predict,-lag_num,1),
                standard_error = 
  sqrt(sigma_hat_squared * (1/n + (static_ft_predict- x_bar)^2 / sum((spei_historical2$static_ft_predict - x_bar)^2, na.rm = TRUE))))

spei_warm_wet2 <- 
 spei_warm_wet %>%
  dplyr::filter(year > 2000) %>%
  dplyr::ungroup() %>%
  dplyr::mutate(y_predict = intercept + slope*(lag_roll(spei_warm_wet,lag_num,roll_num)),
                static_ft_predict = lag_roll(y_predict,-lag_num,1),
                 standard_error = 
  sqrt(sigma_hat_squared * (1/n + (static_ft_predict- x_bar)^2 / sum((spei_historical2$static_ft_predict - x_bar)^2, na.rm = TRUE))))


ggplot() +
  geom_line(data = spei_historical2 %>% dplyr::filter(year > 2000),
            aes(x = time_ag, y = static_ft, color = "Historical (observed)"), lwd = 1) +
  geom_line(data = spei_historical2 %>% dplyr::filter(year > 2000),
            aes(x = time_ag, y = static_ft_predict, color = "Historical (predicted)"), 
            lwd = 1, alpha = .3) +
  geom_line(data = spei_warm_wet2 %>% dplyr::filter(year > 2021),
            aes(x = time_ag, y = static_ft_predict, color = "Warm Wet"), lwd = 1) +
  geom_ribbon(data = spei_warm_wet2 %>% dplyr::filter(year > 2021),
            aes(x = time_ag, 
                ymin = static_ft_predict - standard_error, 
                ymax = static_ft_predict + standard_error), fill = "dodgerblue", 
            lwd = 0, alpha = 0.6) +
  geom_line(data = spei_hot_dry2 %>% dplyr::filter(year > 2021),
            aes(x = time_ag, y = static_ft_predict, color = "Hot Dry"), lwd = 1) +
   geom_ribbon(data = spei_hot_dry2 %>% dplyr::filter(year > 2021),
            aes(x = time_ag, 
                ymin = static_ft_predict - standard_error, 
                ymax = static_ft_predict + standard_error), fill = "tomato", lwd = 0, alpha = 0.6) +
   geom_ribbon(data = spei_hot_dry2 %>% dplyr::filter(year > 2000) %>%
               dplyr::mutate(
                 pump_depth = 
                    ifelse(spei_hot_dry < quantile(spei_hot_dry, 0.3),
                           static_ft_predict-5.5,static_ft_predict)),
            aes(x = time_ag, 
                ymax = static_ft_predict, 
                ymin = pump_depth, 
                fill = "Potential Pumping\n Effects"), lwd = 1, alpha = 0.2) +
  theme_bw() +
  scale_color_manual("",values = c("black","gray20","tomato","dodgerblue")) +
  scale_fill_manual("", values = c("tomato","dodgerblue")) +
  labs(x = "", y = "GW Depth (ft)")
  


# plot(spei(ts(wichita[, "BAL"], freq = 12, start = c(1980, 6)), 12,
# ref.start = c(1980, 1), ref.end = c(2000, 1)
# ))

```