from __future__ import annotations

import numpy as np
import pandas as pd


def _calls_for_expiry(options: pd.DataFrame, expiry: pd.Timestamp) -> pd.DataFrame:
    c = options[
        (options["option_type"] == "call")
        & (options["expiry"] == pd.to_datetime(expiry))
    ].copy()
    c = c.sort_values("strike").drop_duplicates(subset=["strike"], keep="last")
    if "openInterest" not in c.columns:
        c["openInterest"] = 0.0
    return c


def _puts_for_expiry(options: pd.DataFrame, expiry: pd.Timestamp) -> pd.DataFrame:
    p = options[
        (options["option_type"] == "put")
        & (options["expiry"] == pd.to_datetime(expiry))
    ].copy()
    p = p.sort_values("strike").drop_duplicates(subset=["strike"], keep="last")
    if "openInterest" not in p.columns:
        p["openInterest"] = 0.0
    return p


def scan_vertical_arbitrage(
    calls_df: pd.DataFrame,
    tol: float = 1e-8,
    min_edge: float = 0.0,
    min_edge_per_width: float = 0.0,
    min_leg_open_interest: int = 0,
) -> pd.DataFrame:
    g = calls_df.sort_values("strike")
    if "openInterest" not in g.columns:
        g = g.copy()
        g["openInterest"] = 0.0
    k = g["strike"].to_numpy(dtype=float)
    c = g["mid"].to_numpy(dtype=float)
    oi = g["openInterest"].fillna(0.0).to_numpy(dtype=float)
    rows: list[dict[str, object]] = []
    for i in range(len(k) - 1):
        edge = c[i + 1] - c[i]
        if edge > max(tol, min_edge):
            width = float(k[i + 1] - k[i])
            edge_pw = float(edge / max(width, 1e-12))
            leg_oi_min = float(min(oi[i], oi[i + 1]))
            if edge_pw < min_edge_per_width:
                continue
            if leg_oi_min < float(min_leg_open_interest):
                continue
            rows.append(
                {
                    "family": "vertical",
                    "k1": float(k[i]),
                    "k2": float(k[i + 1]),
                    "k3": np.nan,
                    "edge": float(edge),
                    "edge_per_width": edge_pw,
                    "leg_oi_min": leg_oi_min,
                    "notes": "Call should not increase with strike.",
                }
            )
    return pd.DataFrame(rows)


def scan_butterfly_arbitrage(
    calls_df: pd.DataFrame,
    tol: float = 1e-8,
    min_edge: float = 0.0,
    min_edge_per_width: float = 0.0,
    min_leg_open_interest: int = 0,
) -> pd.DataFrame:
    g = calls_df.sort_values("strike")
    if "openInterest" not in g.columns:
        g = g.copy()
        g["openInterest"] = 0.0
    k = g["strike"].to_numpy(dtype=float)
    c = g["mid"].to_numpy(dtype=float)
    oi = g["openInterest"].fillna(0.0).to_numpy(dtype=float)
    rows: list[dict[str, object]] = []
    if len(k) < 3:
        return pd.DataFrame(rows)
    slope_left = (c[1:-1] - c[:-2]) / (k[1:-1] - k[:-2])
    slope_right = (c[2:] - c[1:-1]) / (k[2:] - k[1:-1])
    mismatch = slope_left - slope_right
    for i, mm in enumerate(mismatch, start=1):
        if mm > tol:
            width = float(k[i + 1] - k[i - 1])
            edge_pw = float(mm / max(width, 1e-12))
            leg_oi_min = float(min(oi[i - 1], oi[i], oi[i + 1]))
            if mm < min_edge:
                continue
            if edge_pw < min_edge_per_width:
                continue
            if leg_oi_min < float(min_leg_open_interest):
                continue
            rows.append(
                {
                    "family": "butterfly",
                    "k1": float(k[i - 1]),
                    "k2": float(k[i]),
                    "k3": float(k[i + 1]),
                    "edge": float(mm),
                    "edge_per_width": edge_pw,
                    "leg_oi_min": leg_oi_min,
                    "notes": "Call slope decreases across strikes (convexity violation).",
                }
            )
    return pd.DataFrame(rows)


def scan_arbitrage_candidates(
    options: pd.DataFrame,
    expiry: pd.Timestamp,
    tol: float = 1e-8,
    min_edge: float = 0.0,
    min_edge_per_width: float = 0.0,
    min_leg_open_interest: int = 0,
    spot: float | None = None,
    r: float = 0.0,
    t_years: float = 0.25,
) -> pd.DataFrame:
    calls = _calls_for_expiry(options, expiry)
    puts = _puts_for_expiry(options, expiry)
    if calls.empty:
        return pd.DataFrame(
            columns=[
                "expiry",
                "family",
                "k1",
                "k2",
                "k3",
                "edge",
                "edge_per_width",
                "leg_oi_min",
                "notes",
                "confidence",
            ]
        )

    v = scan_vertical_arbitrage(
        calls,
        tol=tol,
        min_edge=min_edge,
        min_edge_per_width=min_edge_per_width,
        min_leg_open_interest=min_leg_open_interest,
    )
    b = scan_butterfly_arbitrage(
        calls,
        tol=tol,
        min_edge=min_edge,
        min_edge_per_width=min_edge_per_width,
        min_leg_open_interest=min_leg_open_interest,
    )

    parity = scan_put_call_parity_arbitrage(
        calls,
        puts,
        spot=spot,
        r=r,
        t_years=t_years,
        tol=tol,
        min_edge=min_edge,
        min_leg_open_interest=min_leg_open_interest,
    )

    calendar = scan_calendar_arbitrage(
        options,
        tol=tol,
        min_edge=min_edge,
        min_leg_open_interest=min_leg_open_interest,
    )

    out = (
        pd.concat([v, b, parity, calendar], ignore_index=True)
        if not v.empty or not b.empty or not parity.empty or not calendar.empty
        else pd.DataFrame()
    )
    if out.empty:
        return pd.DataFrame(
            columns=[
                "expiry",
                "family",
                "k1",
                "k2",
                "k3",
                "edge",
                "edge_per_width",
                "leg_oi_min",
                "notes",
                "confidence",
            ]
        )
    out.insert(0, "expiry", pd.to_datetime(expiry))
    out["confidence"] = "unrated"
    return out.sort_values("edge", ascending=False).reset_index(drop=True)


def scan_put_call_parity_arbitrage(
    calls_df: pd.DataFrame,
    puts_df: pd.DataFrame,
    spot: float | None,
    r: float,
    t_years: float,
    tol: float = 1e-8,
    min_edge: float = 0.0,
    min_leg_open_interest: int = 0,
) -> pd.DataFrame:
    if spot is None:
        return pd.DataFrame()
    c = calls_df[["strike", "mid", "openInterest"]].rename(
        columns={"mid": "call_mid", "openInterest": "call_oi"}
    )
    p = puts_df[["strike", "mid", "openInterest"]].rename(
        columns={"mid": "put_mid", "openInterest": "put_oi"}
    )
    m = c.merge(p, on="strike", how="inner").sort_values("strike")
    if m.empty:
        return pd.DataFrame()
    disc = float(np.exp(-float(r) * float(t_years)))
    rows: list[dict[str, object]] = []
    for _, row in m.iterrows():
        k = float(row["strike"])
        lhs = float(row["call_mid"] - row["put_mid"])
        rhs = float(spot - disc * k)
        resid = lhs - rhs
        edge = abs(resid)
        leg_oi_min = float(min(row["call_oi"], row["put_oi"]))
        if edge <= max(tol, min_edge):
            continue
        if leg_oi_min < float(min_leg_open_interest):
            continue
        rows.append(
            {
                "family": "parity",
                "k1": k,
                "k2": np.nan,
                "k3": np.nan,
                "edge": edge,
                "edge_per_width": edge,
                "leg_oi_min": leg_oi_min,
                "notes": "Put-call parity residual (American/dividend caveat applies).",
            }
        )
    return pd.DataFrame(rows)


def scan_calendar_arbitrage(
    options: pd.DataFrame,
    tol: float = 1e-8,
    min_edge: float = 0.0,
    min_leg_open_interest: int = 0,
) -> pd.DataFrame:
    calls = options[options["option_type"] == "call"].copy()
    if calls.empty:
        return pd.DataFrame()
    if "openInterest" not in calls.columns:
        calls["openInterest"] = 0.0

    rows: list[dict[str, object]] = []
    for strike, grp in calls.groupby("strike"):
        g = grp.sort_values("expiry")
        if len(g) < 2:
            continue
        mids = g["mid"].to_numpy(dtype=float)
        expiries = pd.to_datetime(g["expiry"]).to_numpy()
        ois = g["openInterest"].fillna(0.0).to_numpy(dtype=float)
        for i in range(len(mids) - 1):
            edge = float(mids[i] - mids[i + 1])
            if edge <= max(tol, min_edge):
                continue
            leg_oi_min = float(min(ois[i], ois[i + 1]))
            if leg_oi_min < float(min_leg_open_interest):
                continue
            rows.append(
                {
                    "family": "calendar",
                    "k1": float(strike),
                    "k2": np.nan,
                    "k3": np.nan,
                    "edge": edge,
                    "edge_per_width": edge,
                    "leg_oi_min": leg_oi_min,
                    "notes": "Longer-dated call cheaper than shorter-dated call at same strike.",
                }
            )
    return pd.DataFrame(rows)


def summarize_arbitrage(candidates: pd.DataFrame) -> pd.DataFrame:
    if candidates.empty:
        return pd.DataFrame(
            [
                {
                    "candidate_count": 0,
                    "max_edge": 0.0,
                    "median_edge": 0.0,
                    "vertical_count": 0,
                    "butterfly_count": 0,
                    "parity_count": 0,
                    "calendar_count": 0,
                    "high_conf_count": 0,
                    "medium_conf_count": 0,
                    "low_conf_count": 0,
                }
            ]
        )
    return pd.DataFrame(
        [
            {
                "candidate_count": int(len(candidates)),
                "max_edge": float(candidates["edge"].max()),
                "median_edge": float(candidates["edge"].median()),
                "vertical_count": int((candidates["family"] == "vertical").sum()),
                "butterfly_count": int((candidates["family"] == "butterfly").sum()),
                "parity_count": int((candidates["family"] == "parity").sum()),
                "calendar_count": int((candidates["family"] == "calendar").sum()),
                "high_conf_count": int((candidates["confidence"] == "high").sum()),
                "medium_conf_count": int((candidates["confidence"] == "medium").sum()),
                "low_conf_count": int((candidates["confidence"] == "low").sum()),
            }
        ]
    )


def assign_candidate_confidence(
    candidates: pd.DataFrame,
    mean_violation_rate: float,
    failed_checks: int,
) -> pd.DataFrame:
    if candidates.empty:
        return candidates
    if failed_checks == 0 and mean_violation_rate <= 0.02:
        conf = "high"
    elif failed_checks <= 1 and mean_violation_rate <= 0.08:
        conf = "medium"
    else:
        conf = "low"
    out = candidates.copy()
    out["confidence"] = conf
    return out