from datetime import timedelta
from AlgorithmImports import *
import pandas as pd
import numpy as np


class EquityVolatilityTargetingUniverse(QCAlgorithm):

    def initialize(self):
        self.set_start_date(self.end_date - timedelta(days=5 * 365))
        self.set_cash(200000)
        self.settings.seed_initial_prices = True

        self.universe_settings.resolution = Resolution.DAILY

        # Universe: SPY constituents
        self._universe = self.add_universe(self.universe.etf("SPY"), self._select_spy)

        # Rebalance daily at 08:00 ET so universe selection is current
        self.schedule.on(
            self.date_rules.every_day("SPY"),
            self.time_rules.at(8, 0),
            self._rebalance,
        )

    def _select_spy(self, constituents):
        return [c.symbol for c in constituents]

    def _rebalance(self):
        symbols = list(self._universe.selected)
        if not symbols:
            return

        # Request 31 days of daily bars to compute 30 daily returns
        history = self.history(TradeBar, symbols, 31, Resolution.DAILY)
        if history.empty:
            return

        # Pivot to time x symbol close-price matrix
        close_prices = history["close"].unstack(level=0)
        if close_prices.empty:
            return

        # Daily returns
        returns = close_prices.pct_change().dropna()
        if returns.empty:
            return

        # Annualized realized volatility (std of daily returns * sqrt(252))
        vol = returns.std() * np.sqrt(252)

        # Keep only names below 15 % annualized vol
        low_vol = vol[vol < 0.15]
        if low_vol.empty:
            return

        # Lowest-vol 30, equal-weighted
        targets = low_vol.nsmallest(30).index.tolist()
        if not targets:
            return

        # Use a 2 % cash buffer to avoid margin exhaustion during rebalance
        weight = 0.98 / len(targets)

        # Build explicit targets for every symbol (current + new) so LEAN
        # computes order deltas instead of relying on liquidate_existing_holdings.
        all_symbols = set(targets)
        for kvp in self.portfolio:
            all_symbols.add(kvp.key)

        target_list = []
        for symbol in all_symbols:
            if symbol in targets:
                target_list.append(PortfolioTarget(symbol, weight))
            else:
                target_list.append(PortfolioTarget(symbol, 0.0))

        self.set_holdings(target_list, liquidate_existing_holdings=False)