from typing import Any, Dict, Optional

import pandas as pd
import statsmodels.api as sm

from pt_trading.rolling_window_fit import RollingFit
from pt_trading.trading_pair import TradingPair

NanoPerMin = 1e9


class ZScoreTradingPair(TradingPair):
    """TradingPair implementation that fits a hedge ratio with OLS and
    computes a standardized spread (z-score).

    The class stores training spread mean/std and hedge ratio so the model
    can be applied to testing data consistently.
    """

    zscore_model_: Optional[sm.regression.linear_model.RegressionResultsWrapper]
    pair_predict_result_: Optional[pd.DataFrame]
    zscore_df_: Optional[pd.Series]
    hedge_ratio_: Optional[float]
    spread_mean_: Optional[float]
    spread_std_: Optional[float]

    def __init__(
        self,
        config: Dict[str, Any],
        market_data: pd.DataFrame,
        symbol_a: str,
        symbol_b: str,
    ):
        super().__init__(config, market_data, symbol_a, symbol_b)
        self.zscore_model_ = None
        self.pair_predict_result_ = None
        self.zscore_df_ = None
        self.hedge_ratio_ = None
        self.spread_mean_ = None
        self.spread_std_ = None

    def _fit_zscore(self) -> None:
        """Fit OLS on the training window and compute training z-score."""
        assert self.training_df_ is not None

        # Extract price series for the two symbols from the training frame.
        px_df = self.training_df_[self.colnames()]
        symbol_a_px = px_df.iloc[:, 0]
        symbol_b_px = px_df.iloc[:, 1]

        # Align indexes and fit OLS: symbol_a ~ const + symbol_b
        symbol_a_px, symbol_b_px = symbol_a_px.align(symbol_b_px, join="inner")
        X = sm.add_constant(symbol_b_px)
        self.zscore_model_ = sm.OLS(symbol_a_px, X).fit()

        # Hedge ratio is the slope on symbol_b
        params = self.zscore_model_.params
        self.hedge_ratio_ = float(params.iloc[1]) if len(params) > 1 else 0.0

        # Training spread and its standardized z-score
        spread = symbol_a_px - self.hedge_ratio_ * symbol_b_px
        self.spread_mean_ = float(spread.mean())
        self.spread_std_ = float(spread.std(ddof=0)) if spread.std(ddof=0) != 0 else 1.0
        self.zscore_df_ = (spread - self.spread_mean_) / self.spread_std_

    def predict(self) -> pd.DataFrame:
        """Apply fitted hedge ratio to the testing frame and return a
        dataframe with canonical columns:
          - disequilibrium: signed z-score
          - scaled_disequilibrium: absolute z-score
          - signed_scaled_disequilibrium: same as disequilibrium (keeps sign)
        """
        # Fit on training window
        self._fit_zscore()
        assert self.zscore_df_ is not None
        assert self.hedge_ratio_ is not None
        assert self.spread_mean_ is not None and self.spread_std_ is not None

        # Keep training columns for inspection
        self.training_df_["disequilibrium"] = self.zscore_df_
        self.training_df_["scaled_disequilibrium"] = self.zscore_df_.abs()

        # Apply model to testing frame
        assert self.testing_df_ is not None
        test_df = self.testing_df_.copy()
        px_test = test_df[self.colnames()]
        a_test = px_test.iloc[:, 0]
        b_test = px_test.iloc[:, 1]
        a_test, b_test = a_test.align(b_test, join="inner")

        # Compute test spread and standardize using training mean/std
        test_spread = a_test - self.hedge_ratio_ * b_test
        test_zscore = (test_spread - self.spread_mean_) / self.spread_std_

        # Attach canonical columns
        # Align back to test_df index if needed
        test_zscore = test_zscore.reindex(test_df.index)
        test_df["disequilibrium"] = test_zscore
        test_df["signed_scaled_disequilibrium"] = test_zscore
        test_df["scaled_disequilibrium"] = test_zscore.abs()

        # Reset index and accumulate results across windows
        test_df = test_df.reset_index(drop=True)
        if self.pair_predict_result_ is None:
            self.pair_predict_result_ = test_df
        else:
            self.pair_predict_result_ = pd.concat(
                [self.pair_predict_result_, test_df], ignore_index=True
            )

        self.pair_predict_result_ = self.pair_predict_result_.reset_index(drop=True)
        return self.pair_predict_result_.dropna()


class ZScoreRollingFit(RollingFit):
    def __init__(self) -> None:
        super().__init__()

    def create_trading_pair(
        self, config: Dict, market_data: pd.DataFrame, symbol_a: str, symbol_b: str
    ) -> TradingPair:
        return ZScoreTradingPair(
            config=config, market_data=market_data, symbol_a=symbol_a, symbol_b=symbol_b
        )
