# AUTOGENERATED! DO NOT EDIT! File to edit: ../../notebooks/utils/31_lbl.ipynb.

# %% auto 0
__all__ = ['HTTP_PROTOCOLS', 'load_lbl', 'LblDataset']

# %% ../../notebooks/utils/31_lbl.ipynb 2
from pathlib import PurePosixPath
import fsspec

from kedro.io import AbstractDataset
from kedro.io.core import get_filepath_str, get_protocol_and_path

import pandas
import pdr
import pooch

from typing import Any, Dict

# %% ../../notebooks/utils/31_lbl.ipynb 3
HTTP_PROTOCOLS = ("http", "https")


def load_lbl(filepath: str, type: str = "table") -> pandas.DataFrame:
    """Load LBL data.

    Args:
        filepath: File path to load the data from.
        type: Type of data to load. Options are 'table' and 'index'.

    Returns:
        A pandas DataFrame containing the loaded data.
    """
    if type == "table":
        df = pdr.read(filepath).TABLE
    elif type == "index":
        df = pandas.read_csv(filepath, delimiter=",", quotechar='"')
        df.columns = df.columns.str.replace(" ", "")

    return df


class LblDataset(AbstractDataset):
    def __init__(
        self,
        filepath: str,
        load_type: str = "table",
        metadata: Dict[str, Any] = None,
    ):
        # parse the path and protocol (e.g. file, http, s3, etc.)
        protocol, path = get_protocol_and_path(filepath)
        self._protocol = protocol
        self._filepath = PurePosixPath(path)

        self._fs = fsspec.filesystem(self._protocol)

        self.load_type = load_type
        self.metadata = metadata

    def _load(self):
        # using get_filepath_str ensures that the protocol and path are appended correctly for different filesystems
        load_path = get_filepath_str(self._filepath, self._protocol)

        if self._protocol in HTTP_PROTOCOLS:
            local_fp = pooch.retrieve(load_path, known_hash=None)
        else:
            local_fp = load_path

        return load_lbl(local_fp, self.load_type)

    def _save(self):
        pass

    def _describe(self):
        """Returns a dict that describes the attributes of the dataset."""
        return dict(filepath=self._filepath, protocol=self._protocol)
