src/option_implied_lab/pipeline.py

from __future__ import annotations

from datetime import datetime

import pandas as pd

from .arbitrage import (
    assign_candidate_confidence,
    scan_arbitrage_candidates,
    summarize_arbitrage,
)
from .data import fetch_option_snapshot
from .density import estimate_rn_density, select_density_slice
from .features import compute_features
from .no_arb import run_all_checks, run_checks_by_expiry, select_best_quality_expiry
from .strategy import generate_candidates, score_candidates, select_best
from .types import OptionSnapshot


def _available_expiries(snapshot: OptionSnapshot) -> pd.DataFrame:
    expiries = (
        snapshot.options[snapshot.options["option_type"] == "call"]["expiry"]
        .dropna()
        .drop_duplicates()
        .sort_values()
    )
    now = pd.Timestamp(snapshot.snapshot_time).tz_localize(None)
    df = pd.DataFrame({"expiry": pd.to_datetime(expiries)})
    df["dte_days"] = (df["expiry"] - now).dt.days
    return df.reset_index(drop=True)


def _select_expiry_by_target_dte(snapshot: OptionSnapshot, target_dte: int) -> datetime:
    exp = _available_expiries(snapshot)
    if exp.empty:
        raise ValueError("No expiries available")
    valid = exp[exp["dte_days"] >= 0].copy()
    if valid.empty:
        valid = exp.copy()
    valid["dist"] = (valid["dte_days"] - int(target_dte)).abs()
    chosen = valid.sort_values(["dist", "dte_days"]).iloc[0]["expiry"]
    return pd.to_datetime(chosen).to_pydatetime()


def analyze_snapshot(
    snapshot: OptionSnapshot,
    expiry: datetime | None = None,
    moneyness_band: float = 0.2,
    min_open_interest: int = 1,
    min_volume: int = 0,
    smooth_window: int = 3,
    risk_lambda: float = 0.5,
    arb_min_edge: float = 0.0,
    arb_min_edge_per_width: float = 0.0,
    arb_min_leg_open_interest: int = 0,
) -> dict[str, object]:
    diagnostics = run_all_checks(snapshot)
    per_expiry_diag = run_checks_by_expiry(snapshot)
    available_expiries = _available_expiries(snapshot)
    diag_df = pd.DataFrame(
        [
            {
                "check": d.name,
                "passed": d.passed,
                "violations": d.violations,
                "comparisons": d.comparisons,
                "violation_rate": d.violation_rate,
                "details": d.details,
            }
            for d in diagnostics
        ]
    )

    per_expiry_summary = (
        per_expiry_diag.groupby("expiry", as_index=False)
        .agg(
            failed_checks=("passed", lambda s: int((~s).sum())),
            total_checks=("passed", "count"),
            mean_violation_rate=("violation_rate", "mean"),
        )
        .sort_values("mean_violation_rate")
        .reset_index(drop=True)
    )

    if not available_expiries.empty:
        per_expiry_summary = per_expiry_summary.merge(
            available_expiries, on="expiry", how="left"
        )

    if expiry is None:
        expiry = select_best_quality_expiry(snapshot)

    strikes, calls, selected_expiry = select_density_slice(
        snapshot.options,
        spot=snapshot.spot,
        expiry=expiry,
        moneyness_band=moneyness_band,
        min_open_interest=min_open_interest,
        min_volume=min_volume,
    )
    density = estimate_rn_density(
        strikes=strikes,
        call_prices=calls,
        expiry=selected_expiry,
        smooth_window=smooth_window,
    )

    features = compute_features(snapshot, density)
    candidates = generate_candidates(snapshot.spot, density, snapshot.options)
    scored = score_candidates(
        candidates, density=density, spot0=snapshot.spot, risk_lambda=risk_lambda
    )
    best = select_best(scored)

    density_df = pd.DataFrame(
        {
            "strike": density.strikes,
            "density": density.density,
        }
    )
    scored_df = pd.DataFrame(
        [
            {
                "strategy": s.candidate.name,
                "expected_payoff": s.expected_payoff,
                "downside_q05": s.downside_q05,
                "objective": s.objective,
            }
            for s in scored
        ]
    )

    arbitrage_candidates = scan_arbitrage_candidates(
        snapshot.options,
        expiry=pd.to_datetime(selected_expiry),
        min_edge=arb_min_edge,
        min_edge_per_width=arb_min_edge_per_width,
        min_leg_open_interest=arb_min_leg_open_interest,
        spot=snapshot.spot,
    )
    exp_quality = per_expiry_summary[
        per_expiry_summary["expiry"] == pd.to_datetime(selected_expiry)
    ]
    if exp_quality.empty:
        mean_violation_rate = (
            float(diag_df["violation_rate"].mean()) if not diag_df.empty else 1.0
        )
        failed_checks = int((~diag_df["passed"]).sum()) if not diag_df.empty else 3
    else:
        mean_violation_rate = float(exp_quality.iloc[0]["mean_violation_rate"])
        failed_checks = int(exp_quality.iloc[0]["failed_checks"])
    arbitrage_candidates = assign_candidate_confidence(
        arbitrage_candidates,
        mean_violation_rate=mean_violation_rate,
        failed_checks=failed_checks,
    )
    arbitrage_summary = summarize_arbitrage(arbitrage_candidates)

    return {
        "ticker": snapshot.ticker,
        "spot": snapshot.spot,
        "snapshot_time": snapshot.snapshot_time,
        "selected_expiry": selected_expiry,
        "diagnostics": diag_df,
        "diagnostics_by_expiry": per_expiry_diag,
        "diagnostics_by_expiry_summary": per_expiry_summary,
        "available_expiries": available_expiries,
        "density": density_df,
        "features": features,
        "scored": scored_df,
        "best_strategy": {
            "name": best.candidate.name,
            "params": best.candidate.params,
            "expected_payoff": best.expected_payoff,
            "downside_q05": best.downside_q05,
            "objective": best.objective,
        },
        "arbitrage_candidates": arbitrage_candidates,
        "arbitrage_summary": arbitrage_summary,
    }


def analyze_ticker(
    ticker: str,
    max_expiries: int = 2,
    expiry: datetime | None = None,
    moneyness_band: float = 0.2,
    min_open_interest: int = 1,
    min_volume: int = 0,
    smooth_window: int = 3,
    risk_lambda: float = 0.5,
    expiry_mode: str = "auto",
    target_dte: int = 30,
    arb_min_edge: float = 0.0,
    arb_min_edge_per_width: float = 0.0,
    arb_min_leg_open_interest: int = 0,
) -> dict[str, object]:
    snapshot = fetch_option_snapshot(ticker=ticker, max_expiries=max_expiries)
    selected_expiry = expiry
    if expiry_mode == "target_dte":
        selected_expiry = _select_expiry_by_target_dte(snapshot, target_dte)
    elif expiry_mode == "auto":
        selected_expiry = expiry
    elif expiry_mode == "manual":
        selected_expiry = expiry
    else:
        raise ValueError(f"Unknown expiry_mode: {expiry_mode}")

    return analyze_snapshot(
        snapshot=snapshot,
        expiry=selected_expiry,
        moneyness_band=moneyness_band,
        min_open_interest=min_open_interest,
        min_volume=min_volume,
        smooth_window=smooth_window,
        risk_lambda=risk_lambda,
        arb_min_edge=arb_min_edge,
        arb_min_edge_per_width=arb_min_edge_per_width,
        arb_min_leg_open_interest=arb_min_leg_open_interest,
    )