from AlgorithmImports import *

class OLMARAlgorithm(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2024, 1, 1)   # Set Start Date
        self.SetEndDate(2024, 5, 1)     # Set End Date
        self.SetCash(100000)            # Set Strategy Cash
        
        self.epsilon = 10.0             # Control parameter for mean reversion strength
        self.windowSize = 5             # Window size for moving average
        self.rebalanceTime = datetime.min
        
        self.symbolData = {}
        self.AddEquity("SPY", Resolution.Daily)  # Use SPY to ensure valid market times for scheduling
        self.AddUniverse(self.CoarseSelectionFunction)
        
        self.Schedule.On(self.DateRules.WeekStart(), self.TimeRules.AfterMarketOpen("SPY", 30), self.Rebalance)
    
    def CoarseSelectionFunction(self, coarse):
        # Select the top 500 stocks by dollar volume
        selected = sorted([x for x in coarse if x.HasFundamentalData], key=lambda x: x.DollarVolume, reverse=True)[:500]
        return [x.Symbol for x in selected]
    
    def OnSecuritiesChanged(self, changes):
        for security in changes.AddedSecurities:
            symbol = security.Symbol
            self.symbolData[symbol] = SymbolData(symbol, self.windowSize)
            history = self.History(symbol, self.windowSize, Resolution.Daily)
            if not history.empty:
                for time, row in history.loc[symbol].iterrows():
                    self.symbolData[symbol].PriceWindow.Add(row['close'])
        
        for security in changes.RemovedSecurities:
            symbol = security.Symbol
            if symbol in self.symbolData:
                del self.symbolData[symbol]
    
    def OnData(self, data):
        if self.Time < self.rebalanceTime:
            return
        
        for symbol, symbolData in self.symbolData.items():
            if data.ContainsKey(symbol) and data[symbol] is not None and data[symbol].Price is not None:
                price = data[symbol].Price
                symbolData.PriceWindow.Add(price)
        
    def Rebalance(self):
        if not all([symbolData.PriceWindow.IsReady for symbolData in self.symbolData.values()]):
            return
        
        weights = {}
        totalWeight = 0
        
        for symbol, symbolData in self.symbolData.items():
            averagePrice = sum(symbolData.PriceWindow) / self.windowSize
            predictedPrice = self.PredictNextPrice(averagePrice)
            weight = self.CalculateWeights(predictedPrice, symbolData.PriceWindow[0])
            
            weights[symbol] = weight
            totalWeight += weight
        
        for symbol in weights:
            weights[symbol] /= totalWeight
        
        for symbol, weight in weights.items():
            self.SetHoldings(symbol, weight)
        
        self.rebalanceTime = self.Time + timedelta(days=1)
    
    def PredictNextPrice(self, averagePrice):
        # Use the simple moving average as the predicted price for the next period
        return averagePrice
    
    def CalculateWeights(self, predictedPrice, currentPrice):
        reversionRatio = predictedPrice / currentPrice
        
        # OLMAR formula: target weight = epsilon * (predictedPrice / currentPrice - 1)
        weight = self.epsilon * (reversionRatio - 1)
        
        # Ensure weight is within bounds [0, 1]
        return max(0, min(1, weight))

class SymbolData:
    def __init__(self, symbol, windowSize):
        self.Symbol = symbol
        self.PriceWindow = RollingWindow[float](windowSize)

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

class StationarySelectionModel(ETFConstituentsUniverseSelectionModel):
    def __init__(self, algorithm, etf, lookback = 10, universe_settings = None):
        self.algorithm = algorithm
        self.lookback = lookback
        self.symbolData = {}

        symbol = Symbol.Create(etf, SecurityType.Equity, Market.USA)
        super().__init__(symbol, universe_settings, self.ETFConstituentsFilter)

    def ETFConstituentsFilter(self, constituents):
        stationarity = {}
        self.algorithm.Debug(f"{self.algorithm.Time}::{len(list(constituents))} tickers in SPY")

        for c in constituents:
            symbol = c.Symbol
            if symbol not in self.symbolData:
                self.symbolData[symbol] = SymbolData(self.algorithm, symbol, self.lookback)
            data = self.symbolData[symbol]

            # Update with the last price
            self.algorithm.Debug(f"{self.algorithm.Time}::{symbol}::{c.MarketValue}::{c.SharesHeld}")
            if c.MarketValue and c.SharesHeld:
                price = c.MarketValue / c.SharesHeld
                data.Update(price)
            # Cache the stationarity test statistics in the dict
            if data.TestStatistics:
                stationarity[symbol] = data.TestStatistics

        # Return the top 10 lowest test statistics stocks (more negative stat means higher prob to have no unit root)
        selected = sorted(stationarity.items(), key=lambda x: x[1])
        return [x[0] for x in selected[:10]]

class SymbolData:
    def __init__(self, algorithm, symbol, lookback):
        # RollingWindow to hold log price series for stationary testing
        self.window = RollingWindow[float](lookback)
        self.model = None

        # Warm up RollingWindow
        history = algorithm.History[TradeBar](symbol, lookback, Resolution.Daily)
        for bar in list(history)[:-1]:
            self.window.Add(np.log(bar.Close))

    def Update(self, value):
        if value == 0: return

        # Update RollingWindow with log price
        self.window.Add(np.log(value))
        if self.window.IsReady:
            # Test stationarity for log price series by augmented dickey-fuller test
            price = np.array(list(self.window))[::-1]
            self.model = adfuller(price, regression='n', autolag='BIC')

    @property
    def TestStatistics(self):
        return self.model[0] if self.model else None