3b_Streamflow_Output_Visualization.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#%run \"0a_Workspace_setup.ipynb\"\n",
    "%run \"0b_Create_poi_files.ipynb\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib as mplib\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.dates as mdates\n",
    "\n",
    "import folium\n",
    "from folium import Choropleth, Circle, Marker\n",
    "\n",
    "from folium import plugins\n",
    "from shapely.geometry import Polygon\n",
    "from folium.features import DivIcon\n",
    "from folium.plugins import MarkerCluster\n",
    "from shapely.geometry import Polygon\n",
    "\n",
    "from folium.plugins import MarkerCluster\n",
    "\n",
    "\n",
    "from tobler.dasymetric import masked_area_interpolate\n",
    "from tobler.model import glm\n",
    "from tobler.area_weighted import area_interpolate\n",
    "\n",
    "from libpysal.examples import load_example\n",
    "\n",
    "import hydroeval as he\n",
    "\n",
    "# import hyswap\n",
    "# from hyswap.percentiles import calculate_variable_percentile_thresholds_by_day\n",
    "# from hyswap.cumulative import calculate_daily_cumulative_values\n",
    "import calendar\n",
    "import statistics\n",
    "from sklearn.metrics import r2_score"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Useful grouping functions for HRUs grouping--adopted from pyPRMS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "def subset_stream_network(dag_ds, uscutoff_seg, dsmost_seg):\n",
    "    \"\"\"Extract subset of stream network\n",
    "\n",
    "    :param dag_ds: Directed, acyclic graph of downstream stream network\n",
    "    :param uscutoff_seg: List of upstream cutoff segments\n",
    "    :param dsmost_seg: List of outlet segments to start extraction from\n",
    "\n",
    "    :returns: Stream network of extracted segments\n",
    "    \"\"\"\n",
    "\n",
    "    # taken from Bandit bandit_helpers.py\n",
    "\n",
    "    # Create the upstream graph\n",
    "    dag_us = dag_ds.reverse()\n",
    "\n",
    "    # Trim the u/s graph to remove segments above the u/s cutoff segments\n",
    "    try:\n",
    "        for xx in uscutoff_seg:\n",
    "            try:\n",
    "                dag_us.remove_nodes_from(nx.dfs_predecessors(dag_us, xx))\n",
    "\n",
    "                # Also remove the cutoff segment itself\n",
    "                dag_us.remove_node(xx)\n",
    "            except KeyError:\n",
    "                print(f\"WARNING: nhm_segment {xx} does not exist in stream network\")\n",
    "    except TypeError:\n",
    "        print(\n",
    "            \"\\nSelected cutoffs should at least be an empty list instead of NoneType.\"\n",
    "        )\n",
    "\n",
    "    # =======================================\n",
    "    # Given a d/s segment (dsmost_seg) create a subset of u/s segments\n",
    "\n",
    "    # Get all unique segments u/s of the starting segment\n",
    "    uniq_seg_us: Set[int] = set()\n",
    "    if dsmost_seg:\n",
    "        for xx in dsmost_seg:\n",
    "            try:\n",
    "                pred = nx.dfs_predecessors(dag_us, xx)\n",
    "                uniq_seg_us = uniq_seg_us.union(\n",
    "                    set(pred.keys()).union(set(pred.values()))\n",
    "                )\n",
    "            except KeyError:\n",
    "                print(f\"KeyError: Segment {xx} does not exist in stream network\")\n",
    "\n",
    "        # Get a subgraph in the dag_ds graph and return the edges\n",
    "        dag_ds_subset = dag_ds.subgraph(uniq_seg_us).copy()\n",
    "\n",
    "        node_outlets = [ee[0] for ee in dag_ds_subset.edges()]\n",
    "        true_outlets = set(dsmost_seg).difference(set(node_outlets))\n",
    "\n",
    "        # Add the downstream segments that exit the subgraph\n",
    "        for xx in true_outlets:\n",
    "            nhm_outlet = list(dag_ds.neighbors(xx))[0]\n",
    "            dag_ds_subset.add_node(\n",
    "                nhm_outlet, style=\"filled\", fontcolor=\"white\", fillcolor=\"grey\"\n",
    "            )\n",
    "            dag_ds_subset.add_edge(xx, nhm_outlet)\n",
    "            dag_ds_subset.nodes[xx][\"style\"] = \"filled\"\n",
    "            dag_ds_subset.nodes[xx][\"fontcolor\"] = \"white\"\n",
    "            dag_ds_subset.nodes[xx][\"fillcolor\"] = \"blue\"\n",
    "    else:\n",
    "        # No outlets specified so pull the full model\n",
    "        dag_ds_subset = dag_ds\n",
    "\n",
    "    return dag_ds_subset\n",
    "\n",
    "\n",
    "def hrus_by_seg(pdb, segs):\n",
    "    # segs: global segment IDs\n",
    "\n",
    "    if isinstance(segs, int):\n",
    "        segs = [segs]\n",
    "    elif isinstance(segs, KeysView):\n",
    "        segs = list(segs)\n",
    "\n",
    "    seg_hrus = {}\n",
    "    seg_to_hru = pdb.seg_to_hru\n",
    "\n",
    "    # Generate stream network for the model\n",
    "    dag_streamnet = pdb.stream_network()\n",
    "\n",
    "    for cseg in segs:\n",
    "        # Lookup segment for the current POI\n",
    "        dsmost_seg = [cseg]\n",
    "\n",
    "        # Get subset of stream network for given POI\n",
    "        dag_ds_subset = subset_stream_network(dag_streamnet, set(), dsmost_seg)\n",
    "\n",
    "        # Create list of segments in the subset\n",
    "        toseg_idx = list(set(xx[0] for xx in dag_ds_subset.edges))\n",
    "\n",
    "        # Build list of HRUs that contribute to the POI\n",
    "        final_hru_list = []\n",
    "\n",
    "        for xx in toseg_idx:\n",
    "            try:\n",
    "                for yy in seg_to_hru[xx]:\n",
    "                    final_hru_list.append(yy)\n",
    "            except KeyError:\n",
    "                # print(f'Segment {xx} has no HRUs connected to it') # comment this out and add pass to not print the KeyError\n",
    "                pass\n",
    "        final_hru_list.sort()\n",
    "\n",
    "        seg_hrus[cseg] = final_hru_list\n",
    "\n",
    "    return seg_hrus\n",
    "\n",
    "\n",
    "def hrus_by_poi(pdb, poi):\n",
    "    if isinstance(poi, str):\n",
    "        poi = [poi]\n",
    "    elif isinstance(poi, KeysView):\n",
    "        poi = list(poi)\n",
    "\n",
    "    poi_hrus = {}\n",
    "    nhm_seg = pdb.get(\"nhm_seg\").data\n",
    "    pois_dict = pdb.poi_to_seg\n",
    "    seg_to_hru = pdb.seg_to_hru\n",
    "\n",
    "    # Generate stream network for the model\n",
    "    dag_streamnet = pdb.stream_network()\n",
    "\n",
    "    for cpoi in poi:\n",
    "        # Lookup global segment id for the current POI\n",
    "        dsmost_seg = [nhm_seg[pois_dict[cpoi] - 1]]\n",
    "\n",
    "        # Get subset of stream network for given POI\n",
    "        dag_ds_subset = subset_stream_network(dag_streamnet, set(), dsmost_seg)\n",
    "\n",
    "        # Create list of segments in the subset\n",
    "        toseg_idx = list(set(xx[0] for xx in dag_ds_subset.edges))\n",
    "\n",
    "        # Build list of HRUs that contribute to the POI\n",
    "        final_hru_list = []\n",
    "\n",
    "        for xx in toseg_idx:\n",
    "            try:\n",
    "                for yy in seg_to_hru[xx]:\n",
    "                    final_hru_list.append(yy)\n",
    "            except KeyError:\n",
    "                # Not all segments have HRUs connected to them\n",
    "                # print(f'{cpoi}: Segment {xx} has no HRUs connected to it')\n",
    "                pass\n",
    "        final_hru_list.sort()\n",
    "        poi_hrus[cpoi] = final_hru_list\n",
    "\n",
    "    return poi_hrus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def stats_table(stats_df):\n",
    "\n",
    "    evaluations = stats_df.discharge\n",
    "    std_evaluations = statistics.stdev(evaluations)\n",
    "\n",
    "    simulations = stats_df.seg_outflow\n",
    "\n",
    "    rmse = np.round(he.evaluator(he.rmse, simulations, evaluations), 2)\n",
    "    nse = np.round(he.evaluator(he.nse, simulations, evaluations), 2)\n",
    "    pbias = np.round(he.evaluator(he.pbias, simulations, evaluations), 2)\n",
    "    kge, r, alpha, beta = np.round(he.evaluator(he.kge, simulations, evaluations), 2)\n",
    "\n",
    "    rsr = np.round(rmse / std_evaluations, 2)\n",
    "    r_sq = np.round(np.array([r2_score(simulations, evaluations)]), 2)\n",
    "\n",
    "    stat_dict = {\n",
    "        \"KGE\": kge[0],\n",
    "        \"NSE\": nse[0],\n",
    "        \"Pbias\": pbias[0],\n",
    "        \"RMSE\": rmse[0],\n",
    "        \"R^2\": r_sq[0],\n",
    "        \"R\": r[0],\n",
    "        \"Alpha\": alpha[0],\n",
    "        \"Beta\": beta[0],\n",
    "        \"RSR\": rsr[0],\n",
    "    }\n",
    "\n",
    "    df = pd.DataFrame(stat_dict, index=[0])\n",
    "\n",
    "    return df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### NHM Calibration Levels for HRUs: (those hrus calibrated in byHW and byHWobs parts)\n",
    "HW basins were descritized using a drainage area maxiumum and minimum; HW HRUs, segments, outlet segment, and drainage area are available. Gages used in byHWobs calibration, Part 3, for selected headwaters are also provided here.  FILES AND TABLES IN THIS SECTION ARE CONUS COVERAGE and will be subsetted later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#### READ table (.csv) of HRU calibration level file\n",
    "hru_cal_levels_df = pd.read_csv(\n",
    "    r\"data_dependencies/NHM_v1_1/nhm_v1_1_HRU_cal_levels.csv\"\n",
    ").fillna(0)\n",
    "hru_cal_levels_df[\"hw_id\"] = hru_cal_levels_df.hw_id.astype(\"int64\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hru_cal_levels_df = pd.merge(\n",
    "    hru_cal_levels_df, hru_gdf, right_on=\"nhm_id\", left_on=\"nhm_id\"\n",
    ")\n",
    "hru_cal_levels_gdf = gpd.GeoDataFrame(\n",
    "    hru_cal_levels_df, geometry=\"geometry\"\n",
    ")  # Creates a Geopandas GeoDataFrame\n",
    "hru_cal_levels_gdf[\"nhm_id\"] = hru_cal_levels_gdf[\"nhm_id\"].astype(str)\n",
    "hru_cal_levels_gdf[\"hw_id\"] = hru_cal_levels_gdf[\"hw_id\"].astype(str)\n",
    "\n",
    "print(\n",
    "    \"The number of HRUs in the byHRU calibration is\",\n",
    "    hru_cal_levels_gdf[hru_cal_levels_gdf[\"level\"] > 0][\"level\"].count(),\n",
    ")\n",
    "print(\n",
    "    \"The number of HRUs in the byHW calibration is\",\n",
    "    hru_cal_levels_gdf[hru_cal_levels_gdf[\"level\"] > 1][\"level\"].count(),\n",
    ")\n",
    "print(\n",
    "    \"The number of HRUs in the byHWobs calibration is\",\n",
    "    hru_cal_levels_gdf[hru_cal_levels_gdf[\"level\"] > 2][\"level\"].count(),\n",
    ")\n",
    "\n",
    "\n",
    "# hru_cal_levels_df #View results to verify"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Add headwater basin (NHM calibration basin) outline layer on the map for referrence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "byHW_basins_gdf = hru_cal_levels_gdf.loc[hru_cal_levels_gdf[\"byHW\"] == 1]\n",
    "HW_basins_gdf = byHW_basins_gdf.dissolve(by=\"hw_id\").to_crs(crs)\n",
    "HW_basins_gdf.reset_index(inplace=True, drop=False)\n",
    "HW_basins = HW_basins_gdf.boundary\n",
    "# HW_basins_gdf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Identify the NHM poi gages that were used in claibration byHWobs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This reads in the csv file that hase the gages used to calibrate the byHWobs part for CONUS.\n",
    "# Read in station file columns needed (You may need to tailor this to the particular file.\n",
    "col_names = [\n",
    "    \"poi_id\",\n",
    "    #'poi_name',\n",
    "    \"latitude\",\n",
    "    \"longitude\",\n",
    "    #'drainage_area',\n",
    "    #'drainage_area_contrib'\n",
    "]\n",
    "col_types = [\n",
    "    np.str_,\n",
    "    # np.str_,\n",
    "    float,\n",
    "    float,\n",
    "    # float,\n",
    "    # float\n",
    "]\n",
    "cols = dict(\n",
    "    zip(col_names, col_types)\n",
    ")  # Creates a dictionary of column header and datatype called below.\n",
    "\n",
    "byHWobs_poi_df = pd.read_csv(\n",
    "    r\"data_dependencies/NHM_v1_1/nhm_v1_1_byhwobs_cal_gages.csv\", sep=\"\\t\", dtype=cols\n",
    ").fillna(0)\n",
    "\n",
    "# byHWobs_poi_df = pd.read_csv(f'{NHM_dir}/nhm_v11_hwobs_pois.csv', sep='\\t').fillna(0)\n",
    "# byHWobs_poi_df['poi_id'] = byHWobs_poi_df.poi_id.astype('str') # makes sure that this is a string,\n",
    "# must have the leading zeros; suggest a more formal read and set like used in prev notebook.\n",
    "\n",
    "# Identify the byHWobs calibration gages in our current poi database (ammended in the model prams file to include more gages)\n",
    "poi_df[\"nhm_calib\"] = \"N\"\n",
    "poi_df.loc[poi_df[\"poi_id\"].isin(byHWobs_poi_df[\"poi_id\"]), \"nhm_calib\"] = \"Y\"\n",
    "# poi_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "poi_df.poi_id"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Compute KGE for all gages to color the icon on the map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set WY start and stop times, and output variable needed for slicing the time series data for plotting\n",
    "WY_start = \"1979-10-01\"\n",
    "WY_end = \"2021-09-30\"\n",
    "# Note that the model start and stop times in the control file should be the same as the observation start and stop times.\n",
    "output_var_sel = \"seg_outflow\"  # Select this output variable to read"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read in simulated flows and write daily ts array and for resample: monthly and annual\n",
    "with xr.open_dataset(out_dir / \"model_custom_output.nc\") as model_output:\n",
    "    # if the start time and end time is in the netcdf, then passs, else print an error.\n",
    "    output = model_output.assign_coords(npoi_gages=model_output.poi_gages)\n",
    "    output_var = getattr(output, output_var_sel)\n",
    "    output_var = output_var.sel(\n",
    "        time=slice(WY_start, WY_end)\n",
    "    )  # .resample(time = 'm').mean()\n",
    "    output_var_monthly = output_var.resample(\n",
    "        time=\"m\"\n",
    "    ).mean()  # resample data to monthly values\n",
    "    output_var_annual = output_var.resample(\n",
    "        time=\"A-SEP\"\n",
    "    ).mean()  # resample data to monthly values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_var"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read in observed flows\n",
    "# Note that the model start and stop times in the control file should be the same as the observation start and stop times.\n",
    "sf_filename = model_dir / 'notebook_output_files' / 'nc_files'  / 'sf_efc.nc'\n",
    "    \n",
    "with xr.open_dataset(sf_filename) as obs_data:\n",
    "    # Make a station name dataframe and station id list from the streamflow file .nc (created in previous notebook)\n",
    "    station_name_df = getattr(obs_data, 'poi_name').to_dataframe()# supporting df for plot labeling\n",
    "    station_id_list = station_name_df.index.to_list()# supporting list for processing\n",
    "\n",
    "    # Resample daily timeseries arrays: monthly and annual \n",
    "    obs_0 = (obs_data.sel(time=slice(WY_start, WY_end))).transpose() #load_dataset will open, read into memory and close the .nc file\n",
    "    obs_efc = getattr(obs_0, 'efc')\n",
    "    obs = getattr(obs_0, 'discharge')\n",
    "    obs_monthly = obs.resample(time = 'm').mean()\n",
    "    obs_annual = obs.resample(time = 'A-SEP').mean()  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_var_monthly\n",
    "#obs_monthly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for idx, row in poi_df.iterrows():\n",
    "#     zzz= output_var_monthly.sel(npoi_gages = row['poi_id']).values.tolist()\n",
    "#     kkk= obs_monthly.sel(poi_id = row['poi_id']).values.tolist()\n",
    "#     #print(len(zzz), len(kkk))\n",
    "#     print(zzz[0], kkk[0])\n",
    "\n",
    "test=poi_df.poi_id[0]\n",
    "zzz= output_var_monthly.sel(npoi_gages = test).values.tolist()\n",
    "kkk= obs_monthly.sel(poi_id = test).values.tolist()\n",
    "#print(len(zzz), len(kkk))\n",
    "#print(zzz, kkk)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "poi_df['kge']= np.nan\n",
    "for idx, row in poi_df.iterrows():\n",
    "    poi_id_sel = row['poi_id']\n",
    "    df_sf_data_sel = (obs.sel(poi_id = poi_id_sel)).to_dataframe()\n",
    "    \n",
    "    # Determine por\n",
    "    por_start = df_sf_data_sel['discharge'].notna().idxmax()# First Day\n",
    "    por_end = df_sf_data_sel['discharge'].notna()[::-1].idxmax()# Last Day\n",
    "        \n",
    "    # Slice to por\n",
    "    df_sf_data_sel = (obs.sel(poi_id = poi_id_sel, time=slice(por_start, por_end))).to_dataframe()\n",
    "    df_sf_data_sel.drop(columns =['poi_id'], inplace =True)# drop unwanted columns\n",
    "        \n",
    "    sim_flow = (output_var.sel(npoi_gages = poi_id_sel, time=slice(por_start, por_end))).to_dataframe()\n",
    "    sim_flow.drop(columns =['npoi_gages', 'poi_gages'], inplace =True)# drop unwanted columns\n",
    "        \n",
    "    \n",
    "    # drop the Nan's from the obs for memory/stats (may want to check back on this later)\n",
    "    daily_stat_df = (df_sf_data_sel.merge(sim_flow, right_index = True, left_index = True, how = 'inner')).dropna()\n",
    "    month_stat_df = daily_stat_df.resample('m').mean().dropna()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    " daily_stat_df.info()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Calculate kge and add value to poi_df\n",
    "poi_df['kge']= np.nan\n",
    "for idx, row in poi_df.iterrows():\n",
    "    poi_id_sel = row['poi_id']\n",
    "    df_sf_data_sel = (obs.sel(poi_id = poi_id_sel)).to_dataframe()\n",
    "    \n",
    "    # Determine por\n",
    "    por_start = df_sf_data_sel['discharge'].notna().idxmax()# First Day\n",
    "    por_end = df_sf_data_sel['discharge'].notna()[::-1].idxmax()# Last Day\n",
    "        \n",
    "    # Slice to por\n",
    "    df_sf_data_sel = (obs.sel(poi_id = poi_id_sel, time=slice(por_start, por_end))).to_dataframe()\n",
    "    df_sf_data_sel.drop(columns =['poi_id'], inplace =True)# drop unwanted columns\n",
    "        \n",
    "    sim_flow = (output_var.sel(npoi_gages = poi_id_sel, time=slice(por_start, por_end))).to_dataframe()\n",
    "    sim_flow.drop(columns =['npoi_gages', 'poi_gages'], inplace =True)# drop unwanted columns\n",
    "        \n",
    "    \n",
    "    # drop the Nan's from the obs for memory/stats (may want to check back on this later)\n",
    "    daily_stat_df = (df_sf_data_sel.merge(sim_flow, right_index = True, left_index = True, how = 'inner')).dropna()\n",
    "    month_stat_df = daily_stat_df.resample('m').mean().dropna()\n",
    "        \n",
    "    # kge_func  = np.round(he.evaluator(he.kge,\n",
    "    #                                   daily_stat_df['seg_outflow'],# simulation data set\n",
    "    #                                   daily_stat_df['discharge'],# observation data set\n",
    "    #                                  ), 2# decimal places for the round() function\n",
    "    #                     )[0]#this grabs only the kge var, in position\"0\" from the list of ke.kge() output vars\n",
    "\n",
    "    # poi_df.loc[idx, 'kge'] = np.array(kge_func[0])# pandas wrangling of the array output from he.evaluator() as an array\n",
    "\n",
    "    kge_func  = np.round(he.evaluator(he.kge,\n",
    "                                      month_stat_df['seg_outflow'],# simulation data set\n",
    "                                      month_stat_df['discharge'],# observation data set\n",
    "                                     ), 2# decimal places for the round() function\n",
    "                        )[0]#this grabs only the kge var, in position\"0\" from the list of ke.kge() output vars\n",
    "\n",
    "    poi_df.loc[idx, 'kge'] = np.array(kge_func[0])# pandas wrangling of the array output from he.evaluator() as an array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "poi_df.kge"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create an interactive map to evaluate streamflow at poi_gages"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Set Map template information"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "poi_id_sel = None\n",
    "# Set map CRS, center, zoom, basemaps, and inset\n",
    "\n",
    "crs = 4326\n",
    "\n",
    "# pfile lat, lon derived for starting point of folium plot windows. Zoom also set here.\n",
    "hru_gdf_map = hru_gdf.to_crs(crs)\n",
    "lat = hru_gdf_map['hru_lat'].mean()\n",
    "lon =  hru_gdf_map['hru_lon'].mean()-1\n",
    "zoom = 7\n",
    "\n",
    "# Set base map options\n",
    "USGStopo_layer = folium.TileLayer(tiles = 'https://basemap.nationalmap.gov/arcgis/rest/services/USGSTopo/MapServer/tile/{z}/{y}/{x}',\n",
    "                 attr = 'USGS_topo', zoom_start = zoom, name = \"USGSTopo\")\n",
    "USGSHydroCached_layer = folium.TileLayer(tiles = 'https://basemap.nationalmap.gov/arcgis/rest/services/USGSHydroCached/MapServer/tile/{z}/{y}/{x}',\n",
    "                 attr = 'USGSHydroCached', zoom_start = zoom, name = \"USGSHydroCached\")\n",
    "\n",
    "# Format the inset map\n",
    "minimap = plugins.MiniMap(tile_layer= 'OpenStreetMap',\n",
    "                          #attr = 'USGS_topo', \n",
    "                          position='topleft',\n",
    "                          #zoom_level_offset=- 4,\n",
    "                          height = 200,\n",
    "                          width = 200,\n",
    "                          collapsed_height = 25,\n",
    "                          collapsed_width = 25,\n",
    "                          zoom_level_fixed = 5,\n",
    "                          toggle_display=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "# Set style functions for map\n",
    "style_function_hru_map = lambda x: {'opacity': 1,\n",
    "                                    'fillColor': '#00000000',#'goldenrod',\n",
    "                                    'color' : 'tan',\n",
    "                                    'weight': 1.5}\n",
    "highlight_function_hru_map = lambda x: {'opacity' : .5,\n",
    "                                        'color':'gray',\n",
    "                                        'fillColor':'gray',\n",
    "                                        'weight': 3}\n",
    "style_function_seg_map = lambda x: {'opacity': 1,\n",
    "                                    'color': '#217de7',\n",
    "                                    'weight': 2}\n",
    "highlight_function_seg_map = lambda x: {'opacity' : 1,\n",
    "                                        'color':'lightblue',\n",
    "                                        'weight': 4}\n",
    "transparent = lambda x: {'fillColor': '#00000000', \n",
    "                         'color': '#00000000',\n",
    "                         'weight': 4,\n",
    "                        }\n",
    "\n",
    "cp_style_function = lambda feature: {\"fillColor\": linear(var_sel_color_dict[feature[\"id\"]]),\n",
    "                                                    \"color\": 'tan',\n",
    "                                                    \"weight\": 1,\n",
    "                                                    #\"dashArray\": \"5, 5\",\n",
    "                                                    \"fillOpacity\": 0.3,\n",
    "                                }\n",
    "hw_basin_style = lambda x: {'fillColor': '#00000000', \n",
    "                               #'fill_opacity' : .8,\n",
    "                               'color': 'white',\n",
    "                               'weight': 2,\n",
    "                               #\"dashArray\": \"5, 5\",\n",
    "                           }\n",
    "cal_style_function=lambda feature: {\n",
    "                                \"fillColor\": \"gray\" if feature[\"properties\"][\"level\"] == 1\n",
    "                                else \"yellow\" if feature[\"properties\"][\"level\"] == 2\n",
    "                                else \"green\",\n",
    "                                \"color\": \"#00000000\",\n",
    "                                \"weight\": 1.5,\n",
    "                                #\"dashArray\": \"5, 5\",\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Zoom to Gage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "# # Zoom to gage, just uncomment the lines below\n",
    "\n",
    "# poi_lookup = poi_id_sel\n",
    "# lat = poi_df.loc[poi_df.poi_id==poi_lookup, 'latitude'].values[0]\n",
    "# lon = poi_df.loc[poi_df.poi_id==poi_lookup, 'longitude'].values[0]\n",
    "# zoom = 12\n",
    "\n",
    "# if poi_id_sel:\n",
    "#     poi_lookup = poi_id_sel\n",
    "#     lat = poi_df.loc[poi_df.poi_id==poi_lookup, 'latitude'].values[0]\n",
    "#     lon = poi_df.loc[poi_df.poi_id==poi_lookup, 'longitude'].values[0]\n",
    "#     zoom = 12\n",
    "# else:\n",
    "#     lat = hru_gdf_map['hru_lat'].mean()\n",
    "#     lon =  hru_gdf_map['hru_lon'].mean()-1\n",
    "#     zoom = 7"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set map\n",
    "if poi_id_sel:\n",
    "    poi_lookup = poi_id_sel\n",
    "    lat = poi_df.loc[poi_df.poi_id==poi_lookup, 'latitude'].values[0]\n",
    "    lon = poi_df.loc[poi_df.poi_id==poi_lookup, 'longitude'].values[0]\n",
    "    zoom = 12\n",
    "    poi_id_sel= None\n",
    "else:\n",
    "    lat = hru_gdf_map['hru_lat'].mean()\n",
    "    lon =  hru_gdf_map['hru_lon'].mean()-1\n",
    "    zoom = 7\n",
    "\n",
    "m = folium.Map(location=[lat, lon],\n",
    "               # width=800, height=600, \n",
    "                tiles = USGSHydroCached_layer, zoom_start = zoom\n",
    "               )\n",
    "folium.TileLayer(tiles = 'https://basemap.nationalmap.gov/arcgis/rest/services/USGSTopo/MapServer/tile/{z}/{y}/{x}',\n",
    "                 attr = 'USGS_Topo', zoom_start = zoom, name = \"USGS Topography\", show = False).add_to(m)\n",
    "\n",
    "\n",
    "################################################   \n",
    "# Create and add hru map\n",
    "hru_gdf_map =HW_basins_gdf.to_crs(crs)\n",
    "hru_map = folium.GeoJson(hru_gdf_map,\n",
    "                         style_function=cal_style_function,\n",
    "                         #highlight_function = highlight_function_hru_map,\n",
    "                         name = \"NHM HRUs\",\n",
    "                         z_index_offset = 40002).add_to(m)\n",
    "\n",
    "tooltip_hru=folium.GeoJsonPopup(fields= [\"hw_id\"],\n",
    "                                  aliases=[\"Headwater id\"],\n",
    "                                  labels=True)\n",
    "    \n",
    "\n",
    "#Add tool tip to map\n",
    "hru_map.add_child(tooltip_hru)\n",
    "\n",
    "################################################\n",
    "#Create and add segments map\n",
    "seg_gdf_map = seg_gdf.to_crs(crs)\n",
    "seg_map = folium.GeoJson(seg_gdf_map, \n",
    "                         style_function = style_function_seg_map,\n",
    "                         highlight_function = highlight_function_seg_map, # lambda feature: {\"fillcolor\": \"white\", \"color\": \"white\"},\n",
    "                         name = \"NHM Segments\",\n",
    "                         control = True,\n",
    "                         z_index_offset = 40003,\n",
    "                         ).add_to(m)\n",
    "\n",
    "tooltip_seg=folium.GeoJsonTooltip(fields= [\"nhm_seg\", \"tosegment_nhm\"],\n",
    "                                  aliases=[\"Segment\", \"flows to segment\"],\n",
    "                                  labels=True)\n",
    "seg_map.add_child(tooltip_seg)\n",
    "\n",
    "################################################\n",
    "#add POI marker clusters (marker and label)\n",
    "#byHRU_Group = folium.FeatureGroup(name='HRUs calibrated by HRU -- brown')\n",
    "marker_cluster = folium.FeatureGroup(\n",
    "    name='All the POIs',\n",
    "    overlay=True,\n",
    "    control=True,\n",
    "    icon_create_function=None,\n",
    "    z_index_offset = 5000,\n",
    ")\n",
    "marker_cluster_label_poi = folium.FeatureGroup(\n",
    "    name='All the POI labels',\n",
    "    overlay=True,\n",
    "    control=True,\n",
    "    show = False,# False will not draw the child upon opening the map, but have it to draw in the Layer control.\n",
    "    icon_create_function=None,\n",
    "    z_index_offset = 4004,\n",
    ")\n",
    "\n",
    "\n",
    "################################################\n",
    "# Add the inset map\n",
    "m.add_child(minimap)\n",
    "\n",
    "################################################\n",
    "\n",
    "\n",
    "\n",
    "################################################\n",
    "\n",
    "\n",
    "#Make Administrative basin labels\n",
    "label_coord_x = 20\n",
    "label_coor_y = 10\n",
    "\n",
    "\n",
    "for idx, row in poi_df.iterrows():\n",
    "    poi_id = row[\"poi_id\"]\n",
    "    var_plot_file = Folium_maps_dir / f'{output_var_sel}_{poi_id}.txt'\n",
    "        \n",
    "    if (row[\"nhm_calib\"] == \"Y\"):# Do this for all the gages used in calibration\n",
    "        if row[\"kge\"] >= 0.7:\n",
    "            \n",
    "            marker = folium.CircleMarker(location=[row['latitude'],row['longitude']],\n",
    "                                            name = row[\"poi_id\"],\n",
    "                                            popup = folium.Popup(f'Gage <b>{row[\"poi_id\"]}</b>, {row[\"poi_name\"]}<br>',\n",
    "                                                                   max_width = 150, \n",
    "                                                                   max_height = 70),\n",
    "                                            radius = 5,\n",
    "                                            weight = 2,\n",
    "                                            color = 'Black',\n",
    "                                            fill = True,\n",
    "                                            fill_color = 'Green',\n",
    "                                            fill_opacity = 1.0,\n",
    "                                            draggable = True,\n",
    "                                            lazy = True,\n",
    "                                            z_index_offset = 4006).add_to(marker_cluster)\n",
    "            \n",
    "            text = f'{row[\"poi_id\"]}'\n",
    "            label_lat = row['latitude']#-0.005\n",
    "            label_lon = row['longitude']\n",
    "            \n",
    "            marker_label = folium.map.Marker([label_lat, label_lon],\n",
    "                                             z_index_offset = 4007,\n",
    "                                             icon=DivIcon(icon_size=(150,36), \n",
    "                                                          icon_anchor=(0,0),\n",
    "                                                          html='<div style=\"font-size: 12pt; font-weight: bold\">%s</div>' % text)).add_to(marker_cluster_label_poi)\n",
    "        if (row[\"kge\"] < 0.7) & (row[\"kge\"] >= 0.5):\n",
    "                \n",
    "            marker = folium.CircleMarker(location=[row['latitude'],row['longitude']],\n",
    "                                            name = row[\"poi_id\"],\n",
    "                                            popup = folium.Popup(f'Gage <b>{row[\"poi_id\"]}</b>, {row[\"poi_name\"]}<br>',\n",
    "                                                                   max_width = 150, \n",
    "                                                                   max_height = 70),\n",
    "                                            radius = 5,\n",
    "                                            weight = 2,\n",
    "                                            color = 'Black',\n",
    "                                            fill = True,\n",
    "                                            fill_color = 'Yellow',\n",
    "                                            fill_opacity = 1.0,\n",
    "                                            draggable = True,\n",
    "                                            lazy=True,\n",
    "                                            z_index_offset = 4006).add_to(marker_cluster)\n",
    "            \n",
    "            #marker_cluster.add_child(marker)\n",
    "            text = f'{row[\"poi_id\"]}'\n",
    "            label_lat = row['latitude']#-0.005\n",
    "            label_lon = row['longitude']\n",
    "            \n",
    "            marker_label = folium.map.Marker([label_lat, label_lon],\n",
    "                                             z_index_offset = 4007,\n",
    "                                             icon=DivIcon(icon_size=(150,36), \n",
    "                                                          icon_anchor=(0,0),\n",
    "                                                          html='<div style=\"font-size: 12pt; font-weight: bold\">%s</div>' % text)).add_to(marker_cluster_label_poi)\n",
    "        if row[\"kge\"] <.5:\n",
    "            \n",
    "            marker = folium.CircleMarker(location=[row['latitude'],row['longitude']],\n",
    "                                            name = row[\"poi_id\"],\n",
    "                                            popup = folium.Popup(f'Gage <b>{row[\"poi_id\"]}</b>, {row[\"poi_name\"]}<br>',\n",
    "                                                                   max_width = 150, \n",
    "                                                                   max_height = 70),\n",
    "                                            radius = 5,\n",
    "                                            weight = 2,\n",
    "                                            color = 'Black',\n",
    "                                            fill = True,\n",
    "                                            fill_color = 'Red',\n",
    "                                            fill_opacity = 1.0,\n",
    "                                            draggable = True,\n",
    "                                            lazy=True,\n",
    "                                            z_index_offset = 4006).add_to(marker_cluster)\n",
    "            \n",
    "            #marker_cluster.add_child(marker)\n",
    "            text = f'{row[\"poi_id\"]}'\n",
    "            label_lat = row['latitude']#-0.005\n",
    "            label_lon = row['longitude']\n",
    "            \n",
    "            marker_label = folium.map.Marker([label_lat, label_lon],\n",
    "                                             z_index_offset = 4007,\n",
    "                                             icon=DivIcon(icon_size=(150,36), \n",
    "                                                          icon_anchor=(0,0),\n",
    "                                                          html='<div style=\"font-size: 12pt; font-weight: bold\">%s</div>' % text)).add_to(marker_cluster_label_poi)\n",
    "################################################\n",
    "        \n",
    "    ###########\n",
    "    if row[\"nhm_calib\"] == \"N\":\n",
    "        if (row[\"kge\"] >= 0.7):\n",
    "            \n",
    "            marker = folium.CircleMarker(location=[row['latitude'],row['longitude']],\n",
    "                                            name = row[\"poi_id\"],\n",
    "                                            popup = folium.Popup(f'Gage <b>{row[\"poi_id\"]}</b>, {row[\"poi_name\"]}<br>',\n",
    "                                                                   max_width = 150, \n",
    "                                                                   max_height = 70),\n",
    "                                            radius = 5,\n",
    "                                            weight = 2,\n",
    "                                            color = 'Lightgray',\n",
    "                                            fill = True,\n",
    "                                            fill_color = 'Green',\n",
    "                                            fill_opacity = 1.0,\n",
    "                                            draggable = True,\n",
    "                                            lazy=True,\n",
    "                                            z_index_offset = 4006).add_to(marker_cluster)\n",
    "            \n",
    "            #marker_cluster.add_child(marker)\n",
    "            text = f'{row[\"poi_id\"]}'\n",
    "            label_lat = row['latitude']#-0.005\n",
    "            label_lon = row['longitude']\n",
    "            \n",
    "            marker_label = folium.map.Marker([label_lat, label_lon],\n",
    "                                             z_index_offset = 4007,\n",
    "                                             icon=DivIcon(icon_size=(150,36), \n",
    "                                                          icon_anchor=(0,0),\n",
    "                                                          html='<div style=\"font-size: 12pt; font-weight: bold\">%s</div>' % text)).add_to(marker_cluster_label_poi)\n",
    "        if (row[\"kge\"] < 0.7) & (row[\"kge\"] >= 0.5):\n",
    "                        \n",
    "            marker = folium.CircleMarker(location=[row['latitude'],row['longitude']],\n",
    "                                            name = row[\"poi_id\"],\n",
    "                                            popup = folium.Popup(f'Gage <b>{row[\"poi_id\"]}</b>, {row[\"poi_name\"]}<br>',\n",
    "                                                                   max_width = 150, \n",
    "                                                                   max_height = 70),\n",
    "                                            radius = 5,\n",
    "                                            weight = 2,\n",
    "                                            color = 'Lightgray',\n",
    "                                            fill = True,\n",
    "                                            fill_color = 'Yellow',\n",
    "                                            fill_opacity = 1.0,\n",
    "                                            draggable = True,\n",
    "                                            lazy=True,\n",
    "                                            z_index_offset = 4006).add_to(marker_cluster)\n",
    "            \n",
    "            #marker_cluster.add_child(marker)\n",
    "            text = f'{row[\"poi_id\"]}'\n",
    "            label_lat = row['latitude']#-0.005\n",
    "            label_lon = row['longitude']\n",
    "                \n",
    "            marker_label = folium.map.Marker([label_lat, label_lon],\n",
    "                                                 z_index_offset = 4007,\n",
    "                                                 icon=DivIcon(icon_size=(150,36), \n",
    "                                                              icon_anchor=(0,0),\n",
    "                                                              html='<div style=\"font-size: 12pt; font-weight: bold\">%s</div>' % text)).add_to(marker_cluster_label_poi)\n",
    "        if (row[\"kge\"] <.5):\n",
    "            \n",
    "            marker = folium.CircleMarker(location=[row['latitude'],row['longitude']],\n",
    "                                            name = row[\"poi_id\"],\n",
    "                                            popup = folium.Popup(f'Gage <b>{row[\"poi_id\"]}</b>, {row[\"poi_name\"]}<br>',\n",
    "                                                                   max_width = 150, \n",
    "                                                                   max_height = 70),\n",
    "                                            radius = 5,\n",
    "                                            weight = 2,\n",
    "                                            color = 'Lightgray',\n",
    "                                            fill = True,\n",
    "                                            fill_color = 'Red',\n",
    "                                            fill_opacity = 1.0,\n",
    "                                            draggable = True,\n",
    "                                            lazy=True,\n",
    "                                            z_index_offset = 4006).add_to(marker_cluster)\n",
    "            \n",
    "            #marker_cluster.add_child(marker)\n",
    "            text = f'{row[\"poi_id\"]}'\n",
    "            label_lat = row['latitude']#-0.005\n",
    "            label_lon = row['longitude']\n",
    "                \n",
    "            marker_label = folium.map.Marker([label_lat, label_lon],\n",
    "                                                 z_index_offset = 4007,\n",
    "                                                 icon=DivIcon(icon_size=(150,36), \n",
    "                                                              icon_anchor=(0,0),\n",
    "                                                              html='<div style=\"font-size: 12pt; font-weight: bold\">%s</div>' % text)).add_to(marker_cluster_label_poi)\n",
    "        if np.isnan(row[\"kge\"]):\n",
    "            \n",
    "            marker = folium.CircleMarker(location=[row['latitude'],row['longitude']],\n",
    "                                                 name = row[\"poi_id\"],\n",
    "                                                 popup = folium.Popup(f'Gage <b>{row[\"poi_id\"]}</b>, {row[\"poi_name\"]}<br> Gage has less than 2yrs of observations.',\n",
    "                                                                   max_width = 150, \n",
    "                                                                   max_height = 70\n",
    "                                                                     ),\n",
    "                                                 radius = 2,\n",
    "                                                 weight = 2,\n",
    "                                                 color = 'Black',\n",
    "                                                 fill = True,\n",
    "                                                 fill_color = 'Black',\n",
    "                                                 fill_opacity = 1.0,\n",
    "                                                 draggable = True,\n",
    "                                                 lazy=True,\n",
    "                                                 z_index_offset = 4006).add_to(marker_cluster)\n",
    "                    \n",
    "            #marker_cluster.add_child(marker)\n",
    "            text = f'{row[\"poi_id\"]}'\n",
    "            label_lat = row['latitude']#-0.005\n",
    "            label_lon = row['longitude']\n",
    "                    \n",
    "            marker_label = folium.map.Marker([label_lat, label_lon],\n",
    "                                                  z_index_offset = 4007,\n",
    "                                                  icon=DivIcon(icon_size=(150,36), \n",
    "                                                               icon_anchor=(0,0),\n",
    "                                                               html='<div style=\"font-size: 12pt; font-weight: bold\">%s</div>' % text)).add_to(marker_cluster_label_poi)\n",
    "    # ################################################\n",
    "# Add hw boundary ref map\n",
    "hw_basins_map = folium.GeoJson(HW_basins, style_function = hw_basin_style, name = 'HW basin boundary').add_to(m)\n",
    "\n",
    "################################################\n",
    "marker_cluster.add_to(m)\n",
    "marker_cluster_label_poi.add_to(m)\n",
    "\n",
    "plugins.Fullscreen(position='topleft').add_to(m)\n",
    "folium.LayerControl(collapsed=True, position = 'bottomright', autoZIndex = True).add_to(m)\n",
    "\n",
    "################################################\n",
    "# Print map header\n",
    "con.print(f'')\n",
    "con.print(f'')\n",
    "con.print(f'')\n",
    "con.print('NHM poi_gages map', style = \"u bold black\")\n",
    "con.print('Click on a poi and copy the gage id into the field below to view hydrographs and flow statistics.', style = \"bold yellow\")\n",
    "\n",
    "m  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    " set(poi_df.nhm_calib.values)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Paste the poi_id in the field below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "v =widgets.Combobox(\n",
    "    # value='John',\n",
    "    placeholder='Enter Gage ID here',\n",
    "    options=poi_df.poi_id.tolist(),\n",
    "    description='Plot Gage:',\n",
    "    ensure_option=True,\n",
    "    disabled=False\n",
    ")\n",
    "\n",
    "def on_change(change):\n",
    "    global poi_id_sel, fig\n",
    "    if change['type'] == 'change' and change['name'] == 'value':\n",
    "        poi_id_sel = v.value\n",
    "        \n",
    "\n",
    "v.observe(on_change)\n",
    "\n",
    "display(v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "if poi_id_sel == None:\n",
    "    con.print('Select a gage ID from the dropdown above or copy/paste from the map into the field.')\n",
    "else:\n",
    "\n",
    "    # Single request   \n",
    "    if len((obs_annual.sel(poi_id = poi_id_sel)).to_dataframe().dropna()) < 2:\n",
    "        con.print (f'The gage {poi_id_sel} has no observation data in the streamflow obs file.')\n",
    "        pass\n",
    "    else:\n",
    "        df_sf_data_sel = (obs.sel(poi_id = poi_id_sel)).to_dataframe()\n",
    "                    \n",
    "        # Determine por\n",
    "        por_start = df_sf_data_sel['discharge'].notna().idxmax()# First Day\n",
    "        por_end = df_sf_data_sel['discharge'].notna()[::-1].idxmax()# Last Day\n",
    "        \n",
    "        # Slice to por\n",
    "        df_sf_data_sel = (obs.sel(poi_id = poi_id_sel, time=slice(por_start, por_end))).to_dataframe()\n",
    "        df_sf_data_sel.drop(columns =['poi_id'], inplace =True)# drop unwanted columns\n",
    "        \n",
    "        obs_efc_sel = (obs_efc.sel(poi_id = poi_id_sel, time=slice(por_start, por_end))).to_dataframe()\n",
    "        obs_efc_sel.drop(columns =['poi_id'], inplace =True)# drop unwanted columns\n",
    "        obs_with_efc_sel = (df_sf_data_sel.merge(obs_efc_sel, right_index = True, left_index = True, how = 'inner'))#.dropna() #how='left' will slice ts with obs range\n",
    "                    \n",
    "        sim_flow = (output_var.sel(npoi_gages = poi_id_sel, time=slice(por_start, por_end))).to_dataframe()\n",
    "        sim_flow.drop(columns =['npoi_gages', 'poi_gages'], inplace =True)# drop unwanted columns\n",
    "        \n",
    "        #Create a dataframe for the NaN's that occur between the beginning and end of por\n",
    "        daily_efc_df = (obs_with_efc_sel.merge(sim_flow, right_index = True, left_index = True, how = 'inner')).dropna()\n",
    "        daily_efc_plot_df = (obs_with_efc_sel.merge(sim_flow, right_index = True, left_index = True, how = 'inner'))\n",
    "        daily =  (df_sf_data_sel.merge(sim_flow, right_index = True, left_index = True, how = 'inner'))\n",
    "        daily_na = daily[daily['discharge'].isnull()]\n",
    "        daily_na['discharge'] = 5.0\n",
    "        \n",
    "        # drop the Nan's from the obs for memory/stats (may want to check back on this later)\n",
    "        daily_stat_df = (df_sf_data_sel.merge(sim_flow, right_index = True, left_index = True, how = 'inner')).dropna()\n",
    "        daily_plot_df = (df_sf_data_sel.merge(sim_flow, right_index = True, left_index = True, how = 'inner'))#.dropna()\n",
    "        \n",
    "        #daily_stat_df_na = daily_stat_df[daily_stat_df['discharge'].isnull()]\n",
    "        #daily_stat_df = daily_stat_df.dropna()\n",
    "        \n",
    "        #.dropna() #how='left' will slice ts with obs range\n",
    "        #daily_stat_df =streamflows_df.copy()#.dropna()\n",
    "        month_stat_df = daily_stat_df.resample('m').mean().dropna()\n",
    "        month_plot_df = daily_plot_df.resample('m').mean()#.dropna()\n",
    "        \n",
    "        water_year_stat_df = daily_stat_df.resample('A-SEP').mean().dropna()\n",
    "        water_year_plot_df = daily_plot_df.resample('A-SEP').mean()#.dropna()\n",
    "    \n",
    "        if len(daily_efc_df) <= 10000:\n",
    "            n = len(daily_efc_df)\n",
    "        else:\n",
    "            n = 10000 #Number of sampled days in records\n",
    "        \n",
    "        ######################################################\n",
    "        # Make timeseries subplot figure\n",
    "        fig = plotly.subplots.make_subplots(rows=3, cols=2, column_widths=[0.5, 0.5], #row_heights=[0., 0.3, 0.3, 0.4],\n",
    "                                shared_xaxes = 'columns',\n",
    "                                #shared_yaxes = 'columns',\n",
    "                                start_cell ='top-left',\n",
    "                                vertical_spacing = 0.1,\n",
    "                                horizontal_spacing = .06,\n",
    "                                y_title= f'Average daily streamflow, {getattr(model_output, output_var_sel).units}',\n",
    "                                subplot_titles = ['Annual mean',f'Flow Exceedence Curve, n = {n}',\n",
    "                                                  'Monthly mean',\n",
    "                                                  'Daily', 'Statistics',\n",
    "                                                 ], \n",
    "                                specs=[[{\"type\":\"scatter\"},{\"type\":\"scatter\", \"rowspan\": 2}],\n",
    "                                       [{\"type\":\"scatter\"}, None],\n",
    "                                       [{\"type\":\"scatter\"},{\"type\":\"table\"}],\n",
    "                                                  ]\n",
    "                                                   )\n",
    "        \n",
    "        station_name = station_name_df.loc[station_name_df.index == poi_id_sel, 'poi_name'].values[0]\n",
    "        date_range =  f'{daily_stat_df.index.month[0]}-{daily_stat_df.index.day[0]}-{daily_stat_df.index.year[0]} to {daily_plot_df.index.month[-1]}-{daily_plot_df.index.day[-1]}-{daily_plot_df.index.year[-1]} '\n",
    "        \n",
    "        fig.update_layout(title_text= f'NHM simulated streamflow at {poi_id_sel},<br>{station_name}, {date_range}', #\n",
    "                        width=900,\n",
    "                        height=700,\n",
    "                        legend=dict(orientation=\"h\",yanchor=\"bottom\",y=-.15, xanchor=\"right\", x=0.7),\n",
    "                        font=dict(family=\"Arial\",size=14,color=\"#7f7f7f\"), #font color\n",
    "                        paper_bgcolor='linen',\n",
    "                        plot_bgcolor='white',\n",
    "                        )\n",
    "        \n",
    "        fig.update_layout(title_automargin=True, \n",
    "                           title_font_color= 'black',\n",
    "                           title_font_size= 20,\n",
    "                           title_x = .5,\n",
    "                           title_y = .945,\n",
    "                           title_xref='container',\n",
    "                           title_xanchor='center')\n",
    "        \n",
    "        fig.update_xaxes(range = [daily_plot_df.index[0], daily_plot_df.index[-1]])\n",
    "        #fig.update_xaxes(range = [(obs[\"time\"][0].dt.strftime(\"%Y-%m-%d\").values.tolist()), (obs[\"time\"][-1].dt.strftime(\"%Y-%m-%d\").values.tolist())])\n",
    "        \n",
    "        #fig.update_layout(legend_grouptitlefont_color='black')\n",
    "        fig.update_layout(font_color='black')\n",
    "        \n",
    "        #fig.update_yaxes(title_text=f'{output_var_sel}, {getattr(model_output, output_var_sel).units}', title_font_color = 'black')\n",
    "        #fig.update_xaxes(title_text=\"Water years, from October 1 to September 31\", title_font_color = 'black')\n",
    "        \n",
    "        fig.update_xaxes(ticks=\"inside\", tickwidth=2, tickcolor='black', ticklen=10)\n",
    "        fig.update_yaxes(ticks=\"inside\", tickwidth=2, tickcolor='black', ticklen=10)\n",
    "            \n",
    "        fig.update_xaxes(showline=True, linewidth=2, linecolor='black', gridcolor='lightgrey')\n",
    "        fig.update_yaxes(showline=True, linewidth=2, linecolor='black', gridcolor='lightgrey')\n",
    "            \n",
    "        fig.update_traces(hovertemplate= None)\n",
    "        fig.update_layout(hovermode=\"x unified\")\n",
    "        fig.update_layout(hoverlabel=dict(bgcolor=\"linen\",\n",
    "                                          font_size=13,\n",
    "                                          font_family=\"Rockwell\",\n",
    "                                         )\n",
    "                         )\n",
    "        #Useful xarray calls\n",
    "        #f'{(obs[\"time\"][0].dt.strftime(\"%Y-%m-%d\").values.tolist())} to {(obs[\"time\"][-1].dt.strftime(\"%Y-%m-%d\").values.tolist())} '\n",
    "        #x_values_annual = (output_var_annual[\"time\"].dt.strftime(\"%Y-%m-%d\").values.tolist())\n",
    "        #sim_values_annual = (output_var_annual.sel(npoi_gages = poi_id_sel).values.tolist())\n",
    "        #obs_values = (obs_annual.sel(poi_id = poi_id_sel).values.tolist())\n",
    "        \n",
    "        ######################################################\n",
    "        # Create annual subplot\n",
    "        annual_plots = [go.Scatter(x=water_year_plot_df.index, \n",
    "                                 y= water_year_plot_df.discharge,\n",
    "                                 mode = 'lines',\n",
    "                                 name='Observed flow, annual',\n",
    "                                 showlegend = False,\n",
    "                                 #marker=dict(color='brown'),\n",
    "                                #xaxis = \n",
    "                                 line = dict(color='deepskyblue',\n",
    "                                     width=4, \n",
    "                                     #dash='dot'\n",
    "                                 )),\n",
    "                         go.Scatter(x= water_year_plot_df.index, \n",
    "                                 y= water_year_plot_df.seg_outflow,\n",
    "                                 mode = 'lines',\n",
    "                                 name='Simulated flow, annual',\n",
    "                                     showlegend = False,\n",
    "                                 #marker=dict(color='brown'), \n",
    "                                 line = dict(color='black',\n",
    "                                     width=1, \n",
    "                                     #dash='dot'\n",
    "                                             )\n",
    "                                )\n",
    "                        ]\n",
    "        annual_fig = go.Figure(data=annual_plots)\n",
    "        \n",
    "        ######################################################\n",
    "        # Create monthly subplot\n",
    "        monthly_plots = [go.Scatter(x= month_plot_df.index, \n",
    "                                 y=month_plot_df.discharge,\n",
    "                                 mode = 'lines',\n",
    "                                 name='Observed flow, monthly',\n",
    "                                 showlegend = False,\n",
    "                                 #marker=dict(color='brown'),\n",
    "                                #xaxis = \n",
    "                                 line = dict(color='deepskyblue',\n",
    "                                     width=4, \n",
    "                                     #dash='dot'\n",
    "                                 )),\n",
    "                         go.Scatter(x= month_plot_df.index,  \n",
    "                                 y= month_plot_df.seg_outflow,\n",
    "                                 mode = 'lines',\n",
    "                                 name='Simulated flow, monthly',\n",
    "                                 showlegend = False,\n",
    "                                 #marker=dict(color='brown'), \n",
    "                                 line = dict(color='black',\n",
    "                                     width=1, \n",
    "                                     #dash='dot'\n",
    "                                             )\n",
    "                                )\n",
    "                        ]\n",
    "        monthly_fig = go.Figure(data=monthly_plots)\n",
    "        \n",
    "        ######################################################\n",
    "        # Create daily subplot\n",
    "        # Make a line set for na values to show no data in the plot.\n",
    "        \n",
    "        #daily_efc_exlow_df = daily_efc_df.loc[daily_efc_df['efc'].isin([5])]\n",
    "        daily_efc_low_plot_df = daily_efc_plot_df.copy()\n",
    "        daily_efc_low_plot_df.loc[daily_efc_low_plot_df['efc']<=3, 'discharge'] = np.nan\n",
    "        \n",
    "        daily_efc_high_plot_df = daily_efc_plot_df.copy()\n",
    "        daily_efc_high_plot_df.loc[daily_efc_high_plot_df['efc']>=4, 'discharge'] = np.nan\n",
    "        \n",
    "        daily_plots = [go.Scatter(x= daily_efc_plot_df.index, #(output_var[\"time\"].dt.strftime(\"%Y-%m-%d\").values.tolist()), \n",
    "                                 y= daily_efc_plot_df.discharge, #(obs.sel(poi_id = poi_id_sel).values.tolist()),\n",
    "                                 mode = 'lines',\n",
    "                                 name='Observed flow',\n",
    "                                 showlegend = True,\n",
    "                                 connectgaps = False,\n",
    "                                 #marker=dict(color='deepskyblue', size = 5),\n",
    "                                 #xaxis = \n",
    "                                 line = dict(color='deepskyblue',\n",
    "                                     width=4, \n",
    "                                     #dash='dot'\n",
    "                                            )\n",
    "                                 ),\n",
    "                        go.Scatter(x= daily_efc_low_plot_df.index, #(output_var[\"time\"].dt.strftime(\"%Y-%m-%d\").values.tolist()), \n",
    "                                 y= daily_efc_low_plot_df.discharge, #(obs.sel(poi_id = poi_id_sel).values.tolist()),\n",
    "                                 mode = 'lines',\n",
    "                                 name='Observed flow, (Low)',\n",
    "                                 showlegend = True,\n",
    "                                 connectgaps = False,\n",
    "                                 #marker=dict(color='deepskyblue', size = 5),\n",
    "                                 #xaxis = \n",
    "                                 line = dict(color='red',\n",
    "                                     width=4, \n",
    "                                     #dash='dot'\n",
    "                                            )\n",
    "                                 ),\n",
    "                        \n",
    "        \n",
    "                         go.Scatter(x= daily_plot_df.index,# (output_var[\"time\"].dt.strftime(\"%Y-%m-%d\").values.tolist()), \n",
    "                                 y= daily_plot_df.seg_outflow,# (output_var.sel(npoi_gages = poi_id_sel).values.tolist()),\n",
    "                                 mode = 'lines',\n",
    "                                 name='Simulated flow, daily',\n",
    "                                 showlegend = True,\n",
    "                                 #marker=dict(color='black', size = 3), \n",
    "                                 line = dict(color='black',\n",
    "                                     width=1, \n",
    "                                     #dash='dot'\n",
    "                                             )\n",
    "                                )\n",
    "                        ]\n",
    "        #######################################################\n",
    "        # EFC classifications\n",
    "            # 1 = Large floods\n",
    "            # 2 = Small floods\n",
    "            # 3 = High flow pulses\n",
    "            # 4 = Low flows\n",
    "            # 5 = Extreme low flows\n",
    "        \n",
    "        daily_df = stats_table(daily_stat_df)\n",
    "        daily_df['time'] = 'daily'\n",
    "        monthly_df = stats_table(month_stat_df)\n",
    "        monthly_df['time'] = 'monthly'\n",
    "        annual_df = stats_table(water_year_stat_df)\n",
    "        annual_df['time'] = 'annual'\n",
    "        \n",
    "        #daily_efc_exlow_df = daily_efc_df.loc[daily_efc_df['efc'].isin([5])]\n",
    "        daily_efc_low_df = daily_efc_df.loc[daily_efc_df['efc'].isin([4, 5])]\n",
    "        daily_efc_high_df =  daily_efc_df.loc[daily_efc_df['efc'].isin([1,2,3])]\n",
    "        \n",
    "        #daily_exlow_tab_df = stats_table(daily_efc_exlow_df)\n",
    "        #daily_exlow_tab_df['time'] = 'exlow'\n",
    "        #daily_exlow_tab_df[['NSE','KGE']] = np.nan\n",
    "        \n",
    "        daily_low_tab_df = stats_table(daily_efc_low_df)\n",
    "        daily_low_tab_df['time'] = 'low'\n",
    "        daily_low_tab_df[['NSE','KGE']] = np.nan\n",
    "        \n",
    "        daily_high_tab_df = stats_table(daily_efc_high_df)\n",
    "        daily_high_tab_df['time'] = 'high'\n",
    "        daily_high_tab_df[['NSE','KGE']] = np.nan\n",
    "        \n",
    "        all_df = pd.concat([daily_df, daily_low_tab_df, daily_high_tab_df, monthly_df, annual_df,])\n",
    "        all_df.set_index('time', inplace=True)\n",
    "        stats_table_df = all_df.T\n",
    "        #stats_table_df\n",
    "        \n",
    "        stats_table_obj = go.Figure(data=[go.Table(header=dict(values=['Statistic','Daily', 'Low', 'High','Monthly', 'Annual']),\n",
    "                         cells=dict(values=[stats_table_df.index, stats_table_df.daily, \n",
    "                                            stats_table_df.low, stats_table_df.high,\n",
    "                                            stats_table_df.monthly, stats_table_df.annual]))\n",
    "        ])\n",
    "        \n",
    "        #######################################################\n",
    "       \n",
    "        obs_data = daily_efc_df.discharge.sample(n = n,#frac=0.25,\n",
    "                                                   replace = False,\n",
    "                                                   random_state = 3)\n",
    "        sim_data = daily_efc_df.seg_outflow.sample(n = n,#frac=0.25,\n",
    "                                                     replace = False,\n",
    "                                                     random_state = 3)\n",
    "        \n",
    "        obs_sort = np.sort(obs_data)[::-1]\n",
    "        sim_sort = np.sort(sim_data)[::-1]\n",
    "        obs_color_sort = daily_efc_df.sort_values('discharge')[::-1]# Makes the color value sort in same order for use in plot.\n",
    "        \n",
    "        obs_exceedence = np.arange(1.,len(obs_sort)+1) / len(obs_sort)\n",
    "        sim_exceedence = np.arange(1.,len(sim_sort)+1) / len(sim_sort)\n",
    "        \n",
    "        efc_colors = {1: 'rgba(0, 191, 255, 0.5)',# Large Floods\n",
    "                       0: 'white',\n",
    "                       2: 'rgba(0, 191, 255, 0.5)',# Small Floods\n",
    "                       3: 'rgba(0, 191, 255, 0.5)',# High Flow Pulse\n",
    "                       4: 'rgba(255, 0, 0, 0.5)',# Low\n",
    "                       5: 'rgba(255, 0, 0, 0.5)',# Extreemly Low\n",
    "                     np.nan: 'yellow'} # missing\n",
    "         # or ...color_discrete_sequence = plotly.colors.sequential.Viridis\n",
    "        \n",
    "        custom_marker_color = obs_color_sort['efc'].map(efc_colors)\n",
    "        \n",
    "        exceed_plot =[go.Scatter(x=obs_exceedence,\n",
    "                                 y= obs_sort,\n",
    "                                 mode = 'markers',\n",
    "                                 name='Observed flow',\n",
    "                                 marker=dict(color=custom_marker_color,\n",
    "                                            size= 3),\n",
    "                                 showlegend = False,\n",
    "                                 #line = dict(color='deepskyblue',\n",
    "                                 #    width=3,\n",
    "                                    # dash='dot'\n",
    "                                            #)\n",
    "                                ),\n",
    "                      go.Scatter(x=sim_exceedence, \n",
    "                                 y= sim_sort,\n",
    "                                 mode = 'lines',\n",
    "                                 name='NHM simulated flow',\n",
    "                                 showlegend = False,\n",
    "                                # marker=dict(#color='brown',\n",
    "                                #            size=1), \n",
    "                                 line = dict(color='black',\n",
    "                                     width=1, \n",
    "                                     #dash='dot'\n",
    "                                            )\n",
    "                                )\n",
    "        ]\n",
    "        \n",
    "        exceed_fig = go.Figure(data=exceed_plot)\n",
    "        \n",
    "        #fig.update_yaxes(title_text=f'Streamflow, {getattr(model_output, \"seg_outflow\").units}', title_font_color = 'black', row=1, col=3)\n",
    "        #fig.update_xaxes(title_text=\"Exceedence, probability\", title_font_color = 'black', row=1, col=3)\n",
    "        \n",
    "        fig.update_yaxes(type=\"log\", col=2)\n",
    "        \n",
    "        tickvals = [0, 1, 2, 5, 10,\n",
    "                    20, 50, 100,\n",
    "                    200, 500, 1000,\n",
    "                    2000, 5000, 10000,\n",
    "                    20000, 50000, 100000,\n",
    "                    200000, 500000, 1000000]\n",
    "        \n",
    "        tickvals_exceed = [0, .25, .5, .75, 1]\n",
    "        \n",
    "        fig.update_xaxes(tickvals=tickvals_exceed,ticks=\"inside\", tickwidth=2, tickcolor='black', showticklabels=True, ticklen=10, col=2)\n",
    "        fig.update_yaxes(tickvals=tickvals, \n",
    "                         ticks=\"inside\", tickwidth=2, tickcolor='black', ticklen=10, col=2)\n",
    "            \n",
    "        fig.update_xaxes(showline=True, linewidth=2, linecolor='black', gridcolor='lightgrey', range = [-0.1,1.1], col=2)\n",
    "        fig.update_yaxes(showline=True, linewidth=2, linecolor='black', gridcolor='lightgrey', col=2)\n",
    "            \n",
    "        #######################################################\n",
    "        # Add plots and stats tables to figure\n",
    "        daily_fig = go.Figure(data=daily_plots)\n",
    "        \n",
    "        for t in annual_fig.data:\n",
    "            fig.append_trace(t, row=1, col=1)\n",
    "        for t in monthly_fig.data:\n",
    "            fig.append_trace(t, row=2, col=1)\n",
    "        for t in daily_fig.data:\n",
    "            fig.append_trace(t, row=3, col=1)               \n",
    "        for t in exceed_fig.data:\n",
    "            fig.append_trace(t, row=1, col=2)\n",
    "        for t in stats_table_obj.data:\n",
    "            fig.append_trace(t, row=3, col=2)\n",
    "    \n",
    "                   \n",
    "        fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add D scores into here (sydney and Tim); maybe use the differenct components of the EFC rating"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}