functions.py

# Import all the necessary libraries
import os
import folium
import json
import requests
from io import StringIO
import pandas as pd
from datetime import datetime
from pathlib import Path
import numpy as np
import geopandas as gpd
from datetime import datetime, timedelta
import calendar
from shapely.geometry import shape, LineString, MultiLineString


# Get the user defined SWORD flowline
def clip_flowline(SWORD, AOI, output_path):
    """
    Clips a river flowline dataset to a user-provided shapefile extent,
    ensuring both datasets have the same CRS.

    Parameters:
    flowline_path (str): Path to the river flowline GPKG file.
    user_shapefile_path (str): Path to the user-provided shapefile (shp, gpkg, geojson, etc.).
    output_path (str): Path to save the clipped flowline.
    """
    flowline_gdf = gpd.read_file(SWORD)
    AOI_gdf = gpd.read_file(AOI)

    if flowline_gdf.crs != AOI_gdf.crs:
        flowline_gdf = flowline_gdf.to_crs(AOI_gdf.crs)

    clipped_flowline = gpd.clip(flowline_gdf, AOI_gdf)

    os.makedirs(output_path, exist_ok=True)
    output_file = os.path.join(output_path, "desired_flowline.gpkg")

    clipped_flowline.to_file(output_file, driver="GPKG")

    print(f"Desired flowline saved at: {output_path}")
    return clipped_flowline


# Get the reach_id based SWOT data
def fetch_hydrocron(feature_id: str, start_date: str, end_date: str, fields: list):
    """
    Fetch Hydrocron timeseries data for a given feature_id and date range.

    Parameters:
    - feature_id (str): The feature ID to query.
    - start_date (str): Start date in flexible format (YYYY, YYYY-MM, YYYY-MM-DD).
    - end_date (str): End date in flexible format (YYYY, YYYY-MM, YYYY-MM-DD).
    - fields (list): List of fields to retrieve.

    Returns:
    - pd.DataFrame: Dataframe containing the retrieved Hydrocron data.
    """
    def format_date(date_str, is_end_date=False):
        try:
            date_obj = datetime.strptime(date_str, "%Y-%m-%d")
        except ValueError:
            try:
                date_obj = datetime.strptime(date_str, "%Y-%m")
                last_day = calendar.monthrange(date_obj.year, date_obj.month)[1]  # Get the last valid day
                date_obj = date_obj.replace(day=last_day if is_end_date else 1)
            except ValueError:
                try:
                    date_obj = datetime.strptime(date_str, "%Y")
                    date_obj = date_obj.replace(month=12 if is_end_date else 1, day=31 if is_end_date else 1)
                except ValueError:
                    raise ValueError("Invalid date format. Use YYYY, YYYY-MM, or YYYY-MM-DD.")
        return date_obj.strftime("%Y-%m-%d")


    # Format dates
    start_date = format_date(start_date)
    end_date = format_date(end_date, is_end_date=True)

    # Construct request URL
    url = "https://soto.podaac.earthdatacloud.nasa.gov/hydrocron/v1/timeseries"
    params = {
        "feature": "Reach",
        "feature_id": feature_id,
        "start_time": f"{start_date}T00:00:00Z",
        "end_time": f"{end_date}T00:00:00Z",
        "output": "csv",
        "fields": ",".join(fields),
    }

    # Make request
    response = requests.get(url, params=params)

    if response.status_code == 200:
        json_data = response.json()
        if "results" in json_data and "csv" in json_data["results"]:
            df = pd.read_csv(StringIO(json_data["results"]["csv"]))
            return df
        else:
            raise Exception(
                "Unexpected response format: Missing 'results' or 'csv' key."
            )
    else:
        raise Exception(f"Error {response.status_code}: {response.text}")


#For the interactive map to display the flowlines
def display_flowlines(geojson_data):
    if isinstance(geojson_data, str):
        geojson_data = json.loads(geojson_data)

    coords = []
    available_fields = set()

    for feature in geojson_data["features"]:
        if "properties" in feature:
            available_fields.update(feature["properties"].keys())
        
            for key, value in feature["properties"].items():
                if isinstance(value, float):
                    feature["properties"][key] = f"{value:.10f}"
                    
        if "geometry" in feature:
            geom = shape(feature["geometry"])

            if geom.geom_type == "MultiLineString":
                for line in geom.geoms:
                    coords.extend(line.coords)
                feature["geometry"] = geom.geoms[0].__geo_interface__  # Corrected

            elif geom.geom_type == "LineString":
                coords.extend(geom.coords)

    if not coords:
        raise ValueError("No valid LineString geometries found in GeoJSON data.")

    lons, lats = zip(*coords)
    min_lon, max_lon = min(lons), max(lons)
    min_lat, max_lat = min(lats), max(lats)
    center_lat = (min_lat + max_lat) / 2
    center_lon = (min_lon + max_lon) / 2

    m = folium.Map(location=[center_lat, center_lon], zoom_start=12, tiles="cartodbpositron")

    desired_fields = ["reach_id", "river_name", "slope", "width", "min_slope[SWOT]", "max_slope[SWOT]", 
                      "median_slope[SWOT]", "mean_slope[SWOT]", "num_records_filtered"]
    desired_aliases = ["Reach ID", "River Name", "Slope", "Width (m)", "Minimum Slope[SWOT]", "Maximum Slope[SWOT]", 
                       "Median Slope[SWOT]", "Mean Slope[SWOT]", "Number of Records"]

    valid_fields = [field for field in desired_fields if field in available_fields]
    valid_aliases = [alias for field, alias in zip(desired_fields, desired_aliases) if field in available_fields]

    folium.GeoJson(
        geojson_data, 
        name="Flowline Data",
        tooltip=folium.GeoJsonTooltip(fields=valid_fields, aliases=valid_aliases),
        popup=folium.GeoJsonPopup(fields=valid_fields, aliases=valid_aliases)
    ).add_to(m)

    m.fit_bounds([[min_lat, min_lon], [max_lat, max_lon]], padding=(10, 10))
    folium.LayerControl().add_to(m)
    return m

# Function to aggregate the reach slope data to get min, max, median, and mean rows
def process_reach_slope(
    df,
    feature_id,
    slope_col="slope",
    slope_u_col="slope_u",
    time_col="time_str",
    reach_q_col="reach_q",
):
    """
    Processes reach slope data to compute min, max, median, mean slope,
    and the number of filtered records following the given conditions.

    Parameters:
    - df (DataFrame): Input dataset containing reach slope data.
    - feature_id (int/str): Reach ID.
    - slope_col (str): Column name for slope.
    - slope_u_col (str): Column name for slope uncertainty.
    - time_col (str): Column name for time (for exclusions).
    - reach_q_col (str): Column name for reach quality flag.

    Returns:
    - DataFrame with computed min, max, median, mean slope, and num_records_filtered.
    """

    # Remove invalid slope values
    df_reach = df[df[slope_col] != -999999999999.0]

    # Remove columns ending in "_units"
    df_reach = df_reach.loc[:, ~df_reach.columns.str.endswith("_units")]

    # Store original count
    num_records_original = len(df_reach)
    print(f"Reach {feature_id} - num_records_original: {num_records_original}")

    if num_records_original == 0:
        return pd.DataFrame(columns=["reach_id", "min_slope[SWOT]", "max_slope[SWOT]", 
                                     "median_slope[SWOT]", "mean_slope[SWOT]", "num_records_filtered"])

    # Apply filtering: Keep records within 95th percentile of slope uncertainty
    if num_records_original >= 2:
        df_reach = df_reach[df_reach[slope_u_col] <= df_reach[slope_u_col].quantile(0.95)]

    num_records_filtered = len(df_reach)
    print(f"Reach {feature_id} - num_records_filtered: {num_records_filtered}")

    if num_records_filtered >= 2:
        min_slope = df_reach[slope_col].min()
        max_slope = df_reach[slope_col].max()
        median_slope = df_reach[slope_col].median()
        mean_slope = df_reach[slope_col].mean()

    elif num_records_filtered == 1:
        if df_reach.iloc[0][reach_q_col] != 3: 
            min_slope = max_slope = median_slope = mean_slope = df_reach.iloc[0][slope_col]
        else:
            print(f"Reach {feature_id} - No valid observation in this period.")
            return pd.DataFrame(columns=["reach_id", "min_slope[SWOT]", "max_slope[SWOT]", 
                                         "median_slope[SWOT]", "mean_slope[SWOT]", "num_records_filtered"])

    else:
        print(f"Reach {feature_id} - Less than 2 valid records after filtering.")
        return pd.DataFrame(columns=["reach_id", "min_slope[SWOT]", "max_slope[SWOT]", 
                                     "median_slope[SWOT]", "mean_slope[SWOT]", "num_records_filtered"])

    # Return results in DataFrame format
    return pd.DataFrame({
        "reach_id": [feature_id],
        "min_slope[SWOT]": [min_slope],
        "max_slope[SWOT]": [max_slope],
        "median_slope[SWOT]": [median_slope],
        "mean_slope[SWOT]": [mean_slope],
        "num_records_filtered": [num_records_filtered]
    })
    
#Calculate the aggregated slope
def getSWOTaggregatedslope(AOI_flowlines, output_gpkg_path, start_date, end_date, fields):
    """
    Processes reach slope data for AOI flowlines, computes slope statistics, 
    merges them into the dataset, and saves the result as a GeoPackage.

    Parameters:
    - AOI_flowlines (GeoDataFrame): The input AOI flowline dataset.
    - output_gpkg_path (str): Path to save the output GeoPackage.
    - start_date (str): Start date for fetching hydro data.
    - end_date (str): End date for fetching hydro data.
    - fields (list): Fields to request from fetch_hydrocron.

    Returns:
    - None (Saves the processed file as a .gpkg)
    """
    unique_reach_ids = AOI_flowlines['reach_id'].unique()
    total_reaches = len(unique_reach_ids)
    
    # Ensuring AOI_flowlines is a GeoDataFrame
    if not isinstance(AOI_flowlines, gpd.GeoDataFrame):
        AOI_flowlines = gpd.GeoDataFrame(AOI_flowlines)

    # Ensure CRS is set
    if AOI_flowlines.crs is None:
        AOI_flowlines.set_crs(epsg=4326, inplace=True)

    # Store processed results
    slope_results = []

    # Process each reach_id (limited to first 10 for testing)
    for i, reach_id in enumerate(AOI_flowlines["reach_id"].unique(), start=1):
        try:
            # Try fetching data
            df = fetch_hydrocron(reach_id, start_date, end_date, fields)

            # Unified check for missing data (None or empty DataFrame)
            if df is None or df.empty:
                print(f"Warning: No data found for reach_id {reach_id}, skipping...")
                continue
            
            # Process slope calculations
            result = process_reach_slope(df, reach_id)
            if not result.empty:
                slope_results.append(result)

        except Exception as e:
            continue  
            # Calculate progress percentage
        progress = (i / total_reaches) * 100

        # Print progress every 10%
        if i % (total_reaches // 10) == 0:  
            print(f"Progress: {progress:.0f}% ({i}/{total_reaches} reach_ids processed)")
            
    # Merge results back to AOI_flowlines
    if slope_results:
        slope_df = pd.concat(slope_results, ignore_index=True)
        AOI_flowlines = AOI_flowlines.merge(slope_df, on="reach_id", how="left")

    # Save the updated dataset to a GeoPackage (.gpkg) file
    AOI_flowlines.to_file(output_gpkg_path, driver="GPKG")
    print(f"Updated GeoPackage saved at: {output_gpkg_path}")


# Define create_nan_row as a separate function
def create_nan_row(df_template, feature_id):
    """
    Create a NaN row with the same structure as df_template.
    Keeps 'reach_id' and sets 'num_records_filtered' to 0.
    
    Parameters:
    - df_template: DataFrame to base the NaN row structure on
    - feature_id: Value to assign to 'reach_id'
    
    Returns:
    - NaN row DataFrame
    """
    nan_row = pd.DataFrame(columns=df_template.columns)
    nan_row.loc[0] = np.nan  # Fill all values with NaN
    nan_row["reach_id"] = feature_id  # Assign reach_id
    nan_row["num_records_filtered"] = 0  # Indicate no valid records
    
    return nan_row

# Function to aggregate the reach slope data to get min, max, median, and mean rows in dataframe
def process_reach_slope_df(df, feature_id, slope_col, slope_u_col, time_col, reach_q_col, df_min_slope, df_max_slope, df_median_slope, df_mean_slope):

    # global df_min_slope, df_max_slope, df_median_slope, df_mean_slope
    
    # Remove invalid slopes
    df_reach = df[df[slope_col] != -999999999999.0]
    
    # Remove columns ending in "_units"
    df_reach = df_reach.loc[:, ~df_reach.columns.str.endswith('_units')]
    
    # Store original count
    num_records_original = len(df_reach)
    # print(f"num_records_original: {num_records_original}")

    # If no records exist, exit function early
    if num_records_original == 0:
        df_min_slope = pd.concat([df_min_slope, create_nan_row(df_reach, feature_id)], ignore_index=True)
        df_max_slope = pd.concat([df_max_slope, create_nan_row(df_reach, feature_id)], ignore_index=True)
        df_median_slope = pd.concat([df_median_slope, create_nan_row(df_reach, feature_id)], ignore_index=True)
        df_mean_slope = pd.concat([df_mean_slope, create_nan_row(df_reach, feature_id)], ignore_index=True)
        return

    # Check if at least 2 records exist
    if num_records_original >= 2:
        df_reach = df_reach[df_reach[slope_u_col] <= df_reach[slope_u_col].quantile(0.95)]
    
        num_records_filtered = len(df_reach)

        # If still at least 2 records
        if num_records_filtered >= 2:

            # print(f"num_records_filtered: {num_records_filtered}")

            row_min_slope = df_reach[df_reach[slope_col] == df_reach[slope_col].min()].copy()
            row_max_slope = df_reach[df_reach[slope_col] == df_reach[slope_col].max()].copy()
            
            # For calculating median row
            if len(df_reach) % 2 == 1:  # Odd number of records
                median_row = df_reach.iloc[[len(df_reach) // 2]].copy()
            else:  # Even number of records
                middle_rows = df_reach.iloc[[len(df_reach) // 2 - 1, len(df_reach) // 2]].copy()
                numeric_columns = df_reach.columns.difference([time_col])  # Exclude time_col for median calculation
                median_values = middle_rows[numeric_columns].median().to_frame().T  # Compute median for numeric columns
                median_row = middle_rows.iloc[[0]].copy()  # Use first middle row as template
                median_row[numeric_columns] = median_values.values  # Assign median values to numeric columns
                
                # Extract the day from the time_col of both middle rows
                days = middle_rows[time_col].apply(lambda x: str(pd.to_datetime(x).day)).tolist()
                median_row[time_col] = f"day{days[0]}, day{days[1]}"  # Assign formatted day string

                # Ensure column order matches df_reach (optional)
                median_row = median_row[df_reach.columns]
            
            # For calculating mean row
            # Drop time column before calculating statistics
            df_numeric = df_reach.drop(columns=[time_col])
            
            # Compute the mean row
            mean_row = df_numeric.mean().to_frame().T

            # Add time_col as NaN to maintain structure consistency
            mean_row[time_col] = np.nan

            # Ensure column order matches df_reach (optional)
            mean_row = mean_row[df_reach.columns]

            # Add the num_records_filtered column
            row_min_slope["num_records_filtered"] = num_records_filtered
            row_max_slope["num_records_filtered"] = num_records_filtered
            median_row["num_records_filtered"] = num_records_filtered
            mean_row["num_records_filtered"] = num_records_filtered

        else:
            print("Less than 2 valid records after filtering, checking single-record case.")

    # If there is only **one** record left (either originally or after filtering)
    if len(df_reach) == 1:

        # print(f"len(df_reach): {len(df_reach)}")

        if df_reach.iloc[0][reach_q_col] != 3:

            # Get rows where slope is min, max, median and mean (for a single record, it's the same row)
            row_min_slope = row_max_slope = df_reach.copy()
            median_row = mean_row = df_reach.copy()

            # Add num_records_filtered column
            row_min_slope["num_records_filtered"] = 1
            row_max_slope["num_records_filtered"] = 1
            median_row["num_records_filtered"] = 1
            mean_row["num_records_filtered"] = 1

        else:
            print("No valid observation in this period.")

            df_min_slope = pd.concat([df_min_slope, create_nan_row(df_reach, feature_id)], ignore_index=True)
            df_max_slope = pd.concat([df_max_slope, create_nan_row(df_reach, feature_id)], ignore_index=True)
            df_median_slope = pd.concat([df_median_slope, create_nan_row(df_reach, feature_id)], ignore_index=True)
            df_mean_slope = pd.concat([df_mean_slope, create_nan_row(df_reach, feature_id)], ignore_index=True)
            return

    # Append results to respective DataFrames
    df_min_slope = pd.concat([df_min_slope, row_min_slope], ignore_index=True)
    df_max_slope = pd.concat([df_max_slope, row_max_slope], ignore_index=True)
    df_median_slope = pd.concat([df_median_slope, median_row], ignore_index=True)
    df_mean_slope = pd.concat([df_mean_slope, mean_row], ignore_index=True)

    return df_min_slope, df_max_slope, df_median_slope, df_mean_slope

# Calculate the aggregated slope for a specific month
def getSWOTaggregatedslopedf(AOI_flowlines, compiled_df, output_df_path, month):
    """
    Aggregates slope data from compiled monthly data and saves the results.

    Parameters:
    - AOI_flowlines (GeoDataFrame): The input AOI flowline dataset.
    - compiled_df (DataFrame): Combined hydro data for the month.
    - output_df_path (Path): Path to save the output CSV.
    - month (str): The month for which the data was aggregated.

    Returns:
    - None (Saves the processed file as CSVs)
    """

    slope_col = 'slope'
    slope_u_col = 'slope_u'  
    time_col = 'time_str'  
    reach_q_col = 'reach_q'  

    # # Ensure AOI_flowlines is a GeoDataFrame
    # if not isinstance(AOI_flowlines, gpd.GeoDataFrame):
    #     AOI_flowlines = gpd.GeoDataFrame(AOI_flowlines)

    # # Ensure CRS is set
    # if AOI_flowlines.crs is None:
    #     AOI_flowlines.set_crs(epsg=4326, inplace=True)

    # Get unique reach_ids
    unique_reach_ids = AOI_flowlines['reach_id'].unique()
    total_reaches = len(unique_reach_ids)

    # Initialize empty DataFrames for storing results
    df_min_slope, df_max_slope, df_median_slope, df_mean_slope = pd.DataFrame(), pd.DataFrame(), pd.DataFrame(), pd.DataFrame()

    for i, reach_id in enumerate(unique_reach_ids, start=1):
        try:
            df = compiled_df[compiled_df['reach_id'] == reach_id]

            if df.empty:
                print(f"Warning: No data for reach_id {reach_id}, skipping...")
                continue

            df_min_slope, df_max_slope, df_median_slope, df_mean_slope = process_reach_slope_df(
                df, reach_id, slope_col, slope_u_col, time_col, reach_q_col, 
                df_min_slope, df_max_slope, df_median_slope, df_mean_slope
            )

        except Exception as e:
            print(f"Error processing reach_id {reach_id}: {e}")
            continue

        if i % (total_reaches // 10) == 0:
            print(f"Progress: {(i / total_reaches) * 100:.0f}% ({i}/{total_reaches} reach_ids processed)")

    # Save results with month-specific filename
    output_df_path.mkdir(parents=True, exist_ok=True)
    df_min_slope.to_csv(output_df_path / f"min_slope_results_{month}.csv", index=False)
    df_max_slope.to_csv(output_df_path / f"max_slope_results_{month}.csv", index=False)
    df_median_slope.to_csv(output_df_path / f"median_slope_results_{month}.csv", index=False)
    df_mean_slope.to_csv(output_df_path / f"mean_slope_results_{month}.csv", index=False)