### Train a Model

Source: https://openhydronet.readthedocs.io/en/latest/_sources/usage/quickstart.rst.txt

Use this command to start the training process for a model. Ensure you have a valid configuration file.

```bash
run train --config-file ~/tutorial/training-config.yml
```

--------------------------------

### Install Google Cloud Storage and GCSFS

Source: https://openhydronet.readthedocs.io/en/latest/tutorial/OpenHydroNet_Tutorial.html

Installs the necessary libraries for Google Cloud Storage integration.

```python
# !pip install -q google-cloud-storage gcsfs

```

--------------------------------

### Print Run Instructions

Source: https://openhydronet.readthedocs.io/en/latest/tutorial/OpenHydroNet_Tutorial.html

Prints detailed instructions for executing fine-tuning and inference commands to the console. This includes configuration file locations and command examples.

```python
finetune_experiment_name = f"finetune-{FINE_TUNING_BASIN}"
finetune_infer_command_display = f"run infer --run-dir={model_run_dir}/{finetune_experiment_name}_YYYYMMDD_HHMMSS"


# Print the run instructions to the screen
print("--------------------------------------------------------------------------------------------------")
print("                       FINE-TUNING AND INFERENCE RUN INSTRUCTIONS                           ")
print("--------------------------------------------------------------------------------------------------")
print("These commands need to be executed in your terminal or a separate process.")
print("Ensure you are in an environment with the necessary Python libraries (e.g., the Conda environment).")
print("")
print("IMPORTANT: The configuration file has been automatically generated for your selected basin:")
print(f"  Config File Location: {finetune_basin_config_file}")
print("  This file contains the specific settings for fine-tuning on basin " + str(FINE_TUNING_BASIN))
print("  It has been pre-configured with the correct base model paths and basin list.")
print("")
print("STEP 1: Run the fine-tuning training command:")
print(f"    {finetune_train_command}")
print("")
print("STEP 2: After training is complete, find the created run directory and run inference.")
print(f"    The output directory will be created inside: {model_run_dir}")
print(f"    It will be named like: {finetune_experiment_name}_YYYYMMDD_HHMMSS")
print("")
print("    Run the inference command using that directory:")
print(f"    {finetune_infer_command_display}")
print("")
print("--------------------------------------------------------------------------------------------------")
```

--------------------------------

### Setup Interactive Visualization Widgets

Source: https://openhydronet.readthedocs.io/en/latest/tutorial/OpenHydroNet_Tutorial.html

Prepares interactive dropdown widgets for selecting a basin and a metric to visualize performance comparisons. It populates these widgets with available basin IDs and metrics from the base model data.

```python
# --- Fine-Tuning: Visualize Metrics ---

# This cell provides an interactive visualization to compare the performance
# of the base model against the fine-tuned model for a specific basin and metric.
# This allows you to assess the impact of fine-tuning on the target basin's predictions
# across different lead times.

# Get the list of all unique basin IDs present in the base model metrics.
# This ensures that the dropdown for basin selection contains only valid basin IDs
# for which data is available.
all_basin_ids = sorted(list(base_model_metrics.index.get_level_values('basin_id').unique()))

# Get the list of available metrics from the base model metrics columns.
# This ensures the dropdown only presents metrics that have actually been calculated and are present.
available_metrics = list(base_model_metrics.columns)

# Create a dropdown widget for selecting the basin.
# The default value is set to the `FINE_TUNING_BASIN`.
# We ensure it matches the format in all_basin_ids.
basin_widget = widgets.Dropdown(
    options=all_basin_ids,
    value=FINE_TUNING_BASIN if FINE_TUNING_BASIN in all_basin_ids else all_basin_ids[0],
    description='Select Basin:',
    disabled=False,
    style = {'description_width': 'initial'}
)

# Create a dropdown widget for selecting the metric.
# The default metric is 'KGE' (Kling-Gupta Efficiency), a common hydrological performance metric.
metric_widget = widgets.Dropdown(
    options=available_metrics,
    value='KGE',
    description='Select Metric:',
    disabled=False,
    style = {'description_width': 'initial'}
)

# Create an interactive widget linking the dropdowns to the plotting function.
# `plot_comparison_metrics_vs_lead_time` will dynamically update the plot
# whenever a new basin or metric is selected from the dropdowns.
# `widgets.fixed` is used to pass the metrics DataFrames as static arguments

```

--------------------------------

### Install Dependencies for OpenHydroNet Tutorial

Source: https://openhydronet.readthedocs.io/en/latest/tutorial/OpenHydroNet_Tutorial.html

Installs necessary libraries for data visualization, scientific computing, and Google Cloud Storage access. These are typically not needed if using the provided conda environment.

```python
# # --- Install Dependencies ---
# # These lines ensure that all required third-party libraries are installed.
# # These are not needed if you use the supplied conda
# # environment in `~/flood-forecasting/environments/conda.yml`

# # Data visualization and interactive widgets
# !pip install -q matplotlib seaborn ipywidgets

# # Scientific computing and data structures
# !pip install -q numpy pandas xarray netCDF4

# # Google Cloud Storage (for MultiMet data access)

```

--------------------------------

### Interactive Metrics Visualization Setup

Source: https://openhydronet.readthedocs.io/en/latest/_sources/tutorial/OpenHydroNet_Tutorial.ipynb.txt

Sets up interactive widgets (dropdowns) for basin and metric selection to visualize and compare model performance. The plot dynamically updates based on user selections.

```python
# Get the list of all unique basin IDs present in the base model metrics.
# This ensures that the dropdown for basin selection contains only valid basin IDs
# for which data is available.
all_basin_ids = sorted(list(base_model_metrics.index.get_level_values('basin_id').unique()))

# Get the list of available metrics from the base model metrics columns.
# This ensures the dropdown only presents metrics that have actually been calculated and are present.
available_metrics = list(base_model_metrics.columns)

# Create a dropdown widget for selecting the basin.
# The default value is set to the `FINE_TUNING_BASIN`.
# We ensure it matches the format in all_basin_ids.
basin_widget = widgets.Dropdown(
    options=all_basin_ids,
    value=FINE_TUNING_BASIN if FINE_TUNING_BASIN in all_basin_ids else all_basin_ids[0],
    description='Select Basin:',
    disabled=False,
    style = {'description_width': 'initial'}
)

# Create a dropdown widget for selecting the metric.
# The default metric is 'KGE' (Kling-Gupta Efficiency), a common hydrological performance metric.
metric_widget = widgets.Dropdown(
    options=available_metrics,
    value='KGE',
    description='Select Metric:',
    disabled=False,
    style = {'description_width': 'initial'}
)

# Create an interactive widget linking the dropdowns to the plotting function.
# `plot_comparison_metrics_vs_lead_time` will dynamically update the plot
# whenever a new basin or metric is selected from the dropdowns.
# `widgets.fixed` is used to pass the metrics DataFrames as static arguments
# because they do not change with dropdown selections.
interactive_plot = interactive(
    backend.plot_comparison_metrics_vs_lead_time,
    base_metrics_df=widgets.fixed(base_model_metrics), # Pass the base metrics DataFrame
    finetune_metrics_df=widgets.fixed(finetune_metrics),
    basin_id=basin_widget,
    metric_name=metric_widget
)

# Display the interactive widget and the plot output.
# The `interactive_plot` object contains the controls (dropdowns) and the output (plot).
# We separate them to display controls above the plot.
uis = interactive_plot.children[:-1] # Controls are all children except the last one (the output)
out = interactive_plot.children[-1] # The output is the last child

display(VBox(uis), out) # Display controls and output in a VBox
```

--------------------------------

### Install Python Dependencies

Source: https://openhydronet.readthedocs.io/en/latest/_sources/tutorial/OpenHydroNet_Tutorial.ipynb.txt

Installs required third-party libraries for data visualization, scientific computing, and cloud storage access. These commands are not needed if using the provided conda environment.

```bash
# Data visualization and interactive widgets
# !pip install -q matplotlib seaborn ipywidgets

# Scientific computing and data structures
# !pip install -q numpy pandas xarray netCDF4

# Google Cloud Storage (for MultiMet data access)
# !pip install -q google-cloud-storage gcsfs
```

--------------------------------

### Example Model Forward Pass

Source: https://openhydronet.readthedocs.io/en/latest/_sources/usage/models.rst.txt

This snippet shows an example of a forward pass within a custom model. It retrieves hindcast, forecast, and static data, and returns a placeholder output. Replace the placeholder logic with your actual forecasting implementation.

```python
# Example forward pass (remove when you add your own modeling logic)
hindcast = data['x_d_hindcast']
forecast = data['x_d_forecast']
statics = data['x_s']

# ... implement your forecasting logic here ...

# Example placeholder output (remove when you add your own modeling logic)
batch_size, seq_len, _ = hindcast.shape
output = torch.zeros(batch_size, seq_len, self.output_size)

return {'y_hat': output}
```

--------------------------------

### Install Package in Editable Mode

Source: https://openhydronet.readthedocs.io/en/latest/_sources/usage/quickstart.rst.txt

Install the OpenHydroNet package in editable mode using pip. This allows immediate reflection of code changes. Run this from the root of the flood-forecasting directory.

```bash
# Run this from the root of the flood-forecasting directory
pip install -e .
```

--------------------------------

### Define Local Data Paths

Source: https://openhydronet.readthedocs.io/en/latest/_sources/tutorial/OpenHydroNet_Tutorial.ipynb.txt

Sets up essential file paths for the tutorial, including shapefiles for basin geometries, attribute files, model run directories, and configuration directories. Also includes flags to control model retraining and statistics recalculation.

```python
# Define the path to the shapefile containing basin geometries.
# This shapefile is used for plotting maps of basin locations.
# Users may need to update this path if their data is stored elsewhere.
SHAPEFILE_PATH = '~/flood-forecasting/tutorial/Caravan-nc/shapefiles/camels/camels_basin_shapes.shx'

# Path to the attributes file that contains additional basin information, such as basin area.
# This file is used for visualizing the distribution of basin areas.
ATTRIBUTES_FILE_PATH = '~/flood-forecasting/tutorial/Caravan-nc/attributes/camels/attributes_other_camels.csv'

# Path to a base directory containing one or more model run directories.
# This directory should house all your trained models.
# The interactive selection widgets in subsequent cells will scan this directory to allow you to choose which model run to evaluate.
MODEL_RUN_DIR = '/home/gsnearing/flood-forecasting/tutorial/model-runs'

# The directory where .yml model configuration files are located.
# These files define parameters for model runs, including data paths and training settings.
CONFIG_DIR = '/home/gsnearing/flood-forecasting/tutorial/configs'

# --- OPTIONAL: Full Training Pipeline Trigger ---
# By default, this tutorial uses pre-trained model outputs to save time.
# Set this flag to 'True' ONLY if you want to perform a fresh training run
# and generate new predictions from scratch.
RUN_NEW_MODELS = False

# A boolean flag to control whether performance statistics (metrics) are recalculated or loaded from pre-existing files.
# Set this to 'True' for the initial run or if model outputs have changed and you need fresh calculations.
# Set to 'False' if metrics have already been computed and saved, to save time during subsequent runs.
# This is particularly useful for live demos or quick comparisons where recalculation is not necessary.
CALCULATE_STATISTICS = False
```

--------------------------------

### Execute Training and Inference

Source: https://openhydronet.readthedocs.io/en/latest/_sources/tutorial/OpenHydroNet_Tutorial.ipynb.txt

This snippet demonstrates how to execute the training command and then run inference using the trained model. It includes loading configuration, finding the most recent run directory, and initiating the inference process.

```python
# WARNING: Full training can take a significant amount of time (minutes to hours)
# depending on your hardware (CPU/GPU) and the number of basins.
print(f"Executing training: {base_train_command}")
os.system(base_train_command)

# Load the base config file to get the actual experiment name
with open(base_config_file, 'r') as f:
    base_config_data = yaml.safe_load(f)
actual_experiment_name = base_config_data.get('experiment_name', 'default_experiment')
print(f"Experiment name from config file: {actual_experiment_name}")

# Get all potential run directories for the base experiment
search_pattern = os.path.join(MODEL_RUN_DIR, f'{actual_experiment_name}_*_*/')
all_run_dirs = glob.glob(search_pattern)

if not all_run_dirs:
    print(f"Error: No run directories found matching pattern {search_pattern}")
else:
    # Sort them by modification time to get the most recent one
    # Or, parse the timestamp from the directory name if it's consistently formatted
    # For now, let's assume the default `googlehydrology` naming `experiment_YYYYMMDD_HHMMSS` makes lexicographical sort work.
    all_run_dirs.sort()
    base_model_actual_run_dir = all_run_dirs[-1].rstrip('/') # Get the last (most recent) and remove trailing slash
    print(f"Found most recent base model run directory: {base_model_actual_run_dir}")

    # 2. Execute Inference:
    # After training completes, this command generates the actual streamflow
    # predictions (.nc or .zarr files) used for the hydrographs later in this notebook.
    infer_command = f"run infer --run-dir={base_model_actual_run_dir}"
    print(f"Executing inference: {infer_command}")
    os.system(infer_command)
else:
    print("Skipping re-training. You will select the base model in the next step.")
```

--------------------------------

### Download Source Code as Zip

Source: https://openhydronet.readthedocs.io/en/latest/_sources/usage/quickstart.rst.txt

Download the source code as a zip file if you do not use git. This method requires manual extraction and navigation.

```bash
# Download the source code zip file
curl -L https://github.com/google-research/flood-forecasting/zipball/master -o flood-forecasting.zip

# Extract the archive
unzip flood-forecasting.zip

# Enter the resulting directory (folder name may vary based on the specific commit)
cd google-research-flood-forecasting-*
```

--------------------------------

### Print Training and Inference Instructions

Source: https://openhydronet.readthedocs.io/en/latest/tutorial/OpenHydroNet_Tutorial.html

This Python code prints instructions for running base model training and inference commands to the console. It highlights the configuration file and the necessary steps.

```python
base_infer_command = f"run infer --run-dir={MODEL_RUN_DIR}<YOUR_BASE_MODEL_RUN_DIR>"

# Print the run instructions to the screen
print("--------------------------------------------------------------------------------------------------")
print("                          BASE MODEL TRAINING AND INFERENCE INSTRUCTIONS                          ")
print("--------------------------------------------------------------------------------------------------")
print("These commands need to be executed in your terminal or a separate process.")
print("Ensure you are in an environment with the necessary Python libraries (e.g., the Conda environment).")
print("")
print("IMPORTANT: Before running these commands, please examine the configuration file:")
print(f"  Config File Location: {base_config_file}")
print("  This file defines crucial parameters for your model run, including data paths, model architecture, and training settings.")
print("  Understanding its contents will help you customize and debug your experiments.")
print("")
print("STEP 1: Run the base model training command FIRST (if you need to train a new base model):")
print(f"    {base_train_command}")
print("")
print("STEP 2: After the base model training is complete, run the inference command:")
print("    (IMPORTANT: Replace '<YOUR_BASE_MODEL_RUN_DIR>' with the actual run directory created by STEP 1, e.g., 'model-runs/5-basin-example_YYYYMMDD_HHMMSS')")
print(f"    {base_infer_command}")
print("")
print("--------------------------------------------------------------------------------------------------")
```

--------------------------------

### Generate Fine-Tuning Configuration File

Source: https://openhydronet.readthedocs.io/en/latest/tutorial/OpenHydroNet_Tutorial.html

Creates a basin-specific configuration file for fine-tuning by first generating a basin list file and then using a template to create the final config.

```python
# --- Configuration File Paths ---
finetune_config_template_path = f'{CONFIG_DIR}/finetune-config.yml'
finetune_basin_config_file = f'{CONFIG_DIR}/finetune-config-{FINE_TUNING_BASIN}.yml'

# --- Execute Config Generation ---

# 1. Create the basin list file
backend.create_basin_list_file(FINE_TUNING_BASIN, output_dir='basin-lists')

# 2. Generate the config file
# Construct the full path to the base model run directory
base_model_run_dir_path = os.path.join(MODEL_RUN_DIR, os.path.basename(model_run_dir))

backend.generate_basin_finetune_config(
    template_path=finetune_config_template_path,
    basin_id=FINE_TUNING_BASIN,
    base_model_dir=base_model_run_dir_path,
    output_path=finetune_basin_config_file
)
```

--------------------------------

### Construct Fine-Tuning Run Command

Source: https://openhydronet.readthedocs.io/en/latest/tutorial/OpenHydroNet_Tutorial.html

Generates the command-line string to execute the fine-tuning training process using a specific configuration file.

```bash
# --- Fine-Tuning: Generate Run Command ---

# This cell generates the commands needed to run the fine-tuning experiment
# and subsequent inference for the selected basin.

# Assumes FINE_TUNING_BASIN and finetune_basin_config_file are available globally.

# Construct the fine-tuning training command.
# This points to the basin-specific config file we just generated.
finetune_train_command = f"run finetune --config-file={finetune_basin_config_file}"

# Construct the fine-tuning inference command.

```

--------------------------------

### Create Dropdown for Base Model Selection

Source: https://openhydronet.readthedocs.io/en/latest/tutorial/OpenHydroNet_Tutorial.html

Creates a dropdown widget to select the base model from available run directories. Sets the initial value to the latest run name or the first available option.

```python
run_dir_options = sorted(available_run_dirs.keys())
base_dropdown = widgets.Dropdown(
    options=run_dir_options,
    description='Select Base Model:',
    disabled=False,
    style = {'description_width': 'initial'},
    value=latest_run_name if latest_run_name in run_dir_options else run_dir_options[0]
)
```

--------------------------------

### Execute Base Model Training

Source: https://openhydronet.readthedocs.io/en/latest/_sources/tutorial/OpenHydroNet_Tutorial.ipynb.txt

Initiates the training process for the base model using a specified YAML configuration file. This command should be run first if a new base model needs to be trained.

```python
if RUN_NEW_MODELS:
    # 1. Execute Base Model Training:
    # This calls the 'run train' command with your specified YAML config.
    
```

--------------------------------

### Define User-Defined Local Paths

Source: https://openhydronet.readthedocs.io/en/latest/tutorial/OpenHydroNet_Tutorial.html

Sets up paths for shapefiles, attribute data, model run directories, and configuration files. Users should update these paths to match their local environment.

```python
# Define the path to the shapefile containing basin geometries.
# This shapefile is used for plotting maps of basin locations.
# Users may need to update this path if their data is stored elsewhere.
SHAPEFILE_PATH = '~/flood-forecasting/tutorial/Caravan-nc/shapefiles/camels/camels_basin_shapes.shx'

# Path to the attributes file that contains additional basin information, such as basin area.
# This file is used for visualizing the distribution of basin areas.
ATTRIBUTES_FILE_PATH = '~/flood-forecasting/tutorial/Caravan-nc/attributes/camels/attributes_other_camels.csv'

# Path to a base directory containing one or more model run directories.
# This directory should house all your trained models.
# The interactive selection widgets in subsequent cells will scan this directory to allow you to choose which model run to evaluate.
MODEL_RUN_DIR = '/home/gsnearing/flood-forecasting/tutorial/model-runs'

# The directory where .yml model configuration files are located.
# These files define parameters for model runs, including data paths and training settings.
CONFIG_DIR = '/home/gsnearing/flood-forecasting/tutorial/configs'

# --- OPTIONAL: Full Training Pipeline Trigger ---
# By default, this tutorial uses pre-trained model outputs to save time.
# Set this flag to 'True' ONLY if you want to perform a fresh training run
# and generate new predictions from scratch.
RUN_NEW_MODELS = False

# A boolean flag to control whether performance statistics (metrics) are recalculated or loaded from pre-existing files.
# Set this to 'True' for the initial run or if model outputs have changed and you need fresh calculations.
# Set to 'False' if metrics have already been computed and saved, to save time during subsequent runs.
# This is particularly useful for live demos or quick comparisons where recalculation is not necessary.
CALCULATE_STATISTICS = False

```

--------------------------------

### Template for Implementing a New Model

Source: https://openhydronet.readthedocs.io/en/latest/usage/models.html

This template provides the basic structure for creating a new forecasting model. Ensure you inherit from BaseModel, define `module_parts`, and implement the `forward` method. The `cfg` dictionary is used for initialization, and the `forward` method accepts dynamic and static inputs, returning predictions under the 'y_hat' key.

```python
import torch

from googlehydrology.modelzoo.basemodel import BaseModel


class TemplateModel(BaseModel):

    # The `module_parts` variable is a list of all of the different model components.
    # You must construct and name these components. This is necessary in order to freeze
    # and unfreeze individual components for fine tuning.
    module_parts = [...] 

    def __init__(self, cfg: dict):
        """Initialize the model

        Each model receives as only input the config dictionary. From this, the entire model can be implemented.

        Each Model inherits from the BaseModel, which implements some universal functionality. The basemodel also
        defines the output_size, which can be used here as a given attribute (self.output_size)

        Parameters
        ----------
        cfg : dict
            Configuration of the run, read from the config file with some additional keys (such as number of basins).
        """
        super(TemplateModel, self).__init__(cfg=cfg)

        ###########################
        # Create model parts here #
        ###########################

    def forward(self, data: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
        """Forward pass through the model

        Parameters
        ----------
        - 'x_d_hindcast': Hindcast dynamic inputs
          Shape: [batch, seq_length, features]
        - 'x_d_forecast': Forecast dynamic inputs
          Shape: [batch, lead_time, features]
        - 'x_s': Static inputs
          Shape: [batch, features]

        Returns
        -------
        dict[str, torch.Tensor]
            The dictionary must contain the key 'y_hat' with predictions.
            Shape: [batch, seq_length, num_targets]
        """

        ###############################
        # Implement forward pass here #
        ###############################

        # Example forward pass (remove when you add your own modeling logic)
        hindcast = data['x_d_hindcast']
        forecast = data['x_d_forecast']
        statics = data['x_s']

        # ... implement your forecasting logic here ...

        # Example placeholder output (remove when you add your own modeling logic)
        batch_size, seq_len, _ = hindcast.shape
        output = torch.zeros(batch_size, seq_len, self.output_size)

        return {'y_hat': output}

```

--------------------------------

### Clone the Repository

Source: https://openhydronet.readthedocs.io/en/latest/_sources/usage/quickstart.rst.txt

Use git to clone the flood-forecasting repository to your local machine. This is the recommended method for obtaining the source code.

```bash
git clone https://github.com/google-research/flood-forecasting.git
cd flood-forecasting
```

--------------------------------

### Select Base Model with Dropdown Widget

Source: https://openhydronet.readthedocs.io/en/latest/_sources/tutorial/OpenHydroNet_Tutorial.ipynb.txt

This code creates an interactive dropdown widget to select a base model from available run directories. It uses `backend.find_model_run_dirs` to discover runs and `widgets.Dropdown` for user selection.

```python
# Find available run directories using the path set in the cell above
available_run_dirs = backend.find_model_run_dirs(MODEL_RUN_DIR)

# Extract the name of the most recent model run for default selection
if 'base_model_actual_run_dir' in globals():
    latest_run_name = os.path.basename(base_model_actual_run_dir)
else:
    latest_run_name = None

# Create a dropdown widget for selecting the base model
run_dir_options = sorted(available_run_dirs.keys())
base_dropdown = widgets.Dropdown(
    options=run_dir_options,
    description='Select Base Model:',
    disabled=False,
    style = {'description_width': 'initial'},
    value=latest_run_name if latest_run_name in run_dir_options else run_dir_options[0]
)

# Create an interactive widget linking the dropdown to the selection function
interactive_selection = interactive(
    select_model,
    model_selection=base_dropdown,
)

# Display the interactive widget
display(interactive_selection)
```

--------------------------------

### Select Fine-Tuning Basin with Interactive Widget

Source: https://openhydronet.readthedocs.io/en/latest/tutorial/OpenHydroNet_Tutorial.html

Creates an interactive dropdown widget to select a basin for fine-tuning. It updates a global variable and plots the selected basin's shapefile.

```python
DEFAULT_FINE_TUNING_BASIN_ID = 'camels_13235000' # Example for small model run in live demo.

# Get the list of test basin IDs from the base model configuration.
# These are the basins on which the base model was evaluated and from which
# we can select a target for fine-tuning.
# `test_basin_ids` is expected to be available globally from previous cells
# (e.g., after `backend.load_model_config_and_basins` in the 'Base Model' section).
fine_tuning_basin_options = sorted(list(test_basin_ids))

# Create a dropdown widget for selecting the fine-tuning basin.
# The options are derived from the `test_basin_ids` of the base model.
# The default value is set to `DEFAULT_FINE_TUNING_BASIN_ID` if it exists
# in the options, otherwise it defaults to the first available basin.
basin_selector_dropdown = widgets.Dropdown(
    options=fine_tuning_basin_options,
    value=DEFAULT_FINE_TUNING_BASIN_ID if DEFAULT_FINE_TUNING_BASIN_ID in fine_tuning_basin_options else fine_tuning_basin_options[0],
    description='Select Fine-Tuning Basin:',
    disabled=False,
    style = {'description_width': 'initial'}
)

# Define a function to update the global FINE_TUNING_BASIN and plot
# when a new basin is selected from the dropdown.
def select_fine_tuning_basin(basin_id_selection):
    global FINE_TUNING_BASIN
    FINE_TUNING_BASIN = basin_id_selection
    print(f"\nSelected Fine-Tuning Basin: {FINE_TUNING_BASIN}")

    # Plot the selected fine-tuning basin using the shapefile function.
    # This visualization helps confirm that the correct basin has been selected
    # and allows for a geographical context of the fine-tuning target.
    # `backend.plot_train_test_shapefile` is a utility function that draws basin boundaries
    # on a map, highlighting specific basins.
    backend.plot_train_test_shapefile(
        shapefile_path=SHAPEFILE_PATH,         # Path to the shapefile containing basin geometries.
        train_basin_ids=set(),                # An empty set because we are not highlighting training basins here.
        test_basin_ids={FINE_TUNING_BASIN},   # A set containing only the chosen fine-tuning basin ID to highlight it.
        model_name=f"Fine-Tuning Basin: {FINE_TUNING_BASIN}" # Updates the plot title to indicate the fine-tuning basin.
    )

# Create an interactive widget linking the dropdown to the selection function.
# This displays the dropdown and automatically calls `select_fine_tuning_basin`
# whenever the dropdown value changes.
interactive_basin_selection = interactive(
    select_fine_tuning_basin,
    basin_id_selection=basin_selector_dropdown,
)

# Display the interactive widget to the user.
display(interactive_basin_selection)

# Initialize FINE_TUNING_BASIN with the default or first value initially.
# This ensures FINE_TUNING_BASIN is set even before the user interacts with the dropdown.
if 'FINE_TUNING_BASIN' not in globals():
    FINE_TUNING_BASIN = basin_selector_dropdown.value
    print(f"Initial Fine-Tuning Basin: {FINE_TUNING_BASIN}")
```

--------------------------------

### Generate Fine-Tuning Run Commands

Source: https://openhydronet.readthedocs.io/en/latest/_sources/tutorial/OpenHydroNet_Tutorial.ipynb.txt

Generates the necessary commands for fine-tuning training and inference. It assumes the fine-tuning basin ID and the generated configuration file path are available. The output includes instructions for running these commands in a terminal.

```python
# --- Fine-Tuning: Generate Run Command ---

# This cell generates the commands needed to run the fine-tuning experiment
# and subsequent inference for the selected basin.

# Assumes FINE_TUNING_BASIN and finetune_basin_config_file are available globally.

# Construct the fine-tuning training command.
# This points to the basin-specific config file we just generated.
finetune_train_command = f"run finetune --config-file={finetune_basin_config_file}"

# Construct the fine-tuning inference command.
# The run directory will be inside the base model directory, starting with the experiment name.
# Note: The actual directory will have a timestamp suffix.
finetune_experiment_name = f"finetune-{FINE_TUNING_BASIN}"
finetune_infer_command_display = f"run infer --run-dir={model_run_dir}/{finetune_experiment_name}_YYYYMMDD_HHMMSS"


# Print the run instructions to the screen
print("--------------------------------------------------------------------------------------------------")
print("                       FINE-TUNING AND INFERENCE RUN INSTRUCTIONS                           ")
print("--------------------------------------------------------------------------------------------------")
print("These commands need to be executed in your terminal or a separate process.")
print("Ensure you are in an environment with the necessary Python libraries (e.g., the Conda environment).")
print("")
print("IMPORTANT: The configuration file has been automatically generated for your selected basin:")
print(f"  Config File Location: {finetune_basin_config_file}")
print("  This file contains the specific settings for fine-tuning on basin " + str(FINE_TUNING_BASIN))
print("  It has been pre-configured with the correct base model paths and basin list.")
print("")
print("STEP 1: Run the fine-tuning training command:")
print(f"    {finetune_train_command}")
print("")
print("STEP 2: After training is complete, find the created run directory and run inference.")
print(f"    The output directory will be created inside: {model_run_dir}")
print(f"    It will be named like: {finetune_experiment_name}_YYYYMMDD_HHMMSS")
print("")
print("    Run the inference command using that directory:")
print(f"    {finetune_infer_command_display}")
print("")
print("--------------------------------------------------------------------------------------------------")
```

--------------------------------

### Configure Model Selection

Source: https://openhydronet.readthedocs.io/en/latest/_sources/tutorial/OpenHydroNet_Tutorial.ipynb.txt

Sets the model run directory and creates a display name based on the user's selection from available options.

```python
global model_run_dir, model_name

    # Get the absolute path for the selected run directory from the pre-populated dictionary
    model_run_dir = available_run_dirs.get(model_selection)
    # Create a display name for the selected model, including its selection identifier
    model_name = f'Model ({model_selection})'

    print("\nConfiguration set:")
    print(f"  Model: {model_name} ({model_run_dir})")
```

--------------------------------

### Import Standard Libraries and Dependencies

Source: https://openhydronet.readthedocs.io/en/latest/tutorial/OpenHydroNet_Tutorial.html

Imports essential Python libraries for data handling, scientific computing, visualization, and interactive widgets.

```python
# Standard Library Imports
import os
import glob
import yaml
from typing import Dict, List, Optional, Set

# Scientific Computing & Data Analysis
import numpy as np
import pandas as pd
import xarray as xr

# Data Visualization
import matplotlib.pyplot as plt
import seaborn as sns

# Interactive Widgets for Notebook Environments
import ipywidgets as widgets
from ipywidgets import HBox, VBox, interactive

# Local Tutorial Module
import backend

```

--------------------------------

### Create Conda Environment

Source: https://openhydronet.readthedocs.io/en/latest/_sources/usage/quickstart.rst.txt

Create a Conda environment from the provided YAML file to manage dependencies. Ensure you are in the root of the cloned repository.

```bash
# Create the environment from the file in the repo
conda env create -f environments/conda.yml

# Activate the environment (MANDATORY)
conda activate googlehydrology
```

--------------------------------

### Select Fine-Tuning Basin with Dropdown

Source: https://openhydronet.readthedocs.io/en/latest/_sources/tutorial/OpenHydroNet_Tutorial.ipynb.txt

Creates an interactive dropdown to select a basin for fine-tuning. It updates a global variable and plots the selected basin's shapefile.

```Python
DEFAULT_FINE_TUNING_BASIN_ID = 'camels_13235000' # Example for small model run in live demo.

fine_tuning_basin_options = sorted(list(test_basin_ids))

basin_selector_dropdown = widgets.Dropdown(
    options=fine_tuning_basin_options,
    value=DEFAULT_FINE_TUNING_BASIN_ID if DEFAULT_FINE_TUNING_BASIN_ID in fine_tuning_basin_options else fine_tuning_basin_options[0],
    description='Select Fine-Tuning Basin:',
    disabled=False,
    style = {'description_width': 'initial'}
)

def select_fine_tuning_basin(basin_id_selection):
    global FINE_TUNING_BASIN
    FINE_TUNING_BASIN = basin_id_selection
    print(f"\nSelected Fine-Tuning Basin: {FINE_TUNING_BASIN}")

    backend.plot_train_test_shapefile(
        shapefile_path=SHAPEFILE_PATH,         
        train_basin_ids=set(),                
        test_basin_ids={FINE_TUNING_BASIN},   
        model_name=f"Fine-Tuning Basin: {FINE_TUNING_BASIN}" 
    )

interactive_basin_selection = interactive(
    select_fine_tuning_basin,
    basin_id_selection=basin_selector_dropdown,
)

display(interactive_basin_selection)

if 'FINE_TUNING_BASIN' not in globals():
    FINE_TUNING_BASIN = basin_selector_dropdown.value
    print(f"Initial Fine-Tuning Basin: {FINE_TUNING_BASIN}")
```

--------------------------------

### Execute Fine-Tuning and Inference

Source: https://openhydronet.readthedocs.io/en/latest/tutorial/OpenHydroNet_Tutorial.html

Programmatically executes the fine-tuning process and then runs inference on the newly fine-tuned model. It dynamically locates the most recent fine-tune run directory.

```python
if RUN_NEW_MODELS:
    # 1. Execute Targeted Fine-Tuning:
    # This calls 'run finetune' using the basin-specific YAML config.
    # It freezes most of the model (LSTMs) and only updates the 'static_attributes_fc'
    # and 'head' layers to better represent the unique characteristics of this basin.
    print(f"Executing fine-tuning: {finetune_train_command}")
    os.system(finetune_train_command)

    # Load the fine-tune config file to get the actual experiment name
    # We use the dynamically generated config file path from the previous steps
    with open(finetune_basin_config_file, 'r') as f:
        finetune_config_data = yaml.safe_load(f)
    actual_finetune_experiment_name = finetune_config_data.get('experiment_name', 'default_finetune_experiment')
    print(f"Finetune experiment name from config file: {actual_finetune_experiment_name}")

    # Get all potential run directories for the fine-tune experiment
    # These are expected to be within the base model's run directory structure.
    finetune_search_pattern = os.path.join(model_run_dir, f'{actual_finetune_experiment_name}_*_*/')
    all_finetune_run_dirs = glob.glob(finetune_search_pattern)

    if not all_finetune_run_dirs:
        print(f"Error: No fine-tune run directories found matching pattern {finetune_search_pattern}")
    else:
        all_finetune_run_dirs.sort()
        finetune_model_actual_run_dir = all_finetune_run_dirs[-1].rstrip('/')
        print(f"Found most recent fine-tune model run directory: {finetune_model_actual_run_dir}")

        # 2. Execute Inference on Fine-Tuned Model:
        # Once fine-tuning is complete, this command generates new predictions
        # using the updated weights.
        finetune_infer_command_actual = f"run infer --run-dir={finetune_model_actual_run_dir}"
        print(f"Executing fine-tuned inference: {finetune_infer_command_actual}")
        os.system(finetune_infer_command_actual)
else:
    print("Skipping fresh fine-tuning. Loading existing fine-tuned artifacts from the run directory.")
```

--------------------------------

### Load Model Configuration and Basin Lists

Source: https://openhydronet.readthedocs.io/en/latest/tutorial/OpenHydroNet_Tutorial.html

Loads the configuration file and extracts training and testing basin IDs for the selected base model. This function reads the model's configuration to determine which basins were used for training and testing.

```python
# Load train and test basin lists
# `backend.load_model_config_and_basins` reads the configuration file of the selected base model
# (from the `model_run_dir`) to extract which basins were used for training and testing.
# It returns the configuration dictionary, and sets of training and testing basin IDs.
base_config, train_basin_ids, test_basin_ids = backend.load_model_config_and_basins(model_run_dir)
```

--------------------------------

### Load Model Configuration and Basins for Fine-Tuning

Source: https://openhydronet.readthedocs.io/en/latest/_sources/tutorial/OpenHydroNet_Tutorial.ipynb.txt

Loads the configuration and test basin IDs for a selected fine-tuned model. This function reads the configuration file associated with the specified run directory.

```python
# --- Fine-Tuning: Calculate Metrics ---

# This cell is responsible for loading the configuration and test basin IDs for the
# selected fine-tuned model, and then either calculating or loading its performance metrics.
# It mirrors the process used for the base model but focuses on the fine-tuned experiment.

# Load configuration and basin IDs for the fine-tuned model.
# `backend.load_model_config_and_basins` reads the configuration file associated with
# the `finetune_run_dir` (which was set in the cell above).
# It extracts the experiment's configuration details, as well as the training and testing
# basin IDs used for this specific fine-tuned run.
finetune_config, finetune_train_basin_ids, finetune_test_basin_ids = backend.load_model_config_and_basins(finetune_run_dir)

# Load simulation data and calculate/load metrics for the fine-tuned model.
# `backend.load_data_and_metrics` performs the heavy lifting of reading raw model outputs
#
```

--------------------------------

### Evaluate Model Performance

Source: https://openhydronet.readthedocs.io/en/latest/_sources/usage/quickstart.rst.txt

Calculate performance metrics on the test set after training. Specify the directory where the model run was saved.

```bash
run evaluate --run-dir /path/to/your/model_run/
```

--------------------------------

### Calculate Metrics for Fine-Tuned Model

Source: https://openhydronet.readthedocs.io/en/latest/tutorial/OpenHydroNet_Tutorial.html

Loads configuration and test basin IDs for the fine-tuned model, then calculates or loads its performance metrics. It uses a flag to determine whether to recalculate statistics or load pre-saved metrics.

```python
# --- Fine-Tuning: Calculate Metrics ---

# This cell is responsible for loading the configuration and test basin IDs for the
# selected fine-tuned model, and then either calculating or loading its performance metrics.
# It mirrors the process used for the base model but focuses on the fine-tuned experiment.

# Load configuration and basin IDs for the fine-tuned model.
# `backend.load_model_config_and_basins` reads the configuration file associated with
# the `finetune_run_dir` (which was set in the cell above).
# It extracts the experiment's configuration details, as well as the training and testing
# basin IDs used for this specific fine-tuned run.
finetune_config, finetune_train_basin_ids, finetune_test_basin_ids = backend.load_model_config_and_basins(finetune_run_dir)

# Load simulation data and calculate/load metrics for the fine-tuned model.
# `backend.load_data_and_metrics` performs the heavy lifting of reading raw model outputs
# and, based on the `CALCULATE_STATISTICS` flag, computes or loads performance metrics.
# - `model_run_dir`: The directory of the fine-tuned model run selected in the previous step.
# - `finetune_test_basin_ids`: The set of basin IDs that were used for testing this fine-tuned model.
# - `CALCULATE_STATISTICS`: A global boolean flag. If `True`, metrics are recalculated;
#   if `False`, pre-saved metrics are loaded to save computation time.
# - `model_name`: A descriptive name for the fine-tuned model, used for display purposes.
finetune_data, finetune_metrics = backend.load_data_and_metrics(
    model_run_dir=finetune_run_dir,
    test_basin_ids=finetune_test_basin_ids,
    calculate_statistics=CALCULATE_STATISTICS,
    model_name=model_name
)

# Display the head of the calculated/loaded fine-tuned model metrics DataFrame.
# This provides a quick overview of the metrics, including basin IDs, lead times, and scores.
display(finetune_metrics.head())
```

--------------------------------

### Load Fine-Tuned Model Data and Metrics

Source: https://openhydronet.readthedocs.io/en/latest/_sources/tutorial/OpenHydroNet_Tutorial.ipynb.txt

Loads data and performance metrics for a fine-tuned model. It can either recalculate metrics or load pre-saved ones based on the CALCULATE_STATISTICS flag.

```python
# - `model_run_dir`: The directory of the fine-tuned model run selected in the previous step.
# - `finetune_test_basin_ids`: The set of basin IDs that were used for testing this fine-tuned model.
# - `CALCULATE_STATISTICS`: A global boolean flag. If `True`, metrics are recalculated;
#   if `False`, pre-saved metrics are loaded to save computation time.
# - `model_name`: A descriptive name for the fine-tuned model, used for display purposes.
finetune_data, finetune_metrics = backend.load_data_and_metrics(
    model_run_dir=finetune_run_dir,
    test_basin_ids=finetune_test_basin_ids,
    calculate_statistics=CALCULATE_STATISTICS,
    model_name=model_name
)

# Display the head of the calculated/loaded fine-tuned model metrics DataFrame.
# This provides a quick overview of the metrics, including basin IDs, lead times, and scores.
display(finetune_metrics.head())
```

--------------------------------

### Plot Base Model KGE vs. Lead Time for Selected Basin

Source: https://openhydronet.readthedocs.io/en/latest/tutorial/OpenHydroNet_Tutorial.html

Generates a bar chart comparing the base model's KGE scores against lead time for a specific fine-tuning basin.

```python
finetune_basin_kge_scores = base_model_metrics.loc[f'{FINE_TUNING_BASIN}', 'KGE']

# Create a simple bar chart
plt.figure(figsize=(10, 6))
plt.bar(finetune_basin_kge_scores.index, finetune_basin_kge_scores.values)

plt.title(f"Base Model KGE vs. Lead Time for Basin {FINE_TUNING_BASIN}")
plt.xlabel("Lead Time (days)")
plt.ylabel("KGE Score")
plt.grid(axis='y', linestyle='--', alpha=0.7)
plt.xticks(finetune_basin_kge_scores.index)
plt.show()
```

--------------------------------

### Import Python Libraries

Source: https://openhydronet.readthedocs.io/en/latest/_sources/tutorial/OpenHydroNet_Tutorial.ipynb.txt

Imports standard library modules, scientific computing packages, data visualization tools, and interactive widget libraries. Also imports a local 'backend' module.

```python
# Standard Library Imports
import os
import glob
import yaml
from typing import Dict, List, Optional, Set

# Scientific Computing & Data Analysis
import numpy as np
import pandas as pd
import xarray as xr

# Data Visualization
import matplotlib.pyplot as plt
import seaborn as sns

# Interactive Widgets for Notebook Environments
import ipywidgets as widgets
from ipywidgets import HBox, VBox, interactive

# Local Tutorial Module
import backend
```

--------------------------------

### Generate Base Model Run Commands

Source: https://openhydronet.readthedocs.io/en/latest/tutorial/OpenHydroNet_Tutorial.html

Constructs shell commands for training and running inference on the base model. These commands are intended to be executed in a separate terminal or process.

```python
# --- Base Model: Generate Run Commands ---

# This cell generates the commands needed to train and run inference for the base model.
# These commands are typically executed outside of this notebook in a terminal or a separate process
# that has access to the `flood-forecasting/googlehydrology` project and its
# required Python environment (e.g., a Conda environment).

# Get the experiment name. We are assuming '5-basin-example' as the base experiment
# because the base_config may not be loaded yet at this stage.
training_config_file = 'train-config.yml'
base_config_file = f"{CONFIG_DIR}/{training_config_file}"

# Construct the base model training command.
# This command initiates the training process for the base model, using its specific configuration file.
base_train_command = f"run train --config-file={base_config_file}"

# Construct the base model inference command.
# After a base model has been trained, this command is used to run predictions
# (inference) using the trained model. It specifies the `--run-dir`
# which points to the output directory of the base model training run.
# NOTE: At this stage, the actual run directory might not exist yet if you are training a new model.

```

--------------------------------

### Find Available Model Run Directories

Source: https://openhydronet.readthedocs.io/en/latest/tutorial/OpenHydroNet_Tutorial.html

This Python code finds available model run directories using a specified path. It is used to populate a selection widget for choosing a base model.

```python
# Find available run directories using the path set in the cell above
available_run_dirs = backend.find_model_run_dirs(MODEL_RUN_DIR)

# Extract the name of the most recent model run for default selection
if 'base_model_actual_run_dir' in globals():
    latest_run_name = os.path.basename(base_model_actual_run_dir)
else:
    latest_run_name = None
```

--------------------------------

### Load Data and Calculate Metrics

Source: https://openhydronet.readthedocs.io/en/latest/tutorial/OpenHydroNet_Tutorial.html

Loads simulation data and calculates or loads performance metrics for the base model. This function requires the model run directory and test basin IDs, and can be configured to recalculate statistics.

```python
# Load simulation data and calculate/load metrics for the base model.
# `backend.load_data_and_metrics` handles reading the raw model outputs
# and, if specified, computes performance metrics across different lead times.
# It requires the `model_run_dir` to locate the data and `test_basin_ids` to filter for relevant basins.
base_model_data, base_model_metrics = backend.load_data_and_metrics(
    model_run_dir=model_run_dir,
    test_basin_ids=test_basin_ids,
    # The 'calculate_statistics' parameter controls whether metrics are re-calculated
    # or loaded from a pre-saved file. Setting it to `True` (as done here temporarily)
    # forces recalculation, which is useful after a new model run or if previous
    # calculations are outdated. If set to `False` (referencing the global
    # CALCULATE_STATISTICS), it will load existing metrics if available, saving time.
    calculate_statistics=CALCULATE_STATISTICS,
    model_name=model_name
)

# `base_model_data` will contain the raw simulation outputs (e.g., streamflow_sim, streamflow_obs)
# typically as an xarray Dataset, allowing for detailed hydrograph analysis.
# `base_model_metrics` will be a pandas DataFrame containing calculated performance scores
# (e.g., KGE, NSE) for each basin and lead time, used for quantitative comparisons.
display(base_model_metrics.head())
```