"""Simple volatility calculator.

基于 pandas 的波动率计算器实现，支持多种 HV 估计方法。
"""

import numpy as np
import pandas as pd

from .base import HVMethod, VolatilityCalculator


class SimpleVolatilityCalculator(VolatilityCalculator):
    """Pure-pandas volatility calculator.

    纯 pandas 实现的波动率计算器，支持 C2C/Parkinson/Garman-Klass/Yang-Zhang。
    """

    def get_hv(
        self,
        ohlc_df: pd.DataFrame,
        method: HVMethod,
        window: int,
        annualization: int = 252,
    ) -> pd.DataFrame:
        """Calculate historical volatility.

        Args:
            ohlc_df: OHLC DataFrame.
            method: 'close_to_close', 'parkinson', 'garman_klass', 'yang_zhang'.
            window: Rolling window length.
            annualization: Trading days per year, default 252.

        Returns:
            pd.DataFrame: [timestamp, hv]
        """
        self._require_columns(ohlc_df, ["close"])
        df = ohlc_df.copy()

        if method == "yang_zhang" and window < 2:
            raise ValueError("Yang-Zhang requires window >= 2.")

        if method in ("parkinson", "garman_klass", "yang_zhang"):
            required = {"parkinson": ["high", "low"],
                        "garman_klass": ["open", "high", "low", "close"],
                        "yang_zhang": ["open", "high", "low", "close"]}
            self._require_columns(df, required[method])
            price_cols = required[method]
            if (df[price_cols] <= 0).any().any():
                raise ValueError("All price columns must be strictly positive.")

        if method == "close_to_close":
            df["_log_ret"] = np.log(df["close"] / df["close"].shift(1))
            hv = df["_log_ret"].rolling(window=window).std() * np.sqrt(annualization)
        elif method == "parkinson":
            self._require_columns(df, ["high", "low"])
            df["_hl_sq"] = np.log(df["high"] / df["low"]) ** 2
            hv = np.sqrt(
                df["_hl_sq"].rolling(window=window).mean() / (4 * np.log(2)) * annualization
            )
        elif method == "garman_klass":
            self._require_columns(df, ["open", "high", "low", "close"])
            df["_gk"] = (
                0.5 * np.log(df["high"] / df["low"]) ** 2
                - (2 * np.log(2) - 1) * np.log(df["close"] / df["open"]) ** 2
            )
            hv = np.sqrt(df["_gk"].rolling(window=window).mean() * annualization)
        elif method == "yang_zhang":
            self._require_columns(df, ["open", "high", "low", "close"])
            df["_log_oc"] = np.log(df["open"] / df["close"].shift(1))
            df["_log_co"] = np.log(df["close"] / df["open"])
            df["_log_ho"] = np.log(df["high"] / df["open"])
            df["_log_lo"] = np.log(df["low"] / df["open"])

            overnight_var = df["_log_oc"].rolling(window=window).var()
            close_var = df["_log_co"].rolling(window=window).var()
            df["_rs"] = (
                df["_log_ho"] * (df["_log_ho"] - df["_log_co"])
                + df["_log_lo"] * (df["_log_lo"] - df["_log_co"])
            )
            rs_var = df["_rs"].rolling(window=window).mean()

            k = 0.34 / (1.34 + (window + 1) / (window - 1))
            yz_var = overnight_var + k * close_var + (1 - k) * rs_var
            hv = np.sqrt(yz_var * annualization)
        else:
            raise ValueError(f"Unknown HV method: {method}")

        df["hv"] = hv
        temp_cols = [c for c in df.columns if c.startswith("_")]
        df = df.drop(columns=temp_cols)
        cols = [c for c in ["timestamp", "hv"] if c in df.columns]
        return df[cols].reset_index(drop=True)

    def get_volatility_cone(
        self,
        ohlc_df: pd.DataFrame,
        windows: list[int],
        percentiles: list[float],
        method: HVMethod = "close_to_close",
        annualization: int = 252,
    ) -> pd.DataFrame:
        """Calculate volatility cone.

        Args:
            ohlc_df: OHLC DataFrame.
            windows: Window lengths.
            percentiles: Percentiles (0-1).
            method: HV method, default 'close_to_close'.
            annualization: Trading days per year, default 252.

        Returns:
            pd.DataFrame: [window, percentile, hv]
        """
        if not windows or not percentiles:
            raise ValueError("windows and percentiles must be non-empty.")
        if any(w <= 0 for w in windows):
            raise ValueError("All windows must be positive.")
        if any(p < 0 or p > 1 for p in percentiles):
            raise ValueError("All percentiles must be in [0, 1].")

        results = []
        for w in windows:
            hv_df = self.get_hv(ohlc_df, method=method, window=w, annualization=annualization)
            hv_values = hv_df["hv"].dropna()
            for p in percentiles:
                results.append({
                    "window": w,
                    "percentile": p,
                    "hv": float(np.quantile(hv_values, p)) if len(hv_values) > 0 else None,
                })

        return pd.DataFrame(results)

    def get_hv_iv_spread(self, hv_df: pd.DataFrame, iv_df: pd.DataFrame) -> pd.DataFrame:
        """Calculate HV-IV spread.

        Args:
            hv_df: DataFrame with [timestamp, hv].
            iv_df: DataFrame with [timestamp, iv].

        Returns:
            pd.DataFrame: [timestamp, hv, iv, hv_iv_spread]
        """
        self._require_columns(hv_df, ["timestamp", "hv"])
        self._require_columns(iv_df, ["timestamp", "iv"])

        df = pd.merge(hv_df, iv_df, on="timestamp", how="inner")
        df["hv_iv_spread"] = df["hv"] - df["iv"]

        return df[["timestamp", "hv", "iv", "hv_iv_spread"]].reset_index(drop=True)

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

        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}")
