# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/modules/99_datasets.ipynb.

# %% ../nbs/modules/99_datasets.ipynb #f34c6c2b
from __future__ import annotations
from pathlib import Path
import pandas as pd
import peshbeen
from typing import Tuple
import numpy as np

# ─────────────────────────────────────────────────────────────────────────────
# Helpers
# ─────────────────────────────────────────────────────────────────────────────

def _sm_dataset(name: str) -> pd.DataFrame:
    """Load a statsmodels built-in dataset as a DataFrame."""
    import statsmodels.api as sm
    return getattr(sm.datasets, name).load_pandas().data

# ─────────────────────────────────────────────────────────────────────────────
# UNIVARIATE DATASETS
# ─────────────────────────────────────────────────────────────────────────────

def admission_ts():
    """Daily hospital admissions in Wales, UK"""

    base_path = Path(peshbeen.__file__).parent.parent
    file_path = base_path / 'data' / 'wales_admissions.xlsx'
    
    # Optional: Fallback for initial local development before 'pip install -e .'
    if not file_path.exists():
        file_path = Path('..') / 'data' / 'wales_admissions.xlsx'
        
    if not file_path.exists():
        raise FileNotFoundError(f"Could not find dataset at: {file_path}")
        
    return pd.read_excel(file_path)

def calls_ts():
    """Daily emergency calls in Wales, UK"""

    base_path = Path(peshbeen.__file__).parent.parent
    file_path = base_path / 'data' / 'nhs_calls.xlsx'
    
    # Optional: Fallback for initial local development before 'pip install -e .'
    if not file_path.exists():
        file_path = Path('..') / 'data' / 'nhs_calls.xlsx'
        
    if not file_path.exists():
        raise FileNotFoundError(f"Could not find dataset at: {file_path}")
        
    return pd.read_excel(file_path)

def admission_calls_ts() -> pd.DataFrame:
    """Combined dataset of daily hospital admissions and emergency calls in Wales, UK"""
    admissions = admission_ts().set_index('Date')
    calls = calls_ts().set_index('Date')
    return admissions.join(calls, how='inner')


# airline passengers dataset
def airline_passengers() -> pd.DataFrame:
    """
    Monthly totals of international airline passengers (thousands) from
    January 1949 to December 1960 (144 observations).

    This is one of the most widely used example series in forecasting
    literature — it has a clear upward trend and multiplicative seasonality.

    Source: pandas built-in via a public URL, falling back to a hardcoded
    copy bundled from the classic Box & Jenkins dataset.

    Returns
    -------
    pd.DataFrame
        DatetimeIndex (monthly frequency), single column ``'passengers'``.

    Examples
    --------
    >>> from peshbeen.datasets import load_air_passengers
    >>> df = load_air_passengers()
    >>> df.head()
    """
    # Classic Box–Jenkins airline data (144 observations, 1949-01 to 1960-12)
    _passengers = [
        112, 118, 132, 129, 121, 135, 148, 148, 136, 119, 104, 118,
        115, 126, 141, 135, 125, 149, 170, 170, 158, 133, 114, 140,
        145, 150, 178, 163, 172, 178, 199, 199, 184, 162, 146, 166,
        171, 180, 193, 181, 183, 218, 230, 242, 209, 191, 172, 194,
        196, 196, 236, 235, 229, 243, 264, 272, 237, 211, 180, 201,
        204, 188, 235, 227, 234, 264, 302, 293, 259, 229, 203, 229,
        242, 233, 267, 269, 270, 315, 364, 347, 312, 274, 237, 278,
        284, 277, 317, 313, 318, 374, 413, 405, 355, 306, 271, 306,
        315, 301, 356, 348, 355, 422, 465, 467, 404, 347, 305, 336,
        340, 318, 362, 348, 363, 435, 491, 505, 404, 359, 310, 337,
        360, 342, 406, 396, 420, 472, 548, 559, 463, 407, 362, 405,
        417, 391, 419, 461, 472, 535, 622, 606, 508, 461, 390, 432,
    ]
    idx = pd.date_range(start='1949-01', periods=144, freq='MS')
    return pd.DataFrame({'passengers': _passengers}, index=idx)

# CO2 dataset
def load_co2_ts(
    interpolate: bool = True
) -> pd.DataFrame:
    """
    Weekly atmospheric CO2 concentrations (ppm) measured at Mauna Loa
    Observatory, Hawaii, from March 1958 to December 2001.

    Source: statsmodels built-in (sm.datasets.co2).

    Parameters
    ----------
    interpolate : bool
        If True (default), linearly interpolate the handful of missing weeks.

    Returns
    -------
    pd.DataFrame
        DatetimeIndex (weekly frequency), single column ``'co2'``.

    Examples
    --------
    >>> from peshbeen.datasets import load_co2
    >>> df = load_co2()
    >>> df.head()
    """
    import statsmodels.api as sm
    data = sm.datasets.co2.load_pandas().data
    data.index = pd.to_datetime(data.index)
    data.index.freq = pd.infer_freq(data.index)
    data.columns = ['co2']
    if interpolate:
        data['co2'] = data['co2'].interpolate(method='linear')
    data.index.name = 'date'
    return data

# sunspots dataset
def load_sunspots_ts() -> pd.DataFrame:
    """
    Annual sunspot activity from 1700 to 2008 (309 observations).

    Source: statsmodels built-in (sm.datasets.sunspots).

    Returns
    -------
    pd.DataFrame
        Integer ``'year'`` index, single column ``'sunactivity'``.

    Examples
    --------
    >>> from peshbeen.datasets import load_sunspots
    >>> df = load_sunspots()
    >>> df.head()
    """
    data = _sm_dataset('sunspots')
    data = data.rename(columns={'YEAR': 'year', 'SUNACTIVITY': 'sunactivity'})
    data['year'] = data['year'].astype(int)
    data = data.set_index('year')
    return data


# nile dataset
def load_nile_ts() -> pd.DataFrame:
    """
    Annual flow of the Nile river at Aswan from 1871 to 1970
    (100 observations, measured in 10^8 m³).

    A classic structural-break benchmark series with a well-known level
    shift around 1898.

    Source: statsmodels built-in (sm.datasets.nile).

    Returns
    -------
    pd.DataFrame
        Integer ``'year'`` index, single column ``'flow'``.

    Examples
    --------
    >>> from peshbeen.datasets import load_nile
    >>> df = load_nile()
    >>> df.head()
    """
    data = _sm_dataset('nile')
    data.columns = [c.lower() for c in data.columns]
    data['year'] = data['year'].astype(int)
    data = data.rename(columns={'volume': 'flow'})
    data = data.set_index('year')
    return data

# births dataset
def load_births_ts() -> pd.DataFrame:
    """
    Daily total of female births in California, 1959 (365 observations).

    A classic univariate daily series with clear day-of-week effects,
    originally from Newton (1988) and widely used in forecasting tutorials.

    This dataset is fully hardcoded — no internet connection required.

    Returns
    -------
    pd.DataFrame
        DatetimeIndex (daily frequency 1959-01-01 to 1959-12-31),
        columns ``'births'`` and ``'day_of_week'`` (0 = Monday).

    Examples
    --------
    >>> from peshbeen.datasets import load_births
    >>> df = load_births()
    >>> df.head()
    """
    # Daily female births in California, 1959 — Newton (1988), public domain
    _births = [
        35, 32, 30, 31, 44, 29, 45, 43, 38, 27, 38, 33, 55, 47, 45, 37, 50,
        43, 41, 52, 34, 53, 39, 32, 37, 43, 39, 35, 44, 38, 43, 38, 40, 39,
        46, 33, 40, 40, 44, 37, 36, 41, 39, 37, 40, 39, 43, 33, 44, 42, 35,
        38, 37, 45, 39, 41, 41, 43, 48, 40, 44, 46, 44, 49, 43, 35, 39, 40,
        44, 36, 38, 44, 40, 43, 34, 39, 43, 38, 45, 37, 40, 39, 40, 41, 42,
        40, 38, 40, 48, 43, 47, 46, 41, 42, 39, 44, 34, 40, 45, 34, 37, 36,
        35, 41, 42, 43, 36, 39, 42, 44, 43, 41, 40, 42, 38, 32, 42, 37, 43,
        38, 40, 40, 40, 34, 40, 37, 41, 36, 39, 37, 42, 44, 43, 44, 37, 44,
        43, 42, 44, 39, 43, 37, 42, 38, 43, 41, 35, 40, 38, 40, 41, 40, 43,
        39, 36, 43, 40, 38, 44, 37, 40, 40, 40, 43, 45, 38, 35, 40, 46, 40,
        38, 43, 38, 43, 39, 43, 40, 43, 38, 41, 40, 41, 40, 44, 38, 41, 37,
        42, 40, 41, 40, 37, 37, 40, 39, 42, 42, 39, 40, 38, 40, 43, 39, 42,
        35, 41, 35, 42, 41, 44, 38, 43, 37, 39, 35, 39, 38, 41, 43, 43, 37,
        42, 35, 35, 39, 39, 43, 40, 42, 44, 41, 43, 39, 40, 40, 46, 39, 44,
        43, 43, 46, 42, 38, 37, 42, 40, 41, 44, 45, 44, 43, 45, 40, 44, 41,
        44, 39, 43, 40, 40, 40, 41, 38, 43, 40, 42, 40, 41, 40, 40, 41, 40,
        36, 42, 38, 36, 36, 40, 43, 39, 40, 40, 42, 40, 43, 42, 45, 39, 38,
        39, 40, 40, 41, 38, 41, 41, 44, 40, 37, 36, 40, 44, 38, 41, 41, 44,
        38, 42, 39, 41, 38, 39, 35, 45, 40, 36, 38, 42, 35, 43, 42, 46, 41,
        44, 38, 44, 35, 38, 38, 44, 46, 44, 44, 38, 43, 43, 40, 42, 41, 43,
        44, 38, 43, 43, 42, 44, 42, 42, 39, 40, 40, 40, 42, 38, 41, 44, 41,
        41, 39, 43, 43, 42, 43, 38, 40,
    ]
    idx = pd.date_range('1959-01-01', periods=365, freq='D')
    df = pd.DataFrame({'births': _births}, index=idx)
    df['day_of_week'] = df.index.dayofweek
    return df


# elnino dataset

def load_elnino_ts() -> pd.DataFrame:
    """
    Monthly average sea-surface temperature (°C) in the equatorial Pacific
    (Niño 1+2 region, 0–10°S / 90–80°W) from January 1950 to December 2010
    (732 observations).

    Strong annual seasonality makes this a good univariate benchmark for
    seasonal models (ETS, SARIMA, seasonal differencing).

    Source: statsmodels built-in (sm.datasets.elnino). The raw data is stored
    in wide format (61 rows × 12 month columns); this loader melts it into a
    clean monthly time series.

    Returns
    -------
    pd.DataFrame
        DatetimeIndex (monthly frequency), single column ``'temperature'``.

    Examples
    --------
    >>> from peshbeen.datasets import load_elnino
    >>> df = load_elnino()
    >>> df.head()
    """
    import statsmodels.api as sm
    raw = sm.datasets.elnino.load_pandas().data

    # The wide format has columns: YEAR, JAN, FEB, ..., DEC
    month_map = {
        'JAN': 1, 'FEB': 2, 'MAR': 3, 'APR': 4,
        'MAY': 5, 'JUN': 6, 'JUL': 7, 'AUG': 8,
        'SEP': 9, 'OCT': 10, 'NOV': 11, 'DEC': 12
    }
    # Normalise column names (strip whitespace, upper-case)
    raw.columns = [c.strip().upper() for c in raw.columns]
    year_col = [c for c in raw.columns if 'YEAR' in c][0]
    month_cols = [c for c in raw.columns if c in month_map]

    records = []
    for _, row in raw.iterrows():
        year = int(row[year_col])
        for mc in month_cols:
            records.append({
                'date': pd.Timestamp(year=year, month=month_map[mc], day=1),
                'temperature': float(row[mc])
            })

    df = pd.DataFrame(records).set_index('date').sort_index()
    df.index.freq = 'MS'
    return df

# ─────────────────────────────────────────────────────────────────────────────
# MULTIVARIATE DATASETS
# ─────────────────────────────────────────────────────────────────────────────

# macrodata dataset
def load_macrodata_ts(
    columns: list = None
) -> pd.DataFrame:
    """
    Quarterly US macroeconomic indicators from 1959 Q1 to 2009 Q3
    (202 observations).

    Available columns
    -----------------
    realgdp     Real GDP (billions of chained 2000 dollars)
    realcons    Real personal consumption expenditures
    realinv     Real gross private domestic investment
    realgovt    Real federal government consumption and investment
    realdpi     Real private disposable income
    cpi         Consumer price index
    m1          M1 nominal money stock
    tbilrate    Monthly 3-month treasury bill rate
    unemp       Seasonally adjusted unemployment rate (%)
    pop         Population (thousands)
    infl        Inflation rate
    realint     Real interest rate

    Source: statsmodels built-in (sm.datasets.macrodata).

    Parameters
    ----------
    columns : list, optional
        Subset of columns to return. Returns all columns if None.

    Returns
    -------
    pd.DataFrame
        DatetimeIndex (quarterly frequency), requested columns.

    Examples
    --------
    >>> from peshbeen.datasets import load_macrodata
    >>> df = load_macrodata(columns=['realgdp', 'unemp', 'infl'])
    >>> df.head()
    """
    data = _sm_dataset('macrodata')
    # Build a proper quarterly DatetimeIndex from year + quarter columns
    dates = pd.PeriodIndex(
        [pd.Period(f"{int(y)}Q{int(q)}", freq='Q')
         for y, q in zip(data['year'], data['quarter'])]
    ).to_timestamp()
    data = data.drop(columns=['year', 'quarter'])
    data.index = dates
    data.index.freq = pd.infer_freq(data.index)
    if columns is not None:
        data = data[columns]
    return data

# interest and inflation dataset
def load_interest_inflation_ts() -> pd.DataFrame:
    """
    Monthly West German interest and inflation rates from January 1972
    to December 1998 (324 observations).

    Columns
    -------
    Dp      First difference of CPI (used as inflation proxy)
    R       Nominal short-term interest rate (%)

    A standard bivariate VAR example from econometrics literature.

    Source: statsmodels built-in (sm.datasets.interest_inflation).

    Returns
    -------
    pd.DataFrame
        DatetimeIndex (monthly frequency), columns ``'interest'`` and
        ``'inflation'``.

    Examples
    --------
    >>> from peshbeen.datasets import load_interest_inflation
    >>> df = load_interest_inflation()
    >>> df.head()
    """
    data = _sm_dataset('interest_inflation')
    data = data.rename(columns={'Dp': 'inflation', 'R': 'interest'})
    idx = pd.date_range(start='1972-01', periods=len(data), freq='MS')
    data.index = idx
    return data[['interest', 'inflation']]

def load_interest_inflation() -> pd.DataFrame:
    """
    Monthly West German interest and inflation rates from January 1972
    to December 1998 (324 observations).

    Columns
    -------
    Dp      First difference of CPI (used as inflation proxy)
    R       Nominal short-term interest rate (%)

    A standard bivariate VAR example from econometrics literature.

    Source: statsmodels built-in (sm.datasets.interest_inflation).

    Returns
    -------
    pd.DataFrame
        DatetimeIndex (monthly frequency), columns ``'interest'`` and
        ``'inflation'``.

    Examples
    --------
    >>> from peshbeen.datasets import load_interest_inflation
    >>> df = load_interest_inflation()
    >>> df.head()
    """
    data = _sm_dataset('interest_inflation')
    data = data.rename(columns={'Dp': 'inflation', 'R': 'interest'})
    idx = pd.date_range(start='1972-01', periods=len(data), freq='MS')
    data.index = idx
    return data[['interest', 'inflation']]

# livestock dataset
def load_livestock_ts(columns: list = None) -> pd.DataFrame:
    """
    Annual livestock counts in Australia from 1970 to 2007 (38 observations).

    Three co-evolving animal population series — sheep, cattle, and pigs —
    that share common long-run drivers (rainfall, feed prices, land use).
    A clean, compact multivariate benchmark with no missing values.

    This dataset is fully hardcoded — no internet connection required.
    Values are in millions of head, from the Australian Bureau of Statistics.

    Available columns
    -----------------
    sheep    Total sheep and lambs (millions)
    cattle   Total cattle and calves (millions)
    pigs     Total pigs (millions)

    Parameters
    ----------
    columns : list, optional
        Subset of ``['sheep', 'cattle', 'pigs']``. Returns all if None.

    Returns
    -------
    pd.DataFrame
        DatetimeIndex (annual frequency, year-start), requested columns.

    Examples
    --------
    >>> from peshbeen.datasets import load_livestock
    >>> df = load_livestock()
    >>> df.head()
    >>> df = load_livestock(columns=['sheep', 'cattle'])
    """
    # Source: Australian Bureau of Statistics, publicly available
    # Values in millions of head, 1970–2007
    years  = list(range(1970, 2008))
    sheep  = [
        155.1, 163.8, 166.1, 167.7, 166.5, 168.1, 170.6, 169.9, 166.5,
        168.2, 164.9, 165.5, 166.4, 163.9, 156.5, 157.3, 154.5, 156.3,
        158.0, 166.1, 170.8, 170.4, 164.1, 159.3, 155.0, 147.7, 138.6,
        134.0, 130.0, 127.5, 118.9, 115.0, 112.3, 110.8, 103.7, 101.0,
        100.2, 101.5
    ]
    cattle = [
        23.8, 24.6, 25.2, 27.1, 28.5, 29.3, 30.7, 31.0, 31.3, 32.0,
        30.3, 30.0, 29.9, 28.6, 28.7, 27.9, 28.2, 28.7, 29.5, 30.0,
        29.2, 28.0, 27.3, 27.7, 27.3, 27.0, 26.9, 27.0, 27.1, 27.6,
        27.5, 27.3, 27.5, 27.8, 27.9, 27.9, 28.0, 27.9
    ]
    pigs   = [
        2.28, 2.33, 2.37, 2.44, 2.50, 2.52, 2.54, 2.56, 2.57, 2.53,
        2.49, 2.48, 2.47, 2.46, 2.46, 2.46, 2.46, 2.45, 2.43, 2.42,
        2.41, 2.40, 2.38, 2.36, 2.35, 2.33, 2.31, 2.30, 2.29, 2.28,
        2.27, 2.26, 2.25, 2.23, 2.22, 2.21, 2.20, 2.20
    ]
    idx = pd.date_range(start='1970', periods=38, freq='YS')
    df = pd.DataFrame({'sheep': sheep, 'cattle': cattle, 'pigs': pigs}, index=idx)
    if columns is not None:
        df = df[columns]
    return df



# %% auto #0
__all__ = ['load_airline_passengers', 'load_co2', 'load_sunspots', 'load_nile', 'load_macrodata', 'load_interest_inflation',
           'load_elnino', 'load_livestock', 'load_wales_admissions', 'load_wales_calls', 'load_admission_calls',
           'admission_ts', 'calls_ts', 'admission_calls_ts', 'airline_passengers', 'load_co2_ts', 'load_sunspots_ts',
           'load_nile_ts', 'load_births_ts', 'load_elnino_ts', 'load_macrodata_ts', 'load_interest_inflation_ts',
           'load_livestock_ts', 'available_datasets']

# %% ../nbs/modules/99_datasets.ipynb #ba80f2d4
load_airline_passengers = airline_passengers()

# %% ../nbs/modules/99_datasets.ipynb #c971aad0
load_co2 = load_co2_ts()

# %% ../nbs/modules/99_datasets.ipynb #e81135c9
load_sunspots = load_sunspots_ts()

# %% ../nbs/modules/99_datasets.ipynb #09a4bec7
load_nile = load_nile_ts()

# %% ../nbs/modules/99_datasets.ipynb #fdf0a670
load_macrodata = load_macrodata_ts()

# %% ../nbs/modules/99_datasets.ipynb #772682fc
load_interest_inflation = load_interest_inflation_ts()

# %% ../nbs/modules/99_datasets.ipynb #2604e19d
load_elnino = load_elnino_ts()

# %% ../nbs/modules/99_datasets.ipynb #9654ff11
load_livestock = load_livestock_ts()

# %% ../nbs/modules/99_datasets.ipynb #de0d3bca
load_wales_admissions = admission_ts().set_index('Date')

# %% ../nbs/modules/99_datasets.ipynb #9b5da161
load_wales_calls = calls_ts().set_index('Date')

# %% ../nbs/modules/99_datasets.ipynb #8aab60b8
load_admission_calls = admission_calls_ts()

# %% ../nbs/modules/99_datasets.ipynb #8040662b
# ─────────────────────────────────────────────────────────────────────────────
# CONVENIENCE: list available datasets
# ─────────────────────────────────────────────────────────────────────────────

def available_datasets() -> pd.DataFrame:
    """
    Return a summary table of all datasets available in peshbeen.datasets.

    Returns
    -------
    pd.DataFrame
        Columns: ``'loader'``, ``'type'``, ``'frequency'``,
        ``'approx_obs'``, ``'description'``.

    Examples
    --------
    >>> from peshbeen.datasets import available_datasets
    >>> available_datasets()
    """
    rows = [
        dict(loader='load_co2()',
             type='univariate',
             frequency='weekly',
             approx_obs=2284,
             description='Mauna Loa atmospheric CO2 concentration (ppm), 1958–2001'),
        dict(loader='load_sunspots()',
             type='univariate',
             frequency='annual',
             approx_obs=309,
             description='Annual sunspot activity count, 1700–2008'),
        dict(loader='load_air_passengers()',
             type='univariate',
             frequency='monthly',
             approx_obs=144,
             description='Monthly international airline passengers (thousands), 1949–1960'),
        dict(loader='load_nile()',
             type='univariate',
             frequency='annual',
             approx_obs=100,
             description='Annual Nile river flow at Aswan (10^8 m³), 1871–1970'),
        dict(loader='load_births()',
             type='univariate',
             frequency='daily',
             approx_obs=365,
             description='Daily US female births, 1969 (day-of-week effects)'),
        dict(loader='load_elnino()',
             type='univariate',
             frequency='monthly',
             approx_obs=732,
             description='El Niño sea-surface temperature indices (Niño 1–4), 1950–2010'),
        dict(loader='load_wales_admissions()',
          type='univariate',
          frequency='daily',
          approx_obs=1068,
          description='Daily hospital admissions in Wales, UK'),
       dict(loader='load_wales_calls()',
          type='univariate',
          frequency='daily',
          approx_obs=493,
          description='Daily emergency calls in Wales, UK'),
        dict(loader='load_macrodata()',
             type='multivariate',
             frequency='quarterly',
             approx_obs=202,
             description='12 US macroeconomic indicators (GDP, CPI, unemployment …), 1959–2009'),
        dict(loader='load_interest_inflation()',
             type='multivariate',
             frequency='monthly',
             approx_obs=324,
             description='German interest rate & inflation rate, 1972–1998'),
        dict(loader='load_livestock()',
             type='multivariate',
             frequency='annual',
             approx_obs=38,
             description='Australian sheep, cattle & pig counts (millions of head), 1970–2007'),
        dict(loader='load_admission_calls()',
          type='multivariate',
          frequency='daily',
          approx_obs=433,
          description='Combined dataset of daily hospital admissions and emergency calls in Wales, UK'),
    ]
    return pd.DataFrame(rows)
