"""National Oceanic And Atmospheric Administration (NOAA) client."""

import logging as lg
from collections.abc import Sequence

import dask
import pandas as pd
import pooch
import requests
from dask import diagnostics
from pyregeon import RegionType

from meteora import settings, utils
from meteora.clients.base import BaseFileClient
from meteora.clients.mixins import (
    StationPartitionedTSMixin,
    StationsEndpointMixin,
    TimePartitionedTSMixin,
    VariablesHardcodedMixin,
)
from meteora.utils import DateTimeType, KwargsType, VariablesType

# disable pooch warnings when providing `None` as "known_hash"
logger = pooch.get_logger()
logger.setLevel("WARNING")


# API endpoints
BASE_URL = "https://www.ncei.noaa.gov/oa/global-historical-climatology-network/hourly"
GHCNH_STATIONS_ENDPOINT = f"{BASE_URL}/doc/ghcnh-station-list.csv"
TS_ENDPOINT = (
    f"{BASE_URL}/access/by-year/"
    + "{period.year}/psv/GHCNh_{station_id}_{period.year}.psv"
)

# for pooch
STATIONS_LIST_KNOWN_HASH = (
    "6b837d8d953aa6ebc172755864e549518944e7bec15b2196ee191219587a96f5"
)

# useful constants
STATIONS_GDF_ID_COL = "GHCN_ID"
TS_DF_STATIONS_ID_COL = "STATION"
TS_DF_TIME_COL = "DATE"
VARIABLES_ID_COL = "code"
VARIABLES_LABEL_COL = "description"

# see section "IV. List of elements/variable" and appendix A of the GHCNh documentation
# www.ncei.noaa.gov/oa/global-historical-climatology-network/hourly/doc/
# ghcnh_DOCUMENTATION.pdf
VARIABLES_DICT = {
    "temperature": "2 meter (circa) Above Ground Level Air (dry bulb) Temperature (⁰C n"
    "to tenths)",
    "relative_humidity": "Depending on the source, relative humidity is either measured"
    " directly or calculated from air (dry bulb) temperature and dew point temperature "
    "(whole percent)",
    "station_level_pressure": "The pressure that is observed at a specific elevation "
    "and is the true barometric pressure of a location. It is the pressure exerted by "
    "the atmosphere at a point as a result of gravity acting upon the 'column' of air "
    "that lies directly above the point. (hPa)",
    "precipitation": "Total liquid precipitation (rain or melted snow). Totals are "
    "nominally for the hour, but may include intermediate reports within the hour. "
    "Please refer to Appendix B for important details on precipitation totals; a `T` in"
    " the measurement code column indicates a trace amount of precipitation "
    "(millimeters).",
    "wind_speed": "Wind speed (meters per second)",
    "wind_direction": "Wind Direction from true north using compass directions (e.g. "
    "360 = true north, 180 = south, 270 = west, etc.). Note: A direction of `000` is "
    "given for calm winds. (whole degrees)",
}
ECV_DICT = {
    # precipitation
    settings.ECV_PRECIPITATION: "precipitation",
    # pressure
    settings.ECV_PRESSURE: "station_level_pressure",
    # temperature
    settings.ECV_TEMPERATURE: "temperature",
    # water vapour
    settings.ECV_RELATIVE_HUMIDITY: "relative_humidity",
    # wind
    settings.ECV_WIND_SPEED: "wind_speed",
    settings.ECV_WIND_DIRECTION: "wind_direction",
}


class GHCNHourlyClient(
    StationPartitionedTSMixin,
    TimePartitionedTSMixin,
    StationsEndpointMixin,
    VariablesHardcodedMixin,
    BaseFileClient,
):
    """NOAA GHCN hourly client.

    Parameters
    ----------
    region : str, Sequence, GeoSeries, GeoDataFrame, PathLike, or IO
        The region to process. This can be either:

        -  A string with a place name (Nominatim query) to geocode.
        -  A sequence with the west, south, east and north bounds.
        -  A geometric object, e.g., shapely geometry, or a sequence of geometric
           objects. In such a case, the value will be passed as the `data` argument of
           the GeoSeries constructor, and needs to be in the same CRS as the one used by
           the client's class (i.e., the `CRS` class attribute).
        -  A geopandas geo-series or geo-data frame.
        -  A filename or URL, a file-like object opened in binary ('rb') mode, or a Path
           object that will be passed to `geopandas.read_file`.
    pooch_kwargs : dict, optional
        Keyword arguments to pass to the `pooch.retrieve` function when downloading the
        stations time series data.
    progress : bool, optional
        Whether to show a tqdm progress bar for partitioned time series fetches. If
        None, the value from `settings.SHOW_PROGRESS` is used.
    sjoin_kwargs : dict, optional
        Keyword arguments to pass to the `geopandas.sjoin` function when filtering the
        stations within the region. If None, the value from `settings.SJOIN_KWARGS` is
        used.
    """

    # geom constants
    X_COL = "LONGITUDE"
    Y_COL = "LATITUDE"
    CRS = utils.LONLAT_CRS

    # time partition frequency
    _time_partition_freq = "YS"

    # API endpoints
    _stations_endpoint = GHCNH_STATIONS_ENDPOINT
    _ts_endpoint = TS_ENDPOINT
    _stations_known_hash = STATIONS_LIST_KNOWN_HASH

    # data frame labels constants
    _stations_gdf_id_col = STATIONS_GDF_ID_COL
    _ts_df_stations_id_col = TS_DF_STATIONS_ID_COL
    _ts_df_time_col = TS_DF_TIME_COL
    _variables_id_col = VARIABLES_ID_COL
    _variables_label_col = VARIABLES_LABEL_COL
    _variables_dict = VARIABLES_DICT
    _ecv_dict = ECV_DICT

    def __init__(
        self,
        region: RegionType,
        *,
        pooch_kwargs: KwargsType | None = None,
        progress: bool | None = None,
        **sjoin_kwargs: KwargsType,
    ) -> None:
        """Initialize GHCN hourly client."""
        self.region = region
        if not sjoin_kwargs:
            sjoin_kwargs = settings.SJOIN_KWARGS.copy()
        self.SJOIN_KWARGS = sjoin_kwargs

        if not pooch_kwargs:
            pooch_kwargs = {}
        self.pooch_kwargs = pooch_kwargs

        # need to call super().__init__() to set the cache
        super().__init__()
        if progress is not None:
            self.progress = progress

    def _ts_params(
        self, variable_ids: Sequence, start: DateTimeType, end: DateTimeType
    ) -> dict:
        start = pd.Timestamp(start)
        end = pd.Timestamp(end)
        return dict(variable_ids=variable_ids, start=start, end=end)

    def _ts_cache(self, ts_params) -> bool:
        return ts_params["period"].year != pd.Timestamp.now().year

    def _ts_df_from_url(self, url, ts_params) -> pd.DataFrame:
        variable_cols = list(ts_params["variable_ids"])
        cols_to_keep = (
            [self._ts_df_stations_id_col] + [self._ts_df_time_col] + variable_cols
        )
        start = ts_params["start"]
        end = ts_params["end"]
        try:
            source = self._ts_source(url, ts_params)
        except requests.HTTPError:
            return pd.DataFrame()
        ts_df = pd.read_csv(source, sep="|", usecols=cols_to_keep)
        ts_df[self._ts_df_time_col] = pd.to_datetime(ts_df[self._ts_df_time_col])
        ts_df = ts_df[ts_df[self._ts_df_time_col].between(start, end)]
        return ts_df.set_index([self._ts_df_stations_id_col, self._ts_df_time_col])

    def _ts_df_from_endpoint(self, ts_params) -> pd.DataFrame:
        # override to parallelize all (station, year) combinations with dask
        variable_cols = list(ts_params["variable_ids"])
        station_partitions = self._iter_station_partitions(ts_params)
        time_partitions = self._iter_time_partitions(ts_params)

        tasks = [
            dask.delayed(self._ts_df_from_url)(
                self._format_ts_endpoint(ts_params | sp | tp),
                ts_params | sp | tp,
            )
            for sp in station_partitions
            for tp in time_partitions
        ]
        with diagnostics.ProgressBar():
            dfs = dask.compute(*tasks)

        non_empty = [df for df in dfs if not df.empty]
        if not non_empty:
            utils.log(
                "No data found for any station in the requested period.",
                level=lg.WARNING,
            )
            return pd.DataFrame(columns=variable_cols)
        return pd.concat(non_empty)

    def get_ts_df(
        self,
        variables: VariablesType,
        start: DateTimeType,
        end: DateTimeType,
    ) -> pd.DataFrame:
        """Get time series data frame.

        Parameters
        ----------
        variables : str, int or list-like of str or int
            Target variables, which can be either an GHCNh variable code (string) or an
            essential climate variable (ECV) following the Meteora nomenclature
            (string).
        start, end : datetime-like, str, int, float
            Values representing the start and end of the requested data period
            respectively. Accepts any datetime-like object that can be passed to
            pandas.Timestamp.

        Returns
        -------
        ts_df : pandas.DataFrame
            Long form data frame with a time series of measurements (second-level index)
            at each station (first-level index) for each variable (column).
        """
        return self._get_ts_df(variables, start, end)
