from __future__ import annotations

import json
import sys
import warnings
from pathlib import Path

import gradio as gr
import matplotlib.pyplot as plt
import pandas as pd

ROOT = Path(__file__).resolve().parent
SRC = ROOT / "src"
if str(SRC) not in sys.path:
    sys.path.insert(0, str(SRC))

from option_implied_lab.pipeline_v2 import analyze_ticker_v2
from option_implied_lab.universe import CURATED_TICKERS


def _plot_density(density_df: pd.DataFrame):
    fig, ax = plt.subplots(figsize=(7, 4))
    ax.plot(density_df["strike"], density_df["density"], linewidth=2)
    ax.set_title("Implied Density")
    ax.set_xlabel("Strike")
    ax.set_ylabel("Density proxy")
    ax.grid(alpha=0.25)
    fig.tight_layout()
    return fig


def _prepare_ticker(ticker_pick: str, custom_ticker: str) -> str:
    ticker = (
        custom_ticker.strip().upper()
        if custom_ticker.strip()
        else ticker_pick.strip().upper()
    )
    if not ticker:
        raise gr.Error("Ticker is required")
    return ticker


def _prepare_expiry(expiry_mode: str, manual_expiry: str):
    mode = "auto"
    expiry = None
    if expiry_mode == "Manual expiry":
        mode = "manual"
        if not manual_expiry.strip():
            raise gr.Error("Manual expiry selected; enter date like 2026-06-19")
        try:
            expiry = pd.to_datetime(manual_expiry).to_pydatetime()
        except Exception as exc:
            raise gr.Error(f"Could not parse expiry date: {exc}")
    return mode, expiry


def _enrich_strategy_table(scored_df: pd.DataFrame, risk_lambda: float) -> pd.DataFrame:
    if scored_df.empty:
        return scored_df
    out = scored_df.copy()
    out["downside_penalty"] = out["downside_q05"].apply(lambda x: max(-float(x), 0.0))
    out["objective_from_formula"] = (
        out["expected_payoff"] - float(risk_lambda) * out["downside_penalty"]
    )
    out = out.sort_values("objective", ascending=False).reset_index(drop=True)
    return out[
        [
            "strategy",
            "expected_payoff",
            "downside_q05",
            "downside_penalty",
            "objective",
            "objective_from_formula",
        ]
    ]


def _diagnostics_note(diag_df: pd.DataFrame) -> str:
    if diag_df.empty:
        return "No diagnostics output available."
    failed_checks = int((~diag_df["passed"]).sum())
    mean_rate = float(diag_df["violation_rate"].mean())
    if failed_checks == 0:
        return "Diagnostics are clean for this snapshot."
    return (
        f"Diagnostics are noisy ({failed_checks} failed checks, mean violation rate {mean_rate:.2%}). "
        "Treat these outputs as research signals, not direct execution instructions."
    )


def _strategy_note(scored_df: pd.DataFrame, risk_lambda: float) -> str:
    if scored_df.empty:
        return "No strategy scores available."
    top = scored_df.iloc[0]
    return (
        "Score formula: objective = E[payoff] - lambda * max(-q05, 0). "
        f"Current lambda = {float(risk_lambda):.2f}. "
        f"Top strategy = {top['strategy']} with objective {float(top['objective']):.4f}."
    )


def _worked_example_md(scored_df: pd.DataFrame, risk_lambda: float) -> str:
    if scored_df.empty:
        return "No worked example available."
    top = scored_df.iloc[0]
    return (
        "### Worked example from this run\n"
        f"- strategy: `{top['strategy']}`\n"
        f"- expected_payoff: `{float(top['expected_payoff']):.4f}`\n"
        f"- downside_q05: `{float(top['downside_q05']):.4f}`\n"
        f"- downside_penalty: `max(-q05, 0) = {float(top['downside_penalty']):.4f}`\n"
        f"- lambda: `{float(risk_lambda):.2f}`\n"
        f"- objective: `{float(top['expected_payoff']):.4f} - {float(risk_lambda):.2f} * {float(top['downside_penalty']):.4f} = {float(top['objective_from_formula']):.4f}`"
    )


def _arbitrage_table_for_display(cands_df: pd.DataFrame) -> pd.DataFrame:
    if cands_df.empty:
        return cands_df
    out = cands_df.copy()
    for col in ["k2", "k3"]:
        if col in out.columns:
            out[col] = out.apply(
                lambda r: "n/a"
                if r.get("family") in {"parity", "calendar"} and pd.isna(r[col])
                else r[col],
                axis=1,
            )
    return out


def run_walkthrough(
    ticker_pick: str,
    custom_ticker: str,
    expiry_mode: str,
    manual_expiry: str,
    moneyness_band: float,
    min_open_interest: int,
    risk_lambda: float,
):
    ticker = _prepare_ticker(ticker_pick, custom_ticker)
    mode, expiry = _prepare_expiry(expiry_mode, manual_expiry)

    try:
        with warnings.catch_warnings():
            warnings.filterwarnings("ignore", category=UserWarning, module="oipd")
            out = analyze_ticker_v2(
                ticker=ticker,
                max_expiries=8,
                expiry=expiry,
                expiry_mode=mode,
                moneyness_band=moneyness_band,
                min_open_interest=min_open_interest,
                min_volume=0,
                risk_lambda=risk_lambda,
            )
    except Exception as exc:
        raise gr.Error(f"Analysis failed: {exc}")

    scored = _enrich_strategy_table(out["scored"], risk_lambda=float(risk_lambda))
    diagnostics_note = _diagnostics_note(out["diagnostics"])
    strategy_note = _strategy_note(scored, risk_lambda=float(risk_lambda))
    worked_example = _worked_example_md(scored, risk_lambda=float(risk_lambda))
    density_fig = _plot_density(out["density"])

    run_summary = (
        f"Ticker: {out['ticker']} | Spot: {out['spot']:.2f} | "
        f"Expiry: {out['selected_expiry'].date()}"
    )

    best_json = json.dumps(out["best_strategy"], indent=2, default=str)

    return (
        run_summary,
        diagnostics_note,
        out["diagnostics"],
        density_fig,
        scored,
        strategy_note,
        worked_example,
        out["arbitrage_summary"],
        _arbitrage_table_for_display(out["arbitrage_candidates"]),
        best_json,
    )


def run_sweep(
    sweep_size: int,
    moneyness_band: float,
    min_open_interest: int,
    risk_lambda: float,
    focus_override: str,
):
    tickers = CURATED_TICKERS[: int(sweep_size)]
    rows: list[dict[str, object]] = []
    errors: list[dict[str, str]] = []
    artifacts: dict[str, dict[str, object]] = {}

    with warnings.catch_warnings():
        warnings.filterwarnings("ignore", category=UserWarning, module="oipd")
        for ticker in tickers:
            try:
                out = analyze_ticker_v2(
                    ticker=ticker,
                    max_expiries=8,
                    moneyness_band=moneyness_band,
                    min_open_interest=min_open_interest,
                    min_volume=0,
                    risk_lambda=risk_lambda,
                )
                artifacts[ticker] = out
                diag = out["diagnostics"]
                arb = out["arbitrage_summary"].iloc[0]
                high = int(arb.get("high_conf_count", 0))
                medium = int(arb.get("medium_conf_count", 0))
                low = int(arb.get("low_conf_count", 0))
                score = 3.0 * high + 1.0 * medium + 0.2 * low
                rows.append(
                    {
                        "ticker": ticker,
                        "status": "ok",
                        "selected_expiry": out["selected_expiry"],
                        "diag_mean_violation_rate": float(
                            diag["violation_rate"].mean()
                        ),
                        "best_strategy": out["best_strategy"]["name"],
                        "best_objective": float(out["best_strategy"]["objective"]),
                        "arb_candidates": int(arb["candidate_count"]),
                        "arb_max_edge": float(arb["max_edge"]),
                        "high_conf": high,
                        "medium_conf": medium,
                        "low_conf": low,
                        "score": score,
                    }
                )
            except Exception as exc:
                msg = str(exc)
                errors.append({"ticker": ticker, "error": msg})
                rows.append({"ticker": ticker, "status": f"error: {msg}"})

    sweep_df = pd.DataFrame(rows)
    ok = sweep_df[sweep_df["status"] == "ok"].copy()
    ranked = (
        ok.sort_values(
            ["score", "arb_candidates", "arb_max_edge"], ascending=[False, False, False]
        ).reset_index(drop=True)
        if not ok.empty
        else pd.DataFrame()
    )
    errors_df = pd.DataFrame(errors)

    if ranked.empty:
        fig, ax = plt.subplots(figsize=(7, 4))
        ax.set_title("No successful sweep results")
        ax.axis("off")
        return (
            f"Sweep complete: 0 success, {len(errors)} errors.",
            pd.DataFrame(),
            errors_df,
            "",
            "{}",
            pd.DataFrame(),
            fig,
        )

    focus_ticker = str(ranked.iloc[0]["ticker"])
    override = (focus_override or "").strip().upper()
    summary = (
        f"Sweep complete: {len(ok)} success, {len(errors)} errors. "
        f"Top ticker by score: {focus_ticker}."
    )
    if override:
        if override in artifacts:
            focus_ticker = override
            summary = summary + f" Focus override applied: {focus_ticker}."
        else:
            summary = (
                summary + f" Override '{override}' not available; kept {focus_ticker}."
            )

    focus_out = artifacts[focus_ticker]
    focus_best = json.dumps(focus_out["best_strategy"], indent=2, default=str)
    focus_candidates = _arbitrage_table_for_display(
        pd.DataFrame(focus_out["arbitrage_candidates"]).head(25)
    )
    focus_density = _plot_density(pd.DataFrame(focus_out["density"]))

    return (
        summary,
        ranked.head(25),
        errors_df,
        focus_ticker,
        focus_best,
        focus_candidates,
        focus_density,
    )


with gr.Blocks(title="Option-Implied Strategy Lab") as demo:
    gr.Markdown(
        """
        # Option-Implied Strategy Lab

        We start from market option prices and work backward to infer what shape of future outcomes the market is implying.
        Then we use that shape to compare strategies and flag possible pricing inconsistencies.

        Live demo: https://huggingface.co/spaces/junaid-hasan/implied-lab
        """
    )

    with gr.Tab("Demo Walkthrough"):
        gr.Markdown(
            r"""
            ## Big-picture introduction

            This app follows the same flow as `notebooks/demo.ipynb`:
            1) run diagnostics, 2) view implied density, 3) compare strategy scores, 4) inspect arbitrage candidates.

            Breeden-Litzenberger idea (under European assumptions):
            $$f_{RN}(K) = e^{rT} \frac{\partial^2 C(K,T)}{\partial K^2}$$

            Plain meaning: the way call prices bend across strikes gives a market-implied probability shape.

            ### Strategy meanings
            - `long_stock`: buy and hold the underlying with no option hedge.
            - `protective_put`: hold stock and buy a put to limit large downside losses.
            - `collar`: hold stock, buy a put, and sell a call to reduce hedge cost while capping upside.

            ### Arbitrage candidate families
            - `vertical`: checks strike ordering consistency.
            - `butterfly`: checks three-strike curvature consistency.
            - `parity`: checks call-put parity at one strike and expiry.
            - `calendar`: checks maturity ordering at the same strike.
            """
        )

        with gr.Row():
            ticker_pick = gr.Dropdown(
                choices=CURATED_TICKERS,
                value="GOOGL",
                label="Ticker (curated)",
                filterable=True,
            )
            custom_ticker = gr.Textbox(
                label="Custom ticker (optional)",
                value="",
                placeholder="Leave blank to use curated ticker",
            )
            expiry_mode = gr.Dropdown(
                choices=["Auto (best quality)", "Manual expiry"],
                value="Auto (best quality)",
                label="Expiry mode",
            )
            manual_expiry = gr.Textbox(
                label="Manual expiry (YYYY-MM-DD)",
                value="",
                placeholder="2026-06-19",
            )

        with gr.Row():
            moneyness_band = gr.Slider(
                0.1,
                0.4,
                value=0.2,
                step=0.05,
                label="Moneyness band (+/-)",
            )
            min_open_interest = gr.Slider(
                0,
                2000,
                value=1,
                step=1,
                label="Min open interest",
            )
            risk_lambda = gr.Slider(
                0.0,
                2.0,
                value=0.5,
                step=0.1,
                label="Risk lambda",
            )

        gr.Examples(
            label="Quick examples",
            examples=[
                ["GOOGL", "", "Auto (best quality)", "", 0.2, 1, 0.5],
                ["IWM", "", "Auto (best quality)", "", 0.2, 1, 0.5],
                ["NVDA", "", "Auto (best quality)", "", 0.2, 1, 0.5],
            ],
            inputs=[
                ticker_pick,
                custom_ticker,
                expiry_mode,
                manual_expiry,
                moneyness_band,
                min_open_interest,
                risk_lambda,
            ],
        )

        run_btn = gr.Button("Run walkthrough")

        run_summary = gr.Textbox(label="Run summary")
        diagnostics_note = gr.Textbox(label="Diagnostics note")
        diagnostics_table = gr.Dataframe(label="Diagnostics")
        density_plot = gr.Plot(label="Implied density")
        strategy_table = gr.Dataframe(label="Strategy ranking")
        strategy_note = gr.Textbox(label="Strategy note")
        worked_example = gr.Markdown(label="Objective worked example")
        arbitrage_summary = gr.Dataframe(label="Arbitrage summary")
        arbitrage_candidates = gr.Dataframe(label="Arbitrage candidates")
        best_strategy_json = gr.Code(label="Best strategy JSON", language="json")

        run_btn.click(
            fn=run_walkthrough,
            inputs=[
                ticker_pick,
                custom_ticker,
                expiry_mode,
                manual_expiry,
                moneyness_band,
                min_open_interest,
                risk_lambda,
            ],
            outputs=[
                run_summary,
                diagnostics_note,
                diagnostics_table,
                density_plot,
                strategy_table,
                strategy_note,
                worked_example,
                arbitrage_summary,
                arbitrage_candidates,
                best_strategy_json,
            ],
        )

    with gr.Tab("Universe Sweep"):
        gr.Markdown(
            r"""
            Sweep the curated universe and rank tickers with confidence-weighted score:
            $$\text{score} = 3\cdot\text{high} + 1\cdot\text{medium} + 0.2\cdot\text{low}$$
            """
        )

        with gr.Row():
            sweep_size = gr.Slider(
                5,
                min(len(CURATED_TICKERS), 60),
                value=12,
                step=1,
                label="Sweep size",
            )
            sweep_band = gr.Slider(
                0.1,
                0.4,
                value=0.2,
                step=0.05,
                label="Moneyness band (+/-)",
            )
            sweep_min_oi = gr.Slider(
                0,
                2000,
                value=1,
                step=1,
                label="Min open interest",
            )
            sweep_lambda = gr.Slider(
                0.0,
                2.0,
                value=0.5,
                step=0.1,
                label="Risk lambda",
            )
            focus_override = gr.Dropdown(
                choices=CURATED_TICKERS,
                value="",
                label="Focus ticker override (optional)",
                filterable=True,
                allow_custom_value=True,
            )

        gr.Examples(
            label="Sweep examples",
            examples=[
                [12, 0.2, 1, 0.5, ""],
                [20, 0.2, 1, 0.5, "GOOGL"],
                [15, 0.25, 1, 0.7, ""],
            ],
            inputs=[
                sweep_size,
                sweep_band,
                sweep_min_oi,
                sweep_lambda,
                focus_override,
            ],
        )

        sweep_btn = gr.Button("Run sweep")
        sweep_summary = gr.Textbox(label="Sweep summary")
        ranked_table = gr.Dataframe(label="Ranked tickers")
        sweep_errors = gr.Dataframe(label="Sweep errors")
        focus_ticker = gr.Textbox(label="Focus ticker")
        focus_best = gr.Code(label="Focus best strategy JSON", language="json")
        focus_candidates = gr.Dataframe(label="Focus arbitrage candidates")
        focus_density = gr.Plot(label="Focus density")

        sweep_btn.click(
            fn=run_sweep,
            inputs=[
                sweep_size,
                sweep_band,
                sweep_min_oi,
                sweep_lambda,
                focus_override,
            ],
            outputs=[
                sweep_summary,
                ranked_table,
                sweep_errors,
                focus_ticker,
                focus_best,
                focus_candidates,
                focus_density,
            ],
        )


if __name__ == "__main__":
    demo.launch()