# region imports
from AlgorithmImports import *
from stat_arb import PairStatArb
from pair import Pair
# endregion


class SpectralStatArb(QCAlgorithm):

    def initialize(self):
        self.set_start_date(self.end_date - timedelta(5 * 365))
        self.set_cash(100_000)
        self.settings.automatic_indicator_warm_up = True
        self.settings.seed_initial_prices = True
        self.universe_settings.resolution = Resolution.DAILY
        self.universe_settings.schedule.on(self.date_rules.month_start("SPY"))
        self._universe = self.add_universe(lambda fundamental: 
            [f.symbol for f in sorted([f for f in fundamental if f.has_fundamental_data and f.market_cap > 5e9],
            key=lambda f: f.dollar_volume)[-100:]]
        )
        self.laplacian_lookback = 252
        self.arb_lookback = 30
        self.entry_z_score = 2.0
        self.exit_z_score = 0.5
        self.max_pairs = 5
        self._pairs = []
        self.set_warm_up(timedelta(45))

    def on_warmup_finished(self):
        # Configure training and prediction schedule.
        trading_time_rule = self.time_rules.at(8, 0)
        training_time_rule = self.time_rules.at(6, 0)
        self.train(self.date_rules.month_end(), training_time_rule, self._train)
        self.schedule.on(self.date_rules.every_day('SPY'), trading_time_rule, self._rebalance)
        # Rebalance today too.
        if self.live_mode:
            self._train()
            self._rebalance()
        else:
            self.train(self.date_rules.today, training_time_rule, self._train)
            self.schedule.on(self.date_rules.today, trading_time_rule, self._rebalance)

    def on_securities_changed(self, changes):
        for security in changes.removed_securities:
            # Drop any active pair that uses a leg leaving the universe.
            for pair in [p for p in self._pairs if security.symbol in p.symbols]:
                self._drop_pair(pair)

    def _train(self):
        # Train pair selection and warm each selected spread indicator.
        model = PairStatArb(self.laplacian_lookback, self.arb_lookback, self.max_pairs)
        spreads = model.train(self, self._universe)
        for pair in list(self._pairs):
            self._drop_pair(pair)
        if not spreads:
            self.log("No cointegrated pairs found in the cluster this month.")
            return
        # Split capital evenly across the pairs selected this month.
        allocation = 1 / len(spreads)
        for spread in spreads:
            pair = Pair(self, spread, allocation)
            self._pairs.append(pair)
            self.register_indicator(spread.symbol_a, spread)
            self.register_indicator(spread.symbol_b, spread)
            self.log(
                f"New Cointegrated Pair: {spread.symbol_a.value} & {spread.symbol_b.value} "
                f"(beta={spread.hedge_ratio:.2f}, alloc={allocation:.2%})"
            )

    def _rebalance(self):
        for pair in self._pairs:
            spread = pair.spread
            if not spread.is_ready:
                continue
            half_life = float(spread.half_life)
            if not np.isfinite(half_life) or half_life > 30:
                continue
            dynamic_entry_z = self.entry_z_score if half_life < 10 else self.entry_z_score + 0.5
            z_score = float(spread.z_score)
            if not np.isfinite(z_score):
                continue
            if not pair.is_invested and abs(z_score) > dynamic_entry_z:
                # Short the spread when z-score is high, and long it when z-score is low.
                pair.enter(np.sign(z_score))
                self.log(
                    f"ENTRY: {pair.symbol_a.value}/{pair.symbol_b.value} | HL: {half_life:.1f}d | Z: {z_score:.2f}"
                )
            elif pair.is_invested and abs(z_score) < self.exit_z_score:
                # Take profit once the spread reverts toward the mean.
                pair.exit()
                self.log("EXIT: Pair converged")

    def _drop_pair(self, pair):
        # Close the pair's positions and release its spread indicator.
        pair.exit()
        self.deregister_indicator(pair.spread)
        self._pairs.remove(pair)

# region imports
from AlgorithmImports import *
from pair_spread import PairSpread
# endregion


class Pair:

    def __init__(self, algorithm: QCAlgorithm, spread: PairSpread, allocation: float):
        self._algorithm = algorithm
        self.spread = spread
        self.symbols = spread.symbols
        self.symbol_a, self.symbol_b = self.symbols
        self._a = algorithm.securities[self.symbol_a]
        self._b = algorithm.securities[self.symbol_b]
        self._allocation = allocation
        self._ticket_by_security = {}

    @property
    def is_invested(self) -> bool:
        return bool(self._ticket_by_security)

    def enter(self, direction: float):
        if self._ticket_by_security:
            return
        gross = 1 + abs(self.spread.hedge_ratio)
        if not np.isfinite(gross) or gross <= 0:
            return
        order_size = self._allocation * self._algorithm.portfolio.total_portfolio_value
        dollar_a = -direction * order_size / gross
        dollar_b = direction * self.spread.hedge_ratio * order_size / gross
        quantity_a = int(dollar_a / self._a.price)
        quantity_b = int(dollar_b / self._b.price)
        if quantity_a and quantity_b:
            self._ticket_by_security = {
                self._a: self._algorithm.market_order(self._a, quantity_a),
                self._b: self._algorithm.market_order(self._b, quantity_b),
            }

    def exit(self):
        for security, ticket in self._ticket_by_security.items():
            self._algorithm.market_order(security, -ticket.quantity)
        self._ticket_by_security.clear()

# region imports
from AlgorithmImports import *
# endregion


class PairSpread(PythonIndicator):

    def __init__(self, symbol_a, symbol_b, hedge_ratio: float, period: int, name: str = "PairSpread"):
        super().__init__()
        self.name = name
        self.symbol_a = symbol_a
        self.symbol_b = symbol_b
        self.symbols = (symbol_a, symbol_b)
        self.hedge_ratio = float(hedge_ratio)
        self.value = 0.0
        self.spread = 0.0
        self.z_score = 0.0
        self.half_life = np.inf
        # Store spread values.
        self._window = RollingWindow[float](period)
        # Cache the latest price and timestamp for each leg.
        self._prices = {}
        self._last_time = None

    def update(self, input_: BaseData):
        price = float(input_.value)
        if input_.symbol not in self.symbols or price <= 0:
            return self.is_ready
        self._prices[input_.symbol] = (input_.end_time, price)
        if self.symbol_a not in self._prices or self.symbol_b not in self._prices:
            return self.is_ready
        time_a, price_a = self._prices[self.symbol_a]
        time_b, price_b = self._prices[self.symbol_b]
        if time_a != time_b or time_a == self._last_time:
            return self.is_ready
        self._last_time = time_a
        # Calculate the log-price spread using the selected hedge ratio.
        self.value = float(np.log(price_a) - self.hedge_ratio * np.log(price_b))
        self.spread = self.value
        # Reset signal values until the spread window can recompute them.
        self.z_score = 0.0
        self.half_life = np.inf
        self._window.add(self.value)
        # Recompute z-score and half-life from the complete trailing spread window.
        if self._window.is_ready:
            spreads = np.array(list(self._window)[::-1], dtype=float)
            sd = np.std(spreads)
            if sd > 0 and np.isfinite(sd):
                self.z_score = float((self.value - np.mean(spreads)) / sd)
            # Estimate mean-reversion speed with a lagged spread regression.
            lagged = spreads[:-1]
            if len(lagged) > 1 and not np.all(lagged == lagged[0]):
                slope = np.polyfit(lagged, spreads[1:] - lagged, 1)[0]
                if np.isfinite(slope) and slope < 0:
                    self.half_life = float(-np.log(2) / slope)
        return self.is_ready

    @property
    def is_ready(self) -> bool:
        return self._window.is_ready

# region imports
from AlgorithmImports import *
from pair_spread import PairSpread
from statsmodels.tsa.stattools import adfuller
# endregion


class PairStatArb:

    # Initialize the strategy parameters for graph clustering and statistical arbitrage lookbacks.
    def __init__(self, laplacian_lookback: int, arb_lookback: int, max_pairs: int):
        self._laplacian_lookback = laplacian_lookback
        self._arb_lookback = arb_lookback
        self._max_pairs = max_pairs

    def train(self, algorithm, universe):
        symbols = sorted(universe.selected, key=lambda s: s.value)
        # Ensure there are at least two valid assets available to form pairs.
        if len(symbols) < 2:
            return []
        # Fetch historical close prices for the whole universe in a single request.
        history = algorithm.history(symbols, self._laplacian_lookback + 1, Resolution.DAILY)
        if history.empty:
            return []
        log_prices = np.log(history["close"].unstack(level=0))
        # diff()'s first row is NaN for every symbol; drop it before dropping
        # columns, or dropna(axis=1) would remove every symbol.
        df_returns = log_prices.diff().iloc[1:].dropna(axis=1)
        if df_returns.shape[1] < 2:
            return []
        valid_symbols = list(df_returns.columns)
        log_prices = log_prices[valid_symbols]
        # Build the weighted correlation graph and its Laplacian.
        weights = np.abs(df_returns.corr().fillna(0).values)
        np.fill_diagonal(weights, 0)
        laplacian = np.diag(np.sum(weights, axis=1)) - weights
        # Compute the eigenvalues and eigenvectors to find the Fiedler vector for spectral clustering.
        _, eigenvectors = np.linalg.eigh(laplacian)
        fiedler_vector = eigenvectors[:, 1]
        # Standardize the Fiedler vector orientation based on its maximum absolute value.
        if fiedler_vector[np.argmax(np.abs(fiedler_vector))] < 0:
            fiedler_vector = -fiedler_vector
        # Select assets belonging to the positive partition of the Fiedler vector.
        cluster_pos = [valid_symbols[i] for i in range(len(valid_symbols)) if fiedler_vector[i] > 0]
        if len(cluster_pos) < 2:
            return []
        cluster_log_prices = log_prices[cluster_pos]
        # Iterate over all asset combinations within the cluster to locate cointegrated candidates.
        candidates = []
        for i in range(len(cluster_pos)):
            for j in range(i + 1, len(cluster_pos)):
                candidate = self._try_create_pair(algorithm, cluster_log_prices, cluster_pos[i], cluster_pos[j])
                if candidate:
                    candidates.append(candidate)
        # Keep the pairs with the strongest cointegration evidence, ranked by ADF p-value.
        candidates.sort(key=lambda candidate: candidate[0])
        return [pair_spread for _, pair_spread in candidates[:self._max_pairs]]

    def _try_create_pair(self, algorithm, log_prices, asset_a, asset_b):
        aligned = log_prices[[asset_a, asset_b]].dropna()
        if len(aligned) < self._arb_lookback:
            return None
        y = aligned[asset_a]
        x = aligned[asset_b]
        if np.all(x == x.iloc[0]):
            return None
        # Determine the hedge ratio using an ordinary least squares linear regression.
        hedge_ratio = float(np.polyfit(x, y, 1)[0])
        # Construct the spread series using the computed hedge ratio.
        spread = y - hedge_ratio * x
        if np.all(spread == spread.iloc[0]):
            return None
        # Verify the stationarity of the spread using the Augmented Dickey-Fuller test.
        p_value = adfuller(spread)[1]
        if p_value >= 0.05:
            return None
        pair_spread = PairSpread(asset_a, asset_b, hedge_ratio, self._arb_lookback)
        # Fetch historical bar data to warm up the indicator windows of the pair.
        pair_history = algorithm.history[TradeBar]([asset_a, asset_b], self._arb_lookback + 5, Resolution.DAILY)
        for bars in pair_history:
            if asset_a in bars and asset_b in bars:
                pair_spread.update(bars[asset_a])
                pair_spread.update(bars[asset_b])
        return p_value, pair_spread