# Filename: D_Volatility_Anomaly.py
from AlgorithmImports import *
import numpy as np

class DVolatilityAnomaly(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2020, 1, 1)  # Start after SPY inception and enough data for 252d lookback
        self.SetEndDate(2022, 1, 1)
        self.SetCash(100000)
        self.rebalance_flag = False
        self.lookback = 252
        self.num_stocks = 1000
        self.quintile = 0.2  # Lowest 20%
        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction)
        self.symbol_data = {}
        self.selected = []
        self.last_month = -1
        self.Schedule.On(self.DateRules.MonthStart(), self.TimeRules.AfterMarketOpen("SPY", 30), self.Rebalance)
        self.rolling_volatility = RollingWindow[float](12)  # For 12 months of realized volatility
        self.rolling_volatility_spy = RollingWindow[float](12)
        self.spy_returns = []
        self.portfolio_returns = []
        self.cum_returns = []
        self.cum_returns_spy = []
        self.last_portfolio_value = self.Portfolio.TotalPortfolioValue
        self.last_spy_price = None

    def CoarseSelectionFunction(self, coarse):
        # Filter for top 1000 by dollar volume, with price > $5
        filtered = [x for x in coarse if x.HasFundamentalData and x.Price > 5 and x.Volume > 0]
        sorted_by_dollar = sorted(filtered, key=lambda x: x.DollarVolume, reverse=True)
        return [x.Symbol for x in sorted_by_dollar[:self.num_stocks]]

    def OnSecuritiesChanged(self, changes):
        # Warm up history for new securities
        for added in changes.AddedSecurities:
            symbol = added.Symbol
            if symbol not in self.symbol_data:
                history = self.History(symbol, self.lookback, Resolution.Daily)
                closes = list(history["close"]) if not history.empty else []
                self.symbol_data[symbol] = closes
        # Remove data for dropped securities
        for removed in changes.RemovedSecurities:
            symbol = removed.Symbol
            if symbol in self.symbol_data:
                del self.symbol_data[symbol]

    def OnData(self, data):
        # Update rolling close prices for all tracked symbols
        for symbol in list(self.symbol_data.keys()):
            if symbol in data and data[symbol] is not None and hasattr(data[symbol], 'Close') and data[symbol].Close > 0:
                self.symbol_data[symbol].append(float(data[symbol].Close))
                if len(self.symbol_data[symbol]) > self.lookback:
                    self.symbol_data[symbol].pop(0)
            # Remove if no data for a while
            elif len(self.symbol_data[symbol]) == 0 or (symbol not in data):
                continue
        # For SPY, update rolling window for volatility and returns
        if self.spy in data and data[self.spy] is not None and hasattr(data[self.spy], 'Close'):
            close = float(data[self.spy].Close)
            if self.last_spy_price is not None:
                ret = (close - self.last_spy_price) / self.last_spy_price
                self.spy_returns.append(ret)
                if len(self.spy_returns) > self.lookback:
                    self.spy_returns.pop(0)
            self.last_spy_price = close
            # Compute rolling 12-month volatility for SPY
            if len(self.spy_returns) == self.lookback:
                vol = np.std(self.spy_returns) * np.sqrt(252)
                self.rolling_volatility_spy.Add(vol)
        # Compute portfolio return
        current_value = self.Portfolio.TotalPortfolioValue
        if self.last_portfolio_value > 0:
            ret = (current_value - self.last_portfolio_value) / self.last_portfolio_value
        else:
            ret = 0
        self.portfolio_returns.append(ret)
        if len(self.portfolio_returns) > self.lookback:
            self.portfolio_returns.pop(0)
        self.last_portfolio_value = current_value
        # Compute rolling 12-month volatility for portfolio
        if len(self.portfolio_returns) == self.lookback:
            vol = np.std(self.portfolio_returns) * np.sqrt(252)
            self.rolling_volatility.Add(vol)
        # Track cumulative returns for plotting
        if len(self.cum_returns) == 0:
            self.cum_returns.append(1)
        else:
            self.cum_returns.append(self.cum_returns[-1] * (1 + ret))
        if self.last_spy_price is not None and len(self.cum_returns_spy) == 0:
            self.cum_returns_spy.append(1)
        elif self.last_spy_price is not None and len(self.spy_returns) > 0:
            self.cum_returns_spy.append(self.cum_returns_spy[-1] * (1 + self.spy_returns[-1]))

    def Rebalance(self):
        # Only rebalance at the start of a new month
        if self.Time.month == self.last_month:
            return
        self.last_month = self.Time.month
        # Calculate volatility for each stock with enough data
        vol_dict = {}
        for symbol, closes in self.symbol_data.items():
            if len(closes) == self.lookback:
                returns = np.diff(closes) / closes[:-1]
                vol = np.std(returns) * np.sqrt(252)
                vol_dict[symbol] = vol
        if len(vol_dict) == 0:
            return
        # Select lowest quintile by volatility
        sorted_syms = sorted(vol_dict.items(), key=lambda x: x[1])
        num_select = max(1, int(len(sorted_syms) * self.quintile))
        selected_syms = [x[0] for x in sorted_syms[:num_select]]
        self.selected = selected_syms
        # Liquidate unselected positions
        for kvp in self.Portfolio:
            symbol = kvp.Key
            if symbol != self.spy and kvp.Value.Invested and symbol not in self.selected:
                self.Liquidate(symbol)
        # Equal weight allocation
        if len(self.selected) == 0:
            return
        weight = 1.0 / len(self.selected)
        for symbol in self.selected:
            if self.Securities.ContainsKey(symbol) and self.Securities[symbol].IsTradable:
                self.SetHoldings(symbol, weight)
        # Ensure no leverage
        for symbol in self.selected:
            if self.Securities.ContainsKey(symbol):
                self.Securities[symbol].SetLeverage(1)
        if self.Securities.ContainsKey(self.spy):
            self.Liquidate(self.spy)  # Do not hold SPY
        # Record stats for plotting
        if self.rolling_volatility.Count > 0 and self.rolling_volatility_spy.Count > 0:
            self.Plot("Volatility", "Portfolio 12m Vol", self.rolling_volatility[0])
            self.Plot("Volatility", "SPY 12m Vol", self.rolling_volatility_spy[0])
        if len(self.cum_returns) > 0 and len(self.cum_returns_spy) > 0:
            self.Plot("Cumulative Return", "Portfolio", self.cum_returns[-1])
            self.Plot("Cumulative Return", "SPY", self.cum_returns_spy[-1])