# region imports
from AlgorithmImports import *
import numpy as np
# endregion


class GreedySharpeAIInfrastructure(QCAlgorithm):
    """
    Cross-Sectional Momentum Strategy — AI Infrastructure Universe

    Source: greedy_sharpe_ai_infrastructure_universe_filtered_20260306_120902

    Signal:  180-day lookback, 20-day skip momentum
    Select:  Top 95th percentile of cross-sectional momentum
    Weight:  Equal weight among selected tickers
    Rebal:   Daily at market close
    Filter:  $2 minimum entry price
    """

    # ── Strategy parameters ──────────────────────────────────────────────
    LOOKBACK = 180          # momentum lookback window (trading days)
    SKIP = 20               # skip most recent N days (mean-reversion filter)
    PERCENTILE = 0.95       # select top 5% of momentum scores
    MIN_ENTRY_PRICE = 2.0   # minimum price to enter a position

    # ── Universe from greedy search results ──────────────────────────────
    TICKERS = [
        "AAOI", "ACLS", "ADI", "AES", "ALAB", "AMAT", "AMKR", "AMZN",
        "ANET", "APH", "APLD", "ASML", "AVGO", "AXTI", "BDC", "CALX",
        "CAMT", "CAT", "CDNS", "CLSK", "COHR", "CORZ", "CRM", "CRUS",
        "CSCO", "DOV", "DUK", "EME", "ENTG", "ETN", "FIX", "FLEX",
        "FN", "GEV", "GFS", "GLW", "GNRC", "GOOG", "HUBB", "HUT",
        "IBM", "IESC", "INTC", "IPGP", "J", "KLAC", "LEU", "LITE",
        "LRCX", "MARA", "META", "MKSI", "MOD", "MPWR", "MRVL", "MSFT",
        "MTSI", "MU", "NEE", "NTAP", "NVDA", "NVMI", "NVT", "NXPI",
        "OKLO", "ONTO", "PCG", "PSTG", "PWR", "QCOM", "QRVO", "SMR",
        "SNPS", "SO", "STX", "SWKS", "TER", "TSM", "TT", "TXN",
        "UCTT", "VIAV", "VRT",
    ]

    def initialize(self):
        self.set_start_date(2019, 1, 1)
        self.set_end_date(2026, 12, 31)
        self.set_cash(50_000)

        # Use daily resolution
        self.universe_settings.resolution = Resolution.DAILY

        # Add all universe symbols
        self.symbols = []
        for ticker in self.TICKERS:
            symbol = self.add_equity(ticker, Resolution.DAILY).symbol
            self.symbols.append(symbol)

        # History window: need LOOKBACK days of data before we can trade
        self.set_warm_up(self.LOOKBACK + self.SKIP + 5, Resolution.DAILY)

        # Schedule daily rebalance 5 minutes before market close
        self.schedule.on(
            self.date_rules.every_day(),
            self.time_rules.before_market_close("SPY", 5),
            self.rebalance,
        )

        # Benchmark
        self.set_benchmark("SPY")

        # Track for logging
        self._last_log_date = None

    def rebalance(self):
        if self.is_warming_up:
            return

        # ── 1. Get historical adjusted close prices ──────────────────
        lookback_total = self.LOOKBACK + self.SKIP + 1
        history = self.history(self.symbols, lookback_total, Resolution.DAILY)

        if history.empty:
            return

        # Pivot to get close prices by symbol
        try:
            close = history["close"].unstack(level=0)
        except Exception:
            return

        if close.shape[0] < lookback_total:
            return

        # ── 2. Compute cross-sectional momentum ─────────────────────
        # momentum = price[t - SKIP] / price[t - LOOKBACK] - 1
        end_prices = close.iloc[-(self.SKIP + 1)]       # price SKIP days ago
        start_prices = close.iloc[0]                      # price LOOKBACK+SKIP days ago

        momentum = {}
        for symbol in close.columns:
            p_end = end_prices.get(symbol, None)
            p_start = start_prices.get(symbol, None)
            if p_end is not None and p_start is not None:
                if not np.isnan(p_end) and not np.isnan(p_start) and p_start > 0:
                    momentum[symbol] = p_end / p_start - 1.0

        if not momentum:
            self.liquidate()
            return

        # ── 3. Select top percentile ─────────────────────────────────
        mom_values = list(momentum.values())
        threshold = np.percentile(mom_values, self.PERCENTILE * 100)

        # Get current prices for min-price filter
        selected = []
        for symbol, mom in momentum.items():
            if mom >= threshold:
                current_price = self.securities[symbol].price
                if current_price >= self.MIN_ENTRY_PRICE:
                    selected.append(symbol)

        if not selected:
            self.liquidate()
            return

        # ── 4. Equal weight allocation ───────────────────────────────
        weight = 1.0 / len(selected)
        selected_set = set(selected)

        # Liquidate positions not in the new selection
        for holding in self.portfolio.Values:
            if holding.invested and holding.symbol not in selected_set:
                self.liquidate(holding.symbol)

        # Set target allocations
        for symbol in selected:
            self.set_holdings(symbol, weight)

        # ── 5. Monthly logging ───────────────────────────────────────
        today = self.time.date()
        if self._last_log_date is None or today.month != self._last_log_date.month:
            self._last_log_date = today
            portfolio_value = self.portfolio.total_portfolio_value
            n_holdings = sum(1 for h in self.portfolio.Values if h.invested)
            self.log(
                f"{today} | Value: ${portfolio_value:,.0f} | "
                f"Holdings: {n_holdings}/{len(selected)} selected | "
                f"Threshold: {threshold:.4f}"
            )

    def on_data(self, data: Slice):
        pass