"""Constructing Continuous Options Series — constant-maturity roll adjustment."""

import plotly.graph_objects as go
import polars as pl
from plotly.subplots import make_subplots

from data import load_sp500_options_eda
from utils.style import COLORS

DEMO_SYMBOL = "AAPL"
DEMO_YEAR = 2019

raw = load_sp500_options_eda(
    symbols=[DEMO_SYMBOL],
    start_date=f"{DEMO_YEAR}-01-01",
    end_date=f"{DEMO_YEAR}-12-31",
).rename({"timestamp": "date"})

print(f"Raw {DEMO_SYMBOL} options ({DEMO_YEAR}): {len(raw):,} rows")
print(f"  Dates: {raw['date'].n_unique()}")
print(
    f"  Unique (strike, expiration): {raw.select(pl.struct('strike', 'expiration')).n_unique():,}"
)

# Selection parameters for constant-maturity straddle
DTE_WINDOW = (25, 35)
TARGET_DELTA = 0.50
DELTA_TOL = 0.15
MIN_BID = 0.01
MAX_REL_SPREAD = 0.30

def select_constant_maturity_straddle(raw_options: pl.DataFrame) -> pl.DataFrame:
    """Select the 'best' 30D ATM straddle per day from raw option chains."""
    opts = raw_options.with_columns(
        pl.col("delta").abs().alias("abs_delta"),
        ((pl.col("ask") - pl.col("bid")) / pl.col("mid_price").clip(lower_bound=0.01)).alias(
            "rel_spread"
        ),
    )
    filtered = opts.filter(
        pl.col("days_to_maturity").is_between(DTE_WINDOW[0], DTE_WINDOW[1])
        & (pl.col("bid") >= MIN_BID)
        & (pl.col("rel_spread") <= MAX_REL_SPREAD)
        & (pl.col("iv_convergence") == "Converged")
        & pl.col("abs_delta").is_between(TARGET_DELTA - DELTA_TOL, TARGET_DELTA + DELTA_TOL)
    )

    def _best_leg(df: pl.DataFrame, cp: str) -> pl.DataFrame:
        leg = df.filter(pl.col("call_put") == cp)
        leg = leg.with_columns((pl.col("abs_delta") - TARGET_DELTA).abs().alias("_dd"))
        return (
            leg.with_columns(pl.col("_dd").rank("ordinal").over("date").alias("_rank"))
            .filter(pl.col("_rank") == 1)
            .drop(["_rank", "_dd"])
        )

    call_cols = [
        "date",
        "strike",
        "expiration",
        "days_to_maturity",
        "underlying_price",
        pl.col("mid_price").alias("call_mid"),
        pl.col("bid").alias("call_bid"),
        pl.col("ask").alias("call_ask"),
        pl.col("delta").alias("call_delta"),
        pl.col("theta").alias("call_theta"),
    ]
    put_cols = [
        "date",
        pl.col("mid_price").alias("put_mid"),
        pl.col("bid").alias("put_bid"),
        pl.col("ask").alias("put_ask"),
        pl.col("delta").alias("put_delta"),
        pl.col("theta").alias("put_theta"),
    ]
    calls = _best_leg(filtered, "C").select(call_cols)
    puts = _best_leg(filtered, "P").select(put_cols)
    return (
        calls.join(puts, on="date", how="inner")
        .with_columns(
            (pl.col("call_mid") + pl.col("put_mid")).alias("instr_mid"),
            (pl.col("call_theta") + pl.col("put_theta")).alias("instr_theta"),
            (pl.col("call_delta") + pl.col("put_delta")).alias("instr_delta"),
        )
        .sort("date")
    )

cm_straddles = select_constant_maturity_straddle(raw)
print(f"Constant-maturity straddle series: {len(cm_straddles)} days")
print(f"  Date range: {cm_straddles['date'].min()} to {cm_straddles['date'].max()}")
cm_straddles.select(["date", "strike", "expiration", "days_to_maturity", "instr_mid"]).head(10)

cm_straddles = cm_straddles.with_columns(
    pl.col("expiration").shift(1).alias("prev_exp"),
    pl.col("strike").shift(1).alias("prev_strike"),
    pl.col("instr_mid").shift(1).alias("prev_mid"),
)

cm_straddles = cm_straddles.with_columns(
    ((pl.col("expiration") != pl.col("prev_exp")) | (pl.col("strike") != pl.col("prev_strike")))
    .fill_null(False)
    .alias("is_roll"),
    ((pl.col("instr_mid") - pl.col("prev_mid")) / pl.col("prev_mid")).alias("daily_return"),
)

n_rolls = cm_straddles.filter(pl.col("is_roll")).height
n_total = len(cm_straddles) - 1  # exclude first row (no return)
print(f"Contract changes: {n_rolls} out of {n_total} days ({n_rolls / n_total:.0%})")

roll_rets = cm_straddles.filter(pl.col("is_roll"))["daily_return"].drop_nulls()
nonroll_rets = cm_straddles.filter(~pl.col("is_roll"))["daily_return"].drop_nulls()

roll_summary = pl.DataFrame(
    {
        "metric": ["count", "mean_return", "median_return", "pct_positive"],
        "roll_days": [
            float(roll_rets.len()),
            roll_rets.mean(),
            roll_rets.median(),
            (roll_rets > 0).mean(),
        ],
        "non_roll_days": [
            float(nonroll_rets.len()),
            nonroll_rets.mean(),
            nonroll_rets.median(),
            (nonroll_rets > 0).mean(),
        ],
    }
)
roll_summary = roll_summary.with_columns(
    (pl.col("roll_days") - pl.col("non_roll_days")).alias("difference"),
)
roll_summary

# Prepare data for sawtooth visualization
dates = cm_straddles["date"].to_list()
prices = cm_straddles["instr_mid"].to_list()
dtes = cm_straddles["days_to_maturity"].to_list()
rolls = cm_straddles["is_roll"].to_list()

roll_dates = [d for d, r in zip(dates, rolls, strict=False) if r]
roll_prices = [p for p, r in zip(prices, rolls, strict=False) if r]
nonroll_dates = [d for d, r in zip(dates, rolls, strict=False) if not r]
nonroll_prices = [p for p, r in zip(prices, rolls, strict=False) if not r]

# Build both panels in a single cell so the inline backend does not flush
# the figure mid-construction.
fig = make_subplots(
    rows=2,
    cols=1,
    shared_xaxes=True,
    subplot_titles=[
        f"{DEMO_SYMBOL} Constant-Maturity Straddle Price",
        "Days to Expiration (DTE)",
    ],
    vertical_spacing=0.08,
)

fig.add_trace(
    go.Scatter(
        x=nonroll_dates,
        y=nonroll_prices,
        mode="markers",
        name="Same contract",
        marker=dict(size=4, color=COLORS["blue"]),
        opacity=0.7,
    ),
    row=1,
    col=1,
)
fig.add_trace(
    go.Scatter(
        x=roll_dates,
        y=roll_prices,
        mode="markers",
        name="Contract switch",
        marker=dict(size=6, color=COLORS["negative"], symbol="diamond"),
        opacity=0.9,
    ),
    row=1,
    col=1,
)
fig.add_trace(
    go.Scatter(
        x=dates,
        y=prices,
        mode="lines",
        name="Price",
        line=dict(width=1, color=COLORS["neutral"]),
        showlegend=False,
    ),
    row=1,
    col=1,
)
fig.add_trace(
    go.Scatter(
        x=dates,
        y=dtes,
        mode="lines+markers",
        name="DTE",
        marker=dict(size=3, color=COLORS["positive"]),
        line=dict(width=1),
    ),
    row=2,
    col=1,
)

fig.update_yaxes(title_text="Straddle Mid ($)", row=1, col=1)
fig.update_yaxes(title_text="DTE", row=2, col=1)
fig.update_xaxes(title_text="Date", row=2, col=1)
fig.update_layout(height=600, legend=dict(x=0.01, y=0.99))
fig.show()

HOLDING_PERIOD = 10  # days


def build_exit_lookup(raw_options: pl.DataFrame) -> pl.DataFrame:
    """Build a lookup table: (date, strike, expiration, call_put) → mid_price.

    Used to find the exit price of a specific contract h days after entry.
    """
    return raw_options.filter(
        pl.col("bid") >= MIN_BID,
        pl.col("iv_convergence") == "Converged",
    ).select(
        [
            "date",
            "strike",
            "expiration",
            "call_put",
            "mid_price",
            "bid",
            "ask",
            "delta",
            "theta",
            "days_to_maturity",
        ]
    )

# Build lookup from the raw chain
lookup = build_exit_lookup(raw)
print(f"Exit price lookup: {len(lookup):,} rows")
print(
    f"  Unique contracts: {lookup.select(pl.struct('strike', 'expiration', 'call_put')).n_unique():,}"
)

# Get the trading calendar (ordered dates)
trading_dates = cm_straddles["date"].unique().sort().to_list()
date_to_idx = {d: i for i, d in enumerate(trading_dates)}


def get_exit_date(entry_date, h: int, cal: list):
    """Get the trading date h business days after entry_date."""
    idx = date_to_idx.get(entry_date)
    if idx is None or idx + h >= len(cal):
        return None
    return cal[idx + h]


# Build entry-exit pairs
entries = cm_straddles.select(
    [
        "date",
        "strike",
        "expiration",
        "instr_mid",
        "call_mid",
        "put_mid",
        "instr_delta",
        "is_roll",
    ]
).rename({"date": "entry_date", "instr_mid": "entry_mid"})

# Add exit dates
exit_dates = [
    get_exit_date(d, HOLDING_PERIOD, trading_dates) for d in entries["entry_date"].to_list()
]
entries = entries.with_columns(pl.Series("exit_date", exit_dates).cast(pl.Date))

# Drop entries where exit date is beyond our data
entries = entries.filter(pl.col("exit_date").is_not_null())
print(f"Entry-exit pairs to look up: {len(entries):,}")

# Look up call exit price
call_exit = lookup.filter(pl.col("call_put") == "C").select(
    [
        pl.col("date").alias("exit_date"),
        "strike",
        "expiration",
        pl.col("mid_price").alias("call_exit_mid"),
    ]
)

# Look up put exit price
put_exit = lookup.filter(pl.col("call_put") == "P").select(
    [
        pl.col("date").alias("exit_date"),
        "strike",
        "expiration",
        pl.col("mid_price").alias("put_exit_mid"),
    ]
)

# Join: entry contract → exit prices for same (strike, expiration)
same_contract = entries.join(call_exit, on=["exit_date", "strike", "expiration"], how="left").join(
    put_exit, on=["exit_date", "strike", "expiration"], how="left"
)

# Compute same-contract straddle exit mid
same_contract = same_contract.with_columns(
    (pl.col("call_exit_mid") + pl.col("put_exit_mid")).alias("exit_mid"),
)

# Same-contract return (short straddle: positive = profitable)
same_contract = same_contract.with_columns(
    ((pl.col("entry_mid") - pl.col("exit_mid")) / pl.col("entry_mid")).alias("same_contract_ret"),
)

# How many lookups succeeded?
found = same_contract.filter(pl.col("exit_mid").is_not_null()).height
print(f"Exit prices found: {found} / {len(same_contract)} ({found / len(same_contract):.1%})")

# Naive return: shift-based on constant-maturity series (what 02_labels.py did)
naive_rets = cm_straddles.with_columns(
    pl.col("instr_mid").shift(-1).alias("naive_entry"),
    pl.col("instr_mid").shift(-(1 + HOLDING_PERIOD)).alias("naive_exit"),
).with_columns(
    ((pl.col("naive_entry") - pl.col("naive_exit")) / pl.col("naive_entry")).alias("naive_ret"),
)

# Align for comparison
comparison = (
    same_contract.filter(pl.col("same_contract_ret").is_not_null())
    .select(["entry_date", "same_contract_ret"])
    .join(
        naive_rets.select([pl.col("date").alias("entry_date"), "naive_ret"]).drop_nulls(),
        on="entry_date",
        how="inner",
    )
)

print(f"Paired observations: {len(comparison):,}")

naive_vs_same = pl.DataFrame(
    {
        "metric": ["mean_return", "std", "median_return", "pct_positive_short"],
        "naive_chained": [
            comparison["naive_ret"].mean(),
            comparison["naive_ret"].std(),
            comparison["naive_ret"].median(),
            (comparison["naive_ret"] > 0).mean(),
        ],
        "same_contract": [
            comparison["same_contract_ret"].mean(),
            comparison["same_contract_ret"].std(),
            comparison["same_contract_ret"].median(),
            (comparison["same_contract_ret"] > 0).mean(),
        ],
    }
)
corr = comparison.select(pl.corr("naive_ret", "same_contract_ret").alias("correlation")).item()
print(f"Correlation(naive, same-contract) = {corr:.3f}")
naive_vs_same

# Scatter plot: naive vs same-contract
fig = go.Figure()
fig.add_trace(
    go.Scatter(
        x=comparison["naive_ret"].to_list(),
        y=comparison["same_contract_ret"].to_list(),
        mode="markers",
        marker=dict(size=3, color=COLORS["blue"], opacity=0.4),
        name="Returns",
    )
)
fig.add_trace(
    go.Scatter(
        x=[-0.5, 0.5],
        y=[-0.5, 0.5],
        mode="lines",
        line=dict(dash="dash", color=COLORS["neutral"]),
        name="45° line",
    )
)
fig.update_layout(
    title=f"{DEMO_SYMBOL}: Naive vs Same-Contract {HOLDING_PERIOD}d Returns",
    xaxis_title="Naive (chained constant-maturity)",
    yaxis_title="Same-contract holding return",
    width=600,
    height=500,
)
fig.show()

def build_continuous_straddle_series(cm_series: pl.DataFrame) -> pl.DataFrame:
    """Build a continuous straddle price series from within-contract returns.

    On roll days, the return is set to 0 (roll is a transaction, not a return).
    The resulting series has no phantom jumps and its returns reflect actual P&L.
    """
    df = cm_series.with_columns(
        ((pl.col("instr_mid") - pl.col("instr_mid").shift(1)) / pl.col("instr_mid").shift(1)).alias(
            "raw_daily_ret"
        ),
    )

    # On roll days, set return to 0
    df = df.with_columns(
        pl.when(pl.col("is_roll"))
        .then(0.0)
        .otherwise(pl.col("raw_daily_ret"))
        .alias("clean_daily_ret"),
    )

    # Reconstruct price from cumulative returns, starting at the first observed price
    start_price = cm_series["instr_mid"][0]
    cum_rets = df["clean_daily_ret"].fill_null(0.0).to_list()

    adj_prices = [start_price]
    for r in cum_rets[1:]:
        adj_prices.append(adj_prices[-1] * (1.0 + r))

    return df.with_columns(pl.Series("price_adjusted", adj_prices))

adjusted = build_continuous_straddle_series(cm_straddles)

# Compare roll-day returns: raw vs cleaned
adj_roll = adjusted.filter(pl.col("is_roll"))
adj_nonroll = adjusted.filter(~pl.col("is_roll"))

raw_vs_clean = pl.DataFrame(
    {
        "metric": ["raw_mean", "clean_mean"],
        "roll_days": [
            adj_roll["raw_daily_ret"].mean(),
            adj_roll["clean_daily_ret"].mean(),
        ],
        "non_roll_days": [
            adj_nonroll["raw_daily_ret"].mean(),
            adj_nonroll["clean_daily_ret"].mean(),
        ],
    }
)
print(
    f"Roll-day bias removed: raw roll mean was {adj_roll['raw_daily_ret'].mean():+.4f}, now 0.0000"
)
raw_vs_clean

# Visualize raw vs adjusted price series
fig = make_subplots(
    rows=2,
    cols=1,
    shared_xaxes=True,
    subplot_titles=["Raw Constant-Maturity Series (sawtooth)", "Continuous (roll-zeroed)"],
    vertical_spacing=0.08,
)

fig.add_trace(
    go.Scatter(
        x=adjusted["date"].to_list(),
        y=adjusted["instr_mid"].to_list(),
        mode="lines",
        name="Raw",
        line=dict(color=COLORS["negative"], width=1.5),
    ),
    row=1,
    col=1,
)
fig.add_trace(
    go.Scatter(
        x=adjusted["date"].to_list(),
        y=adjusted["price_adjusted"].to_list(),
        mode="lines",
        name="Adjusted",
        line=dict(color=COLORS["blue"], width=1.5),
    ),
    row=2,
    col=1,
)

fig.update_yaxes(title_text="Straddle Mid ($)", row=1, col=1)
fig.update_yaxes(title_text="Adjusted Price ($)", row=2, col=1)
fig.update_layout(
    height=500,
    title=f"{DEMO_SYMBOL} Straddle: Raw vs Continuous (Roll-Adjusted)",
)
fig.show()

# Final comparison: all three 10d forward return approaches
# Naive
adjusted = adjusted.with_columns(
    pl.col("instr_mid").shift(-1).alias("raw_entry"),
    pl.col("instr_mid").shift(-(1 + HOLDING_PERIOD)).alias("raw_exit"),
).with_columns(
    ((pl.col("raw_entry") - pl.col("raw_exit")) / pl.col("raw_entry")).alias("naive_fwd_ret"),
)

# Continuous-adjusted
adjusted = adjusted.with_columns(
    pl.col("price_adjusted").shift(-1).alias("adj_entry"),
    pl.col("price_adjusted").shift(-(1 + HOLDING_PERIOD)).alias("adj_exit"),
).with_columns(
    ((pl.col("adj_entry") - pl.col("adj_exit")) / pl.col("adj_entry")).alias("cont_fwd_ret"),
)

# Merge with same-contract returns
comparison_all = (
    adjusted.select([pl.col("date").alias("entry_date"), "naive_fwd_ret", "cont_fwd_ret"])
    .join(
        same_contract.select(["entry_date", "same_contract_ret"]),
        on="entry_date",
        how="left",
    )
    .drop_nulls()
)

print(f"Paired observations: {len(comparison_all):,}")

three_way = pl.DataFrame(
    {
        "metric": ["mean", "std", "corr_with_same_contract"],
        "naive": [
            comparison_all["naive_fwd_ret"].mean(),
            comparison_all["naive_fwd_ret"].std(),
            comparison_all.select(pl.corr("naive_fwd_ret", "same_contract_ret")).item(),
        ],
        "continuous": [
            comparison_all["cont_fwd_ret"].mean(),
            comparison_all["cont_fwd_ret"].std(),
            comparison_all.select(pl.corr("cont_fwd_ret", "same_contract_ret")).item(),
        ],
        "same_contract": [
            comparison_all["same_contract_ret"].mean(),
            comparison_all["same_contract_ret"].std(),
            1.0,
        ],
    }
)
three_way