from collections import defaultdict
from datetime import date
import os
import re

import datasets as ds
import numpy as np
import pandas as pd
import requests
from tqdm import tqdm


def _get_hf_token() -> str:
    token = os.getenv("HF_ACCESS_TOKEN")
    if not token:
        raise ValueError("HF_ACCESS_TOKEN not set")
    return token


def _get_pwb_api_key() -> str | None:
    return os.getenv("PWB_API_KEY")


def _load_dataset_from_pwb(dataset_name: str, split: str, pwb_api_key: str):
    base_url = "https://paperswithbacktest.com/api/v1/datasets"
    resp = requests.get(f"{base_url}/{dataset_name}", params={"pwb_api_key": pwb_api_key, "split": split}, timeout=30)
    resp.raise_for_status()
    files = resp.json().get("files", [])
    if not files:
        raise ValueError(f"No files available for dataset '{dataset_name}' and split '{split}'")

    print(f"Downloading {len(files)} parquet files from PWB...")
    download_config = ds.DownloadConfig(
        max_retries=10,
        storage_options={
            # aiohttp total timeout in seconds
            "timeout": 3600,
        },
    )
    return ds.load_dataset(
        "parquet", 
        data_files={split: files}, 
        download_config=download_config,
        split=split, 
        download_mode="force_redownload", 
        verification_mode="no_checks",
    )


DAILY_PRICE_DATASETS = [
    "Bonds-Daily-Price",
    "Commodities-Daily-Price",
    "Cryptocurrencies-Daily-Price",
    "ETFs-Daily-Price",
    "Forex-Daily-Price",
    "Indices-Daily-Price",
    "Stocks-Daily-Price",
]


SP500_SYMBOLS = [
    "MMM",
    "AOS",
    "ABT",
    "ABBV",
    "ACN",
    "ADBE",
    "AMD",
    "AES",
    "AFL",
    "A",
    "APD",
    "ABNB",
    "AKAM",
    "ALB",
    "ARE",
    "ALGN",
    "ALLE",
    "LNT",
    "ALL",
    "GOOGL",
    "GOOG",
    "MO",
    "AMZN",
    "AMCR",
    "AEE",
    "AEP",
    "AXP",
    "AIG",
    "AMT",
    "AWK",
    "AMP",
    "AME",
    "AMGN",
    "APH",
    "ADI",
    "ANSS",
    "AON",
    "APA",
    "AAPL",
    "AMAT",
    "APTV",
    "ACGL",
    "ADM",
    "ANET",
    "AJG",
    "AIZ",
    "T",
    "ATO",
    "ADSK",
    "ADP",
    "AZO",
    "AVB",
    "AVY",
    "AXON",
    "BKR",
    "BALL",
    "BAC",
    "BK",
    "BBWI",
    "BAX",
    "BDX",
    "BRK-B",
    "BBY",
    "TECH",
    "BIIB",
    "BLK",
    "BX",
    "BA",
    "BKNG",
    "BWA",
    "BSX",
    "BMY",
    "AVGO",
    "BR",
    "BRO",
    "BF-B",
    "BLDR",
    "BG",
    "BXP",
    "CHRW",
    "CDNS",
    "CZR",
    "CPT",
    "CPB",
    "COF",
    "CAH",
    "KMX",
    "CCL",
    "CARR",
    "CTLT",
    "CAT",
    "CBOE",
    "CBRE",
    "CDW",
    "CE",
    "COR",
    "CNC",
    "CNP",
    "CF",
    "CRL",
    "SCHW",
    "CHTR",
    "CVX",
    "CMG",
    "CB",
    "CHD",
    "CI",
    "CINF",
    "CTAS",
    "CSCO",
    "C",
    "CFG",
    "CLX",
    "CME",
    "CMS",
    "KO",
    "CTSH",
    "CL",
    "CMCSA",
    "CAG",
    "COP",
    "ED",
    "STZ",
    "CEG",
    "COO",
    "CPRT",
    "GLW",
    "CPAY",
    "CTVA",
    "CSGP",
    "COST",
    "CTRA",
    "CRWD",
    "CCI",
    "CSX",
    "CMI",
    "CVS",
    "DHR",
    "DRI",
    "DVA",
    "DAY",
    "DECK",
    "DE",
    "DELL",
    "DAL",
    "DVN",
    "DXCM",
    "FANG",
    "DLR",
    "DFS",
    "DG",
    "DLTR",
    "D",
    "DPZ",
    "DOV",
    "DOW",
    "DHI",
    "DTE",
    "DUK",
    "DD",
    "EMN",
    "ETN",
    "EBAY",
    "ECL",
    "EIX",
    "EW",
    "EA",
    "ELV",
    "EMR",
    "ENPH",
    "ETR",
    "EOG",
    "EPAM",
    "EQT",
    "EFX",
    "EQIX",
    "EQR",
    "ERIE",
    "ESS",
    "EL",
    "EG",
    "EVRG",
    "ES",
    "EXC",
    "EXPE",
    "EXPD",
    "EXR",
    "XOM",
    "FFIV",
    "FDS",
    "FICO",
    "FAST",
    "FRT",
    "FDX",
    "FIS",
    "FITB",
    "FSLR",
    "FE",
    "FI",
    "FMC",
    "F",
    "FTNT",
    "FTV",
    "FOXA",
    "FOX",
    "BEN",
    "FCX",
    "GRMN",
    "IT",
    "GE",
    "GEHC",
    "GEV",
    "GEN",
    "GNRC",
    "GD",
    "GIS",
    "GM",
    "GPC",
    "GILD",
    "GPN",
    "GL",
    "GDDY",
    "GS",
    "HAL",
    "HIG",
    "HAS",
    "HCA",
    "DOC",
    "HSIC",
    "HSY",
    "HES",
    "HPE",
    "HLT",
    "HOLX",
    "HD",
    "HON",
    "HRL",
    "HST",
    "HWM",
    "HPQ",
    "HUBB",
    "HUM",
    "HBAN",
    "HII",
    "IBM",
    "IEX",
    "IDXX",
    "ITW",
    "INCY",
    "IR",
    "PODD",
    "INTC",
    "ICE",
    "IFF",
    "IP",
    "IPG",
    "INTU",
    "ISRG",
    "IVZ",
    "INVH",
    "IQV",
    "IRM",
    "JBHT",
    "JBL",
    "JKHY",
    "J",
    "JNJ",
    "JCI",
    "JPM",
    "JNPR",
    "K",
    "KVUE",
    "KDP",
    "KEY",
    "KEYS",
    "KMB",
    "KIM",
    "KMI",
    "KKR",
    "KLAC",
    "KHC",
    "KR",
    "LHX",
    "LH",
    "LRCX",
    "LW",
    "LVS",
    "LDOS",
    "LEN",
    "LLY",
    "LIN",
    "LYV",
    "LKQ",
    "LMT",
    "L",
    "LOW",
    "LULU",
    "LYB",
    "MTB",
    "MRO",
    "MPC",
    "MKTX",
    "MAR",
    "MMC",
    "MLM",
    "MAS",
    "MA",
    "MTCH",
    "MKC",
    "MCD",
    "MCK",
    "MDT",
    "MRK",
    "META",
    "MET",
    "MTD",
    "MGM",
    "MCHP",
    "MU",
    "MSFT",
    "MAA",
    "MRNA",
    "MHK",
    "MOH",
    "TAP",
    "MDLZ",
    "MPWR",
    "MNST",
    "MCO",
    "MS",
    "MOS",
    "MSI",
    "MSCI",
    "NDAQ",
    "NTAP",
    "NFLX",
    "NEM",
    "NWSA",
    "NWS",
    "NEE",
    "NKE",
    "NI",
    "NDSN",
    "NSC",
    "NTRS",
    "NOC",
    "NCLH",
    "NRG",
    "NUE",
    "NVDA",
    "NVR",
    "NXPI",
    "ORLY",
    "OXY",
    "ODFL",
    "OMC",
    "ON",
    "OKE",
    "ORCL",
    "OTIS",
    "PCAR",
    "PKG",
    "PLTR",
    "PANW",
    "PARA",
    "PH",
    "PAYX",
    "PAYC",
    "PYPL",
    "PNR",
    "PEP",
    "PFE",
    "PCG",
    "PM",
    "PSX",
    "PNW",
    "PNC",
    "POOL",
    "PPG",
    "PPL",
    "PFG",
    "PG",
    "PGR",
    "PLD",
    "PRU",
    "PEG",
    "PTC",
    "PSA",
    "PHM",
    "QRVO",
    "PWR",
    "QCOM",
    "DGX",
    "RL",
    "RJF",
    "RTX",
    "O",
    "REG",
    "REGN",
    "RF",
    "RSG",
    "RMD",
    "RVTY",
    "ROK",
    "ROL",
    "ROP",
    "ROST",
    "RCL",
    "SPGI",
    "CRM",
    "SBAC",
    "SLB",
    "STX",
    "SRE",
    "NOW",
    "SHW",
    "SPG",
    "SWKS",
    "SJM",
    "SW",
    "SNA",
    "SOLV",
    "SO",
    "LUV",
    "SWK",
    "SBUX",
    "STT",
    "STLD",
    "STE",
    "SYK",
    "SMCI",
    "SYF",
    "SNPS",
    "SYY",
    "TMUS",
    "TROW",
    "TTWO",
    "TPR",
    "TRGP",
    "TGT",
    "TEL",
    "TDY",
    "TFX",
    "TER",
    "TSLA",
    "TXN",
    "TXT",
    "TMO",
    "TJX",
    "TSCO",
    "TT",
    "TDG",
    "TRV",
    "TRMB",
    "TFC",
    "TYL",
    "TSN",
    "USB",
    "UBER",
    "UDR",
    "ULTA",
    "UNP",
    "UAL",
    "UPS",
    "URI",
    "UNH",
    "UHS",
    "VLO",
    "VTR",
    "VLTO",
    "VRSN",
    "VRSK",
    "VZ",
    "VRTX",
    "VTRS",
    "VICI",
    "V",
    "VST",
    "VMC",
    "WRB",
    "GWW",
    "WAB",
    "WBA",
    "WMT",
    "DIS",
    "WBD",
    "WM",
    "WAT",
    "WEC",
    "WFC",
    "WELL",
    "WST",
    "WDC",
    "WY",
    "WMB",
    "WTW",
    "WYNN",
    "XEL",
    "XYL",
    "YUM",
    "ZBRA",
    "ZBH",
    "ZTS",
]


def load_dataset(
    path,
    symbols=None,
    adjust=True,
    extend=False,
    to_usd=True,
    rate_to_price=True,
    use_hf=False
):
    split = "train"
    pwb_api_key = _get_pwb_api_key()

    if pwb_api_key and not use_hf:
        dataset = _load_dataset_from_pwb(path, split=split, pwb_api_key=pwb_api_key)
    else:
        hf_token = os.getenv("HF_ACCESS_TOKEN")
        if not hf_token:
            raise ValueError("Set PWB_API_KEY or HF_ACCESS_TOKEN to load datasets.")
        dataset = ds.load_dataset(f"paperswithbacktest/{path}", token=hf_token)

    df = (dataset[split] if isinstance(dataset, ds.DatasetDict) else dataset).to_pandas()

    if "date" in df.columns:
        df["date"] = pd.to_datetime(df["date"]).dt.date

    if "datetime" in df.columns:
        df["datetime"] = pd.to_datetime(df["datetime"])

    if isinstance(symbols, list) and "sp500" in symbols:
        symbols.remove("sp500")
        symbols += SP500_SYMBOLS

    if "symbol" in df.columns and isinstance(symbols, list):
        df = df[df["symbol"].isin(symbols)].copy()

    if "symbols" in df.columns and isinstance(symbols, list):
        df = df[
            df["symbols"].apply(lambda x: any(symbol in symbols for symbol in x))
        ].copy()

    if path in DAILY_PRICE_DATASETS:
        if adjust and "adj_close" in df.columns:
            adj_factor = df["adj_close"] / df["close"]
            df["adj_open"] = df["open"] * adj_factor
            df["adj_high"] = df["high"] * adj_factor
            df["adj_low"] = df["low"] * adj_factor
            df.drop(columns=["open", "high", "low", "close"], inplace=True)
            df.rename(
                columns={
                    "adj_open": "open",
                    "adj_high": "high",
                    "adj_low": "low",
                    "adj_close": "close",
                },
                inplace=True,
            )
        else:
            if "adj_close" in df.columns:
                df.drop(columns=["adj_close"])

    if path in DAILY_PRICE_DATASETS and (extend and path == "ETFs-Daily-Price"):
        df = __extend_etfs(df)

    if path in DAILY_PRICE_DATASETS and to_usd:
        if path == "Forex-Daily-Price":
            for index, row in df.iterrows():
                if row["symbol"].endswith("USD"):
                    continue
                df.at[index, "open"] = 1 / row["open"]
                df.at[index, "high"] = 1 / row["high"]
                df.at[index, "low"] = 1 / row["low"]
                df.at[index, "close"] = 1 / row["close"]
                df.at[index, "symbol"] = row["symbol"][3:] + "USD"
        elif path == "Indices-Daily-Price":
            df_forex = load_dataset("Forex-Daily-Price", to_usd=True)
            df = __convert_indices_to_usd(df, df_forex)

    if path in DAILY_PRICE_DATASETS and rate_to_price and path == "Bonds-Daily-Price":
        df = __convert_bond_rates_to_prices(df)

    return df


def __convert_indices_to_usd(
    df_indices: pd.DataFrame, df_forex: pd.DataFrame
) -> pd.DataFrame:
    mapping = {
        "ADSMI": "AED",  # United Arab Emirates
        "AEX": "EUR",  # Netherlands
        "AS30": "AUD",  # Australia
        "AS51": "AUD",  # Australia
        "AS52": "AUD",  # Australia
        "ASE": "EUR",  # Greece
        "ATX": "EUR",  # Austria
        "BEL20": "EUR",  # Belgium
        "BELEX15": "RSD",  # Serbia
        "BGSMDC": "BWP",  # Botswana
        "BHSEEI": "BHD",  # Bahrain
        "BKA": "BAM",  # Bosnia and Herzegovina
        "BLOM": "LBP",  # Lebanon
        "BSX": "BMD",  # Bermuda
        "BUX": "HUF",  # Hungary
        "BVLX": "BOB",  # Bolivia
        "BVPSBVPS": "PAB",  # Panama
        "BVQA": "USD",  # Ecuador
        "CAC": "EUR",  # France
        "CASE": "EGP",  # Egypt
        "CCMP": "USD",  # United States
        "COLCAP": "COP",  # Colombia
        "CRSMBCT": "CRC",  # Costa Rica
        "CSEALL": "LKR",  # Sri Lanka
        "CYSMMAPA": "EUR",  # Cyprus
        "DARSDSEI": "TZS",  # Tanzania
        "DAX": "EUR",  # Germany
        "DFMGI": "AED",  # United Arab Emirates
        "DSEX": "BDT",  # Bangladesh
        "DSM": "QAR",  # Qatar
        "ECU": "USD",  # Ecuador
        "FBMKLCI": "MYR",  # Malaysia
        "FSSTI": "SGD",  # Singapore
        "FTN098": "NAD",  # Namibia
        "FTSEMIB": "EUR",  # Italy
        "GGSECI": "GHS",  # Ghana
        "HEX": "EUR",  # Finland
        "HEX25": "EUR",  # Finland
        "HSI": "HKD",  # Hong Kong
        "IBEX": "EUR",  # Spain
        "IBOV": "BRL",  # Brazil
        "IBVC": "VES",  # Venezuela
        "ICEXI": "ISK",  # Iceland
        "IGPA": "CLP",  # Chile
        "INDEXCF": "RUB",  # Russia
        "INDU": "USD",  # United States
        "INDZI": "IDR",  # Indonesia
        "ISEQ": "EUR",  # Ireland
        "JALSH": "ZAR",  # South Africa
        "JCI": "IDR",  # Indonesia
        "JMSMX": "JMD",  # Jamaica
        "JOSMGNFF": "JOD",  # Jordan
        "KFX": "DKK",  # Denmark
        "KNSMIDX": "KES",  # Kenya
        "KSE100": "PKR",  # Pakistan
        "KZKAK": "KZT",  # Kazakhstan
        "LSXC": "LAK",  # Laos
        "LUXXX": "EUR",  # Luxembourg
        "MALTEX": "EUR",  # Malta
        "MBI": "MKD",  # North Macedonia
        "MERVAL": "ARS",  # Argentina
        "MEXBOL": "MXN",  # Mexico
        "MONEX": "EUR",  # Montenegro
        "MOSENEW": "MAD",  # Morocco
        "MSETOP": "MKD",  # North Macedonia
        "MSM30": "OMR",  # Oman
        "NDX": "USD",  # United States
        "NGSEINDX": "NGN",  # Nigeria
        "NIFTY": "INR",  # India
        "NKY": "JPY",  # Japan
        "NSEASI": "KES",  # Kenya
        "NZSE50FG": "NZD",  # New Zealand
        "OMX": "SEK",  # Sweden
        "OSEAX": "NOK",  # Norway
        "PCOMP": "PHP",  # Philippines
        "PFTS": "UAH",  # Ukraine
        "PSI20": "EUR",  # Portugal
        "PX": "CZK",  # Czech Republic
        "RIGSE": "EUR",  # Latvia
        "RTY": "USD",  # United States
        "SASEIDX": "SAR",  # Saudi Arabia
        "SASX10": "BAM",  # Bosnia and Herzegovina
        "SBITOP": "EUR",  # Slovenia
        "SEMDEX": "MUR",  # Mauritius
        "SENSEX": "INR",  # India
        "SET50": "THB",  # Thailand
        "SHCOMP": "CNY",  # China
        "SHSZ300": "CNY",  # China
        "SKSM": "EUR",  # Slovakia
        "SMI": "CHF",  # Switzerland
        "SOFIX": "BGN",  # Bulgaria
        "SPBLPGPT": "PEN",  # Peru
        "SPTSX": "CAD",  # Canada
        "SPX": "USD",  # United States
        "SSE50": "CNY",  # China
        "SX5E": "EUR",  # Europe
        "TA125": "ILS",  # Israel
    }
    frames = []

    # iterate over the symbols that actually exist in df_indices
    for symbol in df_indices["symbol"].unique():
        df_idx = df_indices[df_indices["symbol"] == symbol].copy()

        # 1) Figure out what currency the index is quoted in.
        ccy = mapping.get(symbol)  # None if not mapped
        if ccy is None or ccy == "USD":
            # Unknown or already USD – just keep the original rows
            frames.append(df_idx)
            continue

        # 2) Find the matching FX rate (home-ccy → USD)
        pair = ccy + "USD"
        df_fx = df_forex[df_forex["symbol"] == pair].copy()

        if df_idx.empty or df_fx.empty:
            # No FX data – keep raw index levels instead of dropping them
            frames.append(df_idx)
            continue

        # 3) Merge on date and convert OHLC
        merged = pd.merge(df_idx, df_fx, on="date", suffixes=("", "_fx"))
        for col in ("open", "high", "low", "close"):
            merged[col] = merged[col] * merged[f"{col}_fx"]

        frames.append(merged[["symbol", "date", "open", "high", "low", "close"]])

    if not frames:
        return pd.DataFrame(columns=df_indices.columns)

    # Combine everything back into one DataFrame
    return pd.concat(frames, ignore_index=True)


def __convert_bond_rates_to_prices(df: pd.DataFrame, face_value: float = 100.0) -> pd.DataFrame:
    """
    Convert bond yields (in percent) to prices, in-place, for rows whose
    symbols encode maturity like 'US10Y', 'US2Y', 'US3M', etc.

    Assumes columns: ['symbol', 'open', 'high', 'low', 'close'] and that
    open/high/low/close are yields in percent.
    """
    # Vectorized extraction of maturity from symbol: last part "10Y", "2Y", "3M", ...
    m = df["symbol"].str.extract(r"(\d+)([YM])$")  # 0: number, 1: unit
    num = pd.to_numeric(m[0], errors="coerce")
    unit = m[1]

    years = pd.Series(np.nan, index=df.index, dtype="float64")
    years[unit == "Y"] = num[unit == "Y"]
    years[unit == "M"] = num[unit == "M"] / 12.0

    # Only apply to rows where we successfully parsed the maturity
    mask = years.notna()
    if not mask.any():
        return df

    y = years[mask]

    # Vectorized rate->price for each OHLC column
    for col in ["open", "high", "low", "close"]:
        r = df.loc[mask, col] / 100.0  # percent -> fraction
        df.loc[mask, col] = face_value / (1.0 + r) ** y

    return df


def __extend_etfs(df_etfs):

    mapping = {
        "AGG": ["Bonds-Daily-Price", "US10Y"],
        "EPP": ["Indices-Daily-Price", "HSI"],
        "EWA": ["Indices-Daily-Price", "AS30"],  # Australia
        "EWO": ["Indices-Daily-Price", "ATX"],  # Austria
        "EWK": ["Indices-Daily-Price", "BEL20"],  # Belgium
        "EWZ": ["Indices-Daily-Price", "IBOV"],  # Brazil
        "EWC": ["Indices-Daily-Price", "SPTSX"],  # Canada
        "FXI": ["Indices-Daily-Price", "SSE50"],  # China
        "EWQ": ["Indices-Daily-Price", "CAC"],  # France
        "EWG": ["Indices-Daily-Price", "DAX"],  # Germany
        "EWH": ["Indices-Daily-Price", "HSI"],  # Hong Kong
        "EWI": ["Indices-Daily-Price", "FTSEMIB"],  # Italy
        "EWJ": ["Indices-Daily-Price", "NKY"],
        "EWM": ["Indices-Daily-Price", "FBMKLCI"],  # Malaysia
        "EWW": ["Indices-Daily-Price", "MEXBOL"],  # Mexico
        "EWN": ["Indices-Daily-Price", "AEX"],  # Netherlands
        "EWS": ["Indices-Daily-Price", "FSSTI"],  # Singapore
        "EZA": ["Indices-Daily-Price", "TOP40"],  # South Africa
        "EWP": ["Indices-Daily-Price", "IBEX"],  # Spain
        "EWD": ["Indices-Daily-Price", "OMX"],  # Sweden
        "EWL": ["Indices-Daily-Price", "SMI"],  # Switzerland
        "EWT": ["Indices-Daily-Price", "TWSE"],  # Taiwan
        "EWU": ["Indices-Daily-Price", "UKX"],  # United Kingdom
        "GLD": ["Commodities-Daily-Price", "GC1"],
        "GSG": ["Commodities-Daily-Price", "GC1"],
        "IEF": ["Bonds-Daily-Price", "US10Y"],
        "IEV": ["Indices-Daily-Price", "SX5E"],
        "IWB": ["Indices-Daily-Price", "SPX"],
        "QQQ": ["Indices-Daily-Price", "NDX"],
        "SHY": ["Bonds-Daily-Price", "US1Y"],
        "SPY": ["Indices-Daily-Price", "SPX"],
        "THD": ["Indices-Daily-Price", "SET50"],  # Thailand
        "USO": ["Commodities-Daily-Price", "CL1"],
    }
    symbols = df_etfs.symbol.unique()
    mapping = {k: v for k, v in mapping.items() if k in symbols}

    # Nothing to extend → just return the input
    if not mapping:
        return df_etfs.copy()

    # ------------------------------------------------------------------ step 2
    grouped = defaultdict(list)  # {path: [proxy1, proxy2, ...]}
    for _, (path, proxy) in mapping.items():
        grouped[path].append(proxy)

    # Load each dataset only if there's at least one proxy symbol
    other_frames = []
    for path, proxies in grouped.items():
        if proxies:  # skip empty lists
            other_frames.append(load_dataset(path, proxies, to_usd=True))

    # If no proxy data could be loaded, fall back to raw ETF data
    if not other_frames:
        return df_etfs.copy()

    df_others = pd.concat(other_frames, ignore_index=True)

    # ------------------------------------------------------------------ step 3
    frames = []
    for etf, (__, proxy) in mapping.items():
        etf_data = df_etfs[df_etfs["symbol"] == etf]
        proxy_data = df_others[df_others["symbol"] == proxy]

        if etf_data.empty or proxy_data.empty:
            frames.append(etf_data)  # keep raw ETF if proxy missing
            continue

        # Find first overlapping date
        first_common = etf_data.loc[
            etf_data["date"].isin(proxy_data["date"]), "date"
        ].min()
        if pd.isna(first_common):
            frames.append(etf_data)  # no overlap → keep raw ETF
            continue

        # Compute adjustment factor on that date
        k = (
            etf_data.loc[etf_data["date"] == first_common, "close"].iloc[0]
            / proxy_data.loc[proxy_data["date"] == first_common, "close"].iloc[0]
        )

        # Scale proxy history before the overlap
        hist = proxy_data[proxy_data["date"] < first_common].copy()
        hist[["open", "high", "low", "close"]] *= k
        hist["symbol"] = etf

        # Combine proxy history + actual ETF data
        frames.append(pd.concat([hist, etf_data]))

    # Add ETFs that were never in the mapping
    untouched = set(symbols) - set(mapping)
    frames.append(df_etfs[df_etfs["symbol"].isin(untouched)])

    return (
        pd.concat(frames, ignore_index=True)
        .sort_values(["date", "symbol"])
        .reset_index(drop=True)
    )


ALLOWED_FIELDS = {"open", "high", "low", "close", "volume"}


def get_pricing(
    symbol_list,
    fields=None,  # ← default set below
    start_date="1980-01-01",
    end_date=date.today().isoformat(),
    extend=False,
    keep_single_level=False,  # backward-compat flag
):
    """
    Fetch OHLC pricing for the requested symbols.

    Parameters
    ----------
    symbol_list : str | list[str]
        One ticker or a list of tickers. Supports suffixed form "SYM.CAT"
        where CAT in {CRYPTO, STOCK, CMDTY, IDX, ETF, FX, BND}.
    fields : list[str] | None
        Any subset of ["open", "high", "low", "close"].
        Defaults to ["close"] for backward compatibility.
    start_date, end_date : str (YYYY-MM-DD)
        Slice the date index (inclusive).
    extend : bool
        Pass-through to `load_dataset(..., extend=extend)`.
    keep_single_level : bool
        If `True` and only one field is requested, flatten the columns so the
        output matches the old behaviour (columns = symbols).  If `False`
        you always get a two-level MultiIndex `(symbol, field)`.

    Returns
    -------
    pd.DataFrame
        * MultiIndex columns (symbol, field) when `len(fields) > 1`
        * Single-level columns     (symbol)      when one field & keep_single_level=True
    """
    if fields is None:
        fields = ["close"]
    if isinstance(symbol_list, str):
        symbol_list = [symbol_list]
    if isinstance(symbol_list, list) and "sp500" in symbol_list:
        symbol_list.remove("sp500")
        symbol_list += SP500_SYMBOLS

    fields = [f.lower() for f in fields]
    bad = [f for f in fields if f not in ALLOWED_FIELDS]
    if bad:
        raise ValueError(f"Invalid field(s): {bad}. Allowed: {sorted(ALLOWED_FIELDS)}")

    # --------------------------------------------------------------- download
    # Direct routing for suffixed symbols (skip Universe lookup)
    SUFFIX_TO_REPO = {
        "CRP": "paperswithbacktest/Cryptocurrencies-Daily-Price",
        "STK": "paperswithbacktest/Stocks-Daily-Price",
        "CMD": "paperswithbacktest/Commodities-Daily-Price",
        "IDX": "paperswithbacktest/Indices-Daily-Price",
        "ETF": "paperswithbacktest/ETFs-Daily-Price",
        "FXR": "paperswithbacktest/Forex-Daily-Price",
        "BND": "paperswithbacktest/Bonds-Daily-Price",
    }

    universe = ds.load_dataset(
        "paperswithbacktest/Universe-Daily-Price",
        token=_get_hf_token(),
    )
    mapping = universe["train"].to_pandas()
    mapping = mapping.set_index("symbol")["repo_id"].to_dict()

    grouped = defaultdict(list)
    unmapped = []

    for sym in symbol_list:
        # Case 1: suffixed like "BTC.CRYPTO" → route directly, use base symbol for the dataset
        if "." in sym:
            base, cat = sym.rsplit(".", 1)
            repo_full = SUFFIX_TO_REPO.get(cat.upper())
            if repo_full:
                repo_id = repo_full.split("/")[1]  # dataset name only
                grouped[repo_id].append(base)  # dataset expects base ticker
                continue
            # unknown suffix → fall back to Universe mapping on the full symbol

        # Case 2: normal symbol → use Universe mapping
        repo_full = mapping.get(sym)
        repo_id = repo_full.split("/")[1] if isinstance(repo_full, str) else None
        if repo_id:
            grouped[repo_id].append(sym)
        else:
            print(f"Warning: No dataset found for symbol '{sym}'")
            unmapped.append(sym)

    frames = []
    for repo_id, syms in grouped.items():
        ext_flag = extend if repo_id != "Indices-Daily-Price" else False
        df_part = load_dataset(repo_id, syms, extend=ext_flag)
        if not df_part.empty:
            frames.append(df_part)

    df = pd.concat(frames, ignore_index=True) if frames else pd.DataFrame()

    df["date"] = pd.to_datetime(df["date"])
    df.set_index("date", inplace=True)
    df.sort_index(inplace=True)
    df = df.loc[start_date:end_date]

    # ------------------------------------------------------------- reshape
    prices = df.pivot_table(values=fields, index=df.index, columns="symbol")
    prices = prices.swaplevel(axis=1).sort_index(axis=1)

    if keep_single_level:
        if isinstance(prices.columns, pd.MultiIndex):
            prices.columns = [
                col if not isinstance(col, tuple) else col[-1] for col in prices.columns
            ]

    return prices
