import polars as pl
from datetime import date

from ocf.data.normalize import *
from ocf.data.align import align_underlying_daily
from ocf.data.enrich import enrich_option_chain_snapshot
from ocf.greeks.api import (
    per_option_greeks,
    exposure_features,
    option_chain_greeks,
)


def test_public_greeks_api_from_raw():
    """
    Public API test for Greeks.

    Verifies:
    - per_option_greeks works
    - exposure_features works
    - option_chain_greeks orchestrates correctly
    """

    raw_dir = "data/raw/bloomberg/spy"
    symbol = "SPY"
    snapshot_date = date(2024, 1, 19)

    # Load raw CSVs
    raw = {
        "ohlcv": pl.read_csv(f"{raw_dir}/ohlcv.csv"),
        "oi": pl.read_csv(f"{raw_dir}/oi.csv"),
        "iv30": pl.read_csv(f"{raw_dir}/ivol_30_mny.csv"),
        "iv60": pl.read_csv(f"{raw_dir}/ivol_60_mny.csv"),
        "sofr": pl.read_csv(f"{raw_dir}/riskfree_sofr.csv"),
        "usgg": pl.read_csv(f"{raw_dir}/riskfree_usgg.csv"),
        "vix": pl.read_csv(f"{raw_dir}/vix.csv"),
        "chain": pl.read_csv(f"{raw_dir}/chain.csv"),
    }

    # Normalize
    normalized = {
        "ohlcv": normalize_ohlcv(raw["ohlcv"], symbol),
        "oi": normalize_oi(raw["oi"], symbol),
        "iv30": normalize_ivol_30(raw["iv30"], symbol),
        "iv60": normalize_ivol_60(raw["iv60"], symbol),
        "sofr": normalize_sofr(raw["sofr"], symbol),
        "usgg": normalize_usgg(raw["usgg"], symbol),
        "vix": normalize_vix(raw["vix"], symbol),
        "chain": normalize_chain(
            raw["chain"],
            symbol=symbol,
            snapshot_date=snapshot_date,
        ),
    }

    # Align underlying (EXPLICIT inputs only)
    underlying = align_underlying_daily(
        symbol=symbol,
        ohlcv=normalized["ohlcv"],
        oi=normalized["oi"],
        iv30=normalized["iv30"],
        iv60=normalized["iv60"],
        sofr=normalized["sofr"],
        usgg=normalized["usgg"],
        vix=normalized["vix"],
    )

    underlying_row = (
        underlying
        .sort("date")
        .tail(1)
        .to_dicts()[0]
    )

    # Enrich option chain
    chain = enrich_option_chain_snapshot(
        chain=normalized["chain"],
        underlying_row=underlying_row,
    )

    # 1. Per-option Greeks
    greeks = per_option_greeks(
        chain_snapshot=chain,
        underlying_row=underlying_row,
    )
    assert not greeks.is_empty()

    # 2. Exposure features
    exposure = exposure_features(greeks)
    assert exposure.height == 1

    # 3. Composite API
    out = option_chain_greeks(
        chain_snapshot=chain,
        underlying_row=underlying_row,
    )

    assert "per_option_greeks" in out
    assert "exposure_features" in out