from __future__ import annotations

from datetime import datetime

import numpy as np
import pandas as pd

from .density import DensityEstimate
from .payoff import payoff_collar, payoff_protective_put, payoff_stock, risk_metrics
from .types import ScoredStrategy, StrategyCandidate


def _nearest_price(
    options: pd.DataFrame, expiry: datetime, option_type: str, strike: float
) -> float:
    o = options[
        (options["expiry"] == pd.to_datetime(expiry))
        & (options["option_type"] == option_type)
    ].copy()
    if o.empty:
        raise ValueError(f"No {option_type} options for expiry={expiry}")
    o["dist"] = (o["strike"] - strike).abs()
    row = o.sort_values("dist").iloc[0]
    return float(row["mid"])


def generate_candidates(
    spot: float, density: DensityEstimate, options: pd.DataFrame
) -> list[StrategyCandidate]:
    expiry = density.expiry
    k_put = 0.95 * spot
    k_call = 1.05 * spot
    put_premium = _nearest_price(options, expiry, "put", k_put)
    call_premium = _nearest_price(options, expiry, "call", k_call)
    return [
        StrategyCandidate(name="long_stock", params={"units": 1.0}),
        StrategyCandidate(
            name="protective_put",
            params={"units": 1.0, "strike_put": k_put, "put_premium": put_premium},
        ),
        StrategyCandidate(
            name="collar",
            params={
                "units": 1.0,
                "strike_put": k_put,
                "put_premium": put_premium,
                "strike_call": k_call,
                "call_premium": call_premium,
            },
        ),
    ]


def _strategy_payoff(
    name: str, params: dict[str, float], s_t: np.ndarray, spot0: float
) -> np.ndarray:
    if name == "long_stock":
        return payoff_stock(s_t, spot0, units=params.get("units", 1.0))
    if name == "protective_put":
        return payoff_protective_put(
            s_t,
            spot0,
            strike_put=params["strike_put"],
            put_premium=params["put_premium"],
            units=params.get("units", 1.0),
        )
    if name == "collar":
        return payoff_collar(
            s_t,
            spot0,
            strike_put=params["strike_put"],
            put_premium=params["put_premium"],
            strike_call=params["strike_call"],
            call_premium=params["call_premium"],
            units=params.get("units", 1.0),
        )
    raise ValueError(f"Unknown strategy {name}")


def score_candidates(
    candidates: list[StrategyCandidate],
    density: DensityEstimate,
    spot0: float,
    risk_lambda: float = 0.5,
) -> list[ScoredStrategy]:
    s_t = density.strikes.to_numpy(dtype=float)
    probs = density.density.to_numpy(dtype=float)
    out: list[ScoredStrategy] = []
    for c in candidates:
        pay = _strategy_payoff(c.name, c.params, s_t, spot0)
        metrics = risk_metrics(pay, probs=probs)
        objective = metrics["expected_payoff"] - risk_lambda * abs(
            min(metrics["q05"], 0.0)
        )
        out.append(
            ScoredStrategy(
                candidate=c,
                expected_payoff=metrics["expected_payoff"],
                downside_q05=metrics["q05"],
                objective=float(objective),
            )
        )
    return sorted(out, key=lambda x: x.objective, reverse=True)


def select_best(scored: list[ScoredStrategy]) -> ScoredStrategy:
    if not scored:
        raise ValueError("No scored candidates")
    return scored[0]