from typing import List
from pandas import DataFrame
import numpy as np
from quant.strategy import Strategy, HistoricalData, StrategyOutput


class MeanStrategy(Strategy):
    def __init__(
        self,
        stocks: List[str],
        lookback: int = 50,          # ile punktów wstecz do liczenia z-score
        min_lookback: int = 5,       # minimalna liczba obserwacji, żeby w ogóle robić MR
        use_equal_fallback: bool = False,  # czy w razie problemów robić równy portfel
    ) -> None:
        self.stocks = stocks
        self.lookback = lookback
        self.min_lookback = min_lookback
        self.use_equal_fallback = use_equal_fallback

    def _equal_weights(self, available_stocks: List[str]) -> StrategyOutput:
        n = len(available_stocks)
        if n == 0:
            raise ValueError("No available stocks to assign equal weights to.")
        equal_w = np.ones(n) / n
        weights_df = DataFrame({
            "stock_name": available_stocks,
            "weight": equal_w
        })
        return StrategyOutput(weights=weights_df)

    def get_weights(self, historical_data: HistoricalData) -> StrategyOutput:
        prices = historical_data.stock_prices

        # Upewniamy się, że używamy tylko tych kolumn, które faktycznie są w DataFrame
        available_stocks = [s for s in self.stocks if s in prices.columns]
        if len(available_stocks) == 0:
            raise ValueError("None of the requested stocks are present in stock_prices.")

        prices = prices[available_stocks].dropna(how="all")

        if len(prices) == 0:
            if self.use_equal_fallback:
                return self._equal_weights(available_stocks)
            else:
                raise ValueError("No price data available for the selected stocks.")

        # --- Mean reversion na bazie z-score cen ---
        # realne okno = min(lookback, liczba dostępnych punktów)
        lookback = min(self.lookback, len(prices))

        if lookback < self.min_lookback:
            if self.use_equal_fallback:
                return self._equal_weights(available_stocks)
            else:
                raise ValueError(
                    f"Not enough data points for lookback={self.lookback}, "
                    f"got only {len(prices)} rows."
                )

        roll_mean = prices.rolling(window=lookback).mean()
        roll_std = prices.rolling(window=lookback).std()

        zscore = (prices - roll_mean) / roll_std
        # używamy ostatnie dostępne z-score (bez look-ahead)
        z_last = zscore.iloc[-1]

        # Mean reversion: jeśli coś jest wysoko (z>0), spodziewamy się spadku → short,
        # jeśli nisko (z<0), spodziewamy się odbicia → long
        signal = -z_last.values  # minus = mean reversion

        # Jeżeli wszystko NaN lub 0 → równy portfel (albo błąd, zależnie od flagi)
        if np.all(np.isnan(signal)) or np.allclose(signal, 0):
            if self.use_equal_fallback:
                return self._equal_weights(available_stocks)
            else:
                raise ValueError("Signal is degenerate (all NaN or zero).")

        # Zastępujemy NaN zerami (brak sygnału na danym instrumencie)
        signal = np.nan_to_num(signal, nan=0.0)

        # Jeśli nadal same zera → fallback
        if np.allclose(signal, 0):
            if self.use_equal_fallback:
                return self._equal_weights(available_stocks)
            else:
                raise ValueError("Signal is zero after NaN handling.")

        # Normalizacja: sum(|w|) = 1
        abs_sum = np.sum(np.abs(signal))
        if abs_sum == 0:
            if self.use_equal_fallback:
                return self._equal_weights(available_stocks)
            else:
                raise ValueError("Absolute sum of signal is zero.")

        weights = signal / abs_sum

        weights_df = DataFrame({
            "stock_name": available_stocks,
            "weight": weights
        })

        return StrategyOutput(weights=weights_df)