# 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(365))
        self.set_cash(100_000)
        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._pairs = []
        self.set_warm_up(timedelta(45))

    def on_warmup_finished(self):
        training_time_rule = self.time_rules.at(6, 0)
        trading_time_rule = self.time_rules.at(8, 0)
        self.train(self.date_rules.month_end("SPY"), training_time_rule, self._train)
        self.schedule.on(self.date_rules.every_day("SPY"), trading_time_rule, self._rebalance)
        # Train and rebalance today too (live mode vs backtest).
        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.spread.symbols]:
                self._drop_pair(pair)

    def _train(self):
        # Select cointegrated pairs within the spectral clusters for this month's trading.
        spreads = PairStatArb(252, 10).train(self, self._universe)
        # Retire flat pairs, but let invested pairs keep trading toward convergence like the paper's overlapping trading periods.
        for pair in [p for p in self._pairs if not p.is_invested]:
            self._drop_pair(pair)
        if not spreads:
            self.log("No cointegrated pairs found in the clusters this month.")
            return
        allocation = 1 / len(spreads)
        live_symbols = [pair.spread.symbols for pair in self._pairs]
        # Cap the live pair count so carried positions cannot push gross exposure past margin limits.
        for spread in [s for s in spreads if s.symbols not in live_symbols][:15 - len(self._pairs)]:
            self._pairs.append(Pair(self, spread, allocation))
            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} (beta={spread.hedge_ratio:.2f}, alloc={allocation:.2%})")

    def _rebalance(self):
        for pair in list(self._pairs):
            spread = pair.spread
            z_score = float(spread.z_score)
            if not np.isfinite(z_score):
                continue
            if not pair.is_invested and 2 < abs(z_score) <= 4:
                # Short the spread when the z-score is high, and long it when the z-score is low.
                pair.enter(np.sign(z_score))
                self.log(f"ENTRY: {spread.symbol_a.value}/{spread.symbol_b.value} | Z: {z_score:.2f}")
            elif pair.is_invested and abs(z_score) > 4:
                # Cut the loss and retire the pair once the spread breaks far beyond the entry band.
                self.log(f"STOP: {spread.symbol_a.value}/{spread.symbol_b.value} | Z: {z_score:.2f}")
                self._drop_pair(pair)
            elif pair.is_invested and abs(z_score) < 0.5:
                # 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._a = algorithm.securities[spread.symbol_a]
        self._b = algorithm.securities[spread.symbol_b]
        self._allocation = allocation
        self._ticket_by_security = {}

    def enter(self, direction: float):
        if self._ticket_by_security or not (self._a.price and self._b.price):
            return
        # Split the pair's gross allocation between the two legs in the hedge ratio proportion.
        gross = 1 + abs(self.spread.hedge_ratio)
        order_size = self._allocation * self._algorithm.portfolio.total_portfolio_value
        quantity_a = int(-direction * order_size / gross / self._a.price)
        quantity_b = int(direction * self.spread.hedge_ratio * order_size / gross / 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):
        # Unwind only what actually filled so a rejected entry cannot create a naked position.
        for security, ticket in self._ticket_by_security.items():
            if ticket.quantity_filled:
                self._algorithm.market_order(security, -ticket.quantity_filled)
        self._ticket_by_security.clear()

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

# region imports
from AlgorithmImports import *
# endregion


class PairSpread(PythonIndicator):

    def __init__(self, symbol_a, symbol_b, hedge_ratio: float, mean: float, sd: float):
        super().__init__()
        self.name = f"PairSpread({symbol_a.value},{symbol_b.value})"
        self.symbol_a = symbol_a
        self.symbol_b = symbol_b
        self.symbols = (symbol_a, symbol_b)
        self.hedge_ratio = float(hedge_ratio)
        self._mean = float(mean)
        self._sd = float(sd)
        self.value = 0.0
        self.z_score = np.nan
        # Cache the latest price and timestamp for each leg.
        self._prices = {}

    def update(self, input_: BaseData):
        price = float(input_.value)
        if input_.symbol not in self.symbols or not np.isfinite(price) 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:
            return self.is_ready
        # Score the current log-price spread against the formation-period distribution.
        self.value = float(np.log(price_a) - self.hedge_ratio * np.log(price_b))
        self.z_score = float((self.value - self._mean) / self._sd)
        return self.is_ready

    @property
    def is_ready(self) -> bool:
        return bool(np.isfinite(self.z_score))

# region imports
from AlgorithmImports import *
from pair_spread import PairSpread
from sklearn.cluster import KMeans
from statsmodels.tsa.stattools import coint
# endregion


class PairStatArb:

    def __init__(self, laplacian_lookback: int, max_pairs: int):
        self._laplacian_lookback = laplacian_lookback
        self._max_pairs = max_pairs

    def train(self, algorithm, universe):
        symbols = sorted(universe.selected, key=lambda s: s.value)
        if len(symbols) < 2:
            return []
        history = algorithm.history(symbols, self._laplacian_lookback + 1, Resolution.DAILY)
        if history.empty:
            return []
        log_prices = np.log(history["close"].unstack(level=0))
        df_returns = log_prices.diff().iloc[1:].dropna(axis=1)
        if df_returns.shape[1] < 2:
            return []
        valid_symbols = list(df_returns.columns)
        # Build the similarity graph from positive daily log-return correlations, since economic pairs must co-move.
        weights = np.maximum(df_returns.corr().fillna(0).values, 0)
        np.fill_diagonal(weights, 0)
        # Spectral clustering per Ng et al.: embed with the top eigenvectors of the normalized similarity matrix, then k-means.
        inv_sqrt_degree = 1 / np.sqrt(np.maximum(np.sum(weights, axis=1), 1e-12))
        k = min(10, len(valid_symbols))
        embedding = np.linalg.eigh(weights * np.outer(inv_sqrt_degree, inv_sqrt_degree))[1][:, -k:]
        embedding = embedding / np.maximum(np.linalg.norm(embedding, axis=1, keepdims=True), 1e-12)
        labels = KMeans(n_clusters=k, n_init=10, random_state=0).fit_predict(embedding)
        # Iterate over all asset combinations within each cluster to locate cointegrated candidates.
        candidates = []
        for cluster in [[valid_symbols[i] for i in np.flatnonzero(labels == label)] for label in range(k)]:
            for i in range(len(cluster)):
                for j in range(i + 1, len(cluster)):
                    candidate = self._try_create_pair(log_prices, cluster[i], cluster[j])
                    if candidate:
                        candidates.append(candidate)
        # Keep the pairs with the strongest cointegration evidence, allowing each symbol in at most one pair for diversification.
        selected = []
        used_symbols = []
        for _, pair_spread in sorted(candidates, key=lambda candidate: candidate[0]):
            if len(selected) == self._max_pairs:
                break
            if pair_spread.symbol_a in used_symbols or pair_spread.symbol_b in used_symbols:
                continue
            selected.append(pair_spread)
            used_symbols += pair_spread.symbols
        return selected

    def _try_create_pair(self, log_prices, asset_a, asset_b):
        y = log_prices[asset_a]
        x = log_prices[asset_b]
        if np.all(x == x.iloc[0]):
            return None
        # Determine the hedge ratio using an ordinary least squares regression of the log prices.
        hedge_ratio = float(np.polyfit(x, y, 1)[0])
        # Require a positive two-sided hedge so the trade is a genuine relative-value pair rather than a near-naked single-name bet.
        if not 0.25 <= hedge_ratio <= 4:
            return None
        spread = (y - hedge_ratio * x).values
        sd = float(np.std(spread))
        # Require enough spread volatility for a two-sigma entry to clear trading costs, filtering near-duplicates like dual share classes.
        if sd < 0.01:
            return None
        # Require a mean-reversion half-life inside the trading horizon, estimated with a lagged spread regression.
        slope = float(np.polyfit(spread[:-1], np.diff(spread), 1)[0])
        if slope >= 0 or -np.log(2) / slope > 30:
            return None
        # Verify spread stationarity with the Engle-Granger test, whose critical values account for the estimated hedge ratio.
        p_value = coint(y, x)[1]
        if p_value >= 0.05:
            return None
        return p_value, PairSpread(asset_a, asset_b, hedge_ratio, float(np.mean(spread)), sd)