Source code for bhf_smarthealth.utils.data_io

"""This module contains functions used for data input and output"""

import pickle
import json
import geopandas as gpd
import pandas as pd
import logging

from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)




def load_data_from_file(file_path: str | Path, **kwargs) -> Any:
    """Load data

    Parameters
    ----------
    file_path : str | Path
        Path to data file

    Returns
    -------
    Any
        Data object
    """
    file_path = Path(file_path) if isinstance(file_path, str) else file_path

    # Adjust the mode and function used to load the data
    # based on the file extension
    match file_path.suffix:
        case ".pkl":
            mode = "rb"
            load_func = pickle.load
        case ".csv":
            mode = "r"
            load_func = pd.read_csv
        case ".json":
            mode = "r"
            load_func = json.load
        case ".parquet":
            mode = "rb"
            load_func = pd.read_parquet
        case _:
            raise ValueError("File extension not recognised.")

    with open(
        file_path,
        mode,
    ) as file:
        data = load_func(file, **kwargs)

    logger.debug(f"Data successfully loaded from {file_path}")

    return data





def save_dataframe_ouput(
    df: gpd.GeoDataFrame | pd.DataFrame,
    file_out_path: str | Path,
    **kwargs: Any,
) -> None:
    """Save DataFrame to output file. File format determined by the file extension.

    Parameters
    ----------
    df : gpd.GeoDataFrame | pd.DataFrame
        DataFrame containing data to be exported
    file_out_path : str | Path
        Path to output file to

    Returns
    -------
    None
    """
    file_out_path = (
        Path(file_out_path) if isinstance(file_out_path, str) else file_out_path
    )

    # Adjust the mode of file output based on the file extension
    match file_out_path.suffix:
        case ".parquet":
            df.to_parquet(
                file_out_path,
                **kwargs,
            )
            logger.debug(f"File saved to {file_out_path}")
        case ".csv":
            df.to_csv(
                file_out_path,
                **kwargs,
            )
            logger.debug(f"File saved to {file_out_path}")
        case _:
            raise ValueError("File format not recognised.")