"""Simple options indicator calculator.

基于 pandas 的期权指标计算器实现。
"""

import numpy as np
import pandas as pd

from .base import OptionsIndicatorCalculator, PCRBy


class SimpleOptionsIndicatorCalculator(OptionsIndicatorCalculator):
    """Pure-pandas options indicator calculator.

    纯 pandas 实现的期权指标计算器。
    """

    def get_pcr(self, chain_df: pd.DataFrame, by: PCRBy) -> pd.DataFrame:
        """Calculate Put-Call Ratio.

        Args:
            chain_df: Options chain with columns [option_type, volume/open_interest].
            by: 'volume' or 'open_interest'.

        Returns:
            pd.DataFrame: [pcr, put_value, call_value]
        """
        self._require_columns(chain_df, ["option_type", by])

        df = chain_df.copy()
        df[by] = pd.to_numeric(df[by], errors="coerce").fillna(0)
        put_val = df.loc[df["option_type"] == "put", by].sum()
        call_val = df.loc[df["option_type"] == "call", by].sum()
        pcr = put_val / call_val if call_val > 0 else float("inf")

        return pd.DataFrame([{"pcr": pcr, "put_value": put_val, "call_value": call_val}])

    def get_max_pain(self, chain_df: pd.DataFrame) -> pd.DataFrame:
        """Calculate Max Pain.

        Args:
            chain_df: Options chain with [option_type, strike, open_interest].

        Returns:
            pd.DataFrame: [max_pain_strike, total_pain]
        """
        self._require_columns(chain_df, ["option_type", "strike", "open_interest"])

        df = chain_df.copy()
        df["open_interest"] = pd.to_numeric(df["open_interest"], errors="coerce").fillna(0)

        strikes = sorted(df["strike"].dropna().unique())
        min_pain = float("inf")
        max_pain_strike = strikes[0] if strikes else 0

        calls = df[df["option_type"] == "call"][["strike", "open_interest"]]
        puts = df[df["option_type"] == "put"][["strike", "open_interest"]]

        for settlement in strikes:
            call_pain = calls.apply(
                lambda r, s=settlement: max(0, s - r["strike"]) * r["open_interest"],
                axis=1,
            ).sum()
            put_pain = puts.apply(
                lambda r, s=settlement: max(0, r["strike"] - s) * r["open_interest"],
                axis=1,
            ).sum()
            total = call_pain + put_pain
            if total < min_pain:
                min_pain = total
                max_pain_strike = settlement

        return pd.DataFrame([{"max_pain_strike": max_pain_strike, "total_pain": min_pain}])

    def get_iv_rank(self, iv_series: pd.DataFrame, window: int) -> pd.DataFrame:
        """Calculate IV Rank: (current - min) / (max - min).

        Args:
            iv_series: DataFrame with [timestamp, iv].
            window: Lookback window.

        Returns:
            pd.DataFrame: [timestamp, iv, iv_rank] where iv_rank is in [0, 1].
        """
        self._require_columns(iv_series, ["iv"])
        df = iv_series.copy()
        rolling_min = df["iv"].rolling(window=window, min_periods=1).min()
        rolling_max = df["iv"].rolling(window=window, min_periods=1).max()
        range_ = rolling_max - rolling_min
        df["iv_rank"] = np.where(range_ > 0, (df["iv"] - rolling_min) / range_, 0.0)
        cols = [c for c in ["timestamp", "iv", "iv_rank"] if c in df.columns]
        return df[cols].reset_index(drop=True)

    def get_iv_percentile(self, iv_series: pd.DataFrame, window: int) -> pd.DataFrame:
        """Calculate IV Percentile: fraction of prior values below current.

        Args:
            iv_series: DataFrame with [timestamp, iv].
            window: Lookback window.

        Returns:
            pd.DataFrame: [timestamp, iv, iv_percentile] where iv_percentile is in [0, 1].
        """
        self._require_columns(iv_series, ["iv"])
        df = iv_series.copy()
        df["iv_percentile"] = (
            df["iv"]
            .rolling(window=window, min_periods=1)
            .apply(lambda x: (x.iloc[:-1] < x.iloc[-1]).sum() / max(len(x) - 1, 1), raw=False)
        )
        cols = [c for c in ["timestamp", "iv", "iv_percentile"] if c in df.columns]
        return df[cols].reset_index(drop=True)

    def get_iv_smile(
        self,
        chain_df: pd.DataFrame,
        expiration: str,
        underlying_price: float | None = None,
    ) -> pd.DataFrame:
        """Get IV smile curve for a specific expiration using OTM contracts.

        Standard approach: use OTM call (K >= ATM) + OTM put (K <= ATM) to
        construct the smile, avoiding illiquid deep-ITM data.

        Args:
            chain_df: Options chain with [expiration, strike, option_type,
                implied_volatility/iv]. Optionally includes ``iv_source``.
            expiration: Target expiration date.
            underlying_price: Spot/underlying price for moneyness calculation
                and OTM filtering. Falls back to median strike if not provided.

        Returns:
            pd.DataFrame: [expiration, strike, iv, moneyness, iv_source]
        """
        iv_col = "implied_volatility" if "implied_volatility" in chain_df.columns else "iv"
        self._require_columns(chain_df, ["expiration", "strike", "option_type", iv_col])

        df = chain_df[chain_df["expiration"] == expiration].copy()
        if df.empty:
            return pd.DataFrame(columns=["expiration", "strike", "iv", "moneyness", "iv_source"])
        df = df.rename(columns={iv_col: "iv"}) if iv_col != "iv" else df

        atm = underlying_price if underlying_price is not None else df["strike"].median()

        otm_calls = df[(df["option_type"] == "call") & (df["strike"] >= atm)]
        otm_puts = df[(df["option_type"] == "put") & (df["strike"] <= atm)]
        df = pd.concat([otm_puts, otm_calls], ignore_index=True)
        df = df.drop_duplicates(subset=["strike"], keep="first")

        df = df[df["iv"] > 0].copy()
        df = df.sort_values("strike").reset_index(drop=True)
        df["moneyness"] = df["strike"] / atm if atm > 0 else df["strike"]

        if "iv_source" not in df.columns:
            df["iv_source"] = "sse"

        out_cols = ["expiration", "strike", "iv", "moneyness", "iv_source"]
        return df[[c for c in out_cols if c in df.columns]].reset_index(drop=True)

    def get_iv_term_structure(
        self, chain_df: pd.DataFrame, underlying_price: float | None = None
    ) -> pd.DataFrame:
        """Calculate IV term structure.

        Args:
            chain_df: Options chain with [expiration, strike, implied_volatility/iv].
            underlying_price: Underlying/spot price for ATM identification. Optional.

        Returns:
            pd.DataFrame: [expiration, iv_atm]
        """
        iv_col = "implied_volatility" if "implied_volatility" in chain_df.columns else "iv"
        self._require_columns(chain_df, ["expiration", "strike", iv_col])

        results = []
        for exp, group in chain_df.groupby("expiration"):
            atm_ref = underlying_price if underlying_price is not None else group["strike"].median()
            atm_row = group.iloc[(group["strike"] - atm_ref).abs().argsort()[:1]]
            iv_atm = atm_row[iv_col].iloc[0] if not atm_row.empty else None
            results.append({"expiration": exp, "iv_atm": iv_atm})

        return pd.DataFrame(results)

    @staticmethod
    def _require_columns(df: pd.DataFrame, columns: list[str]) -> None:
        """Validate required columns exist.

        Raises:
            ValueError: If any column is missing.
        """
        missing = [c for c in columns if c not in df.columns]
        if missing:
            raise ValueError(f"Missing required columns: {missing}")
