from AlgorithmImports import *

class RollingWindowAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2003, 1, 2)
        self.SetEndDate(2015, 1, 2)
        self.SetCash(100000)

        # Define the period for the KER calculation
        self.kerPeriod = 20

        # Create a dictionary to hold the KER indicators and rolling windows for each symbol
        self.kerDict = {}
        self.kerWinDict = {}

        # Add 20 equities at daily resolution
        symbols = ["AAPL", "MSFT", "AMZN", "FB", "GOOG", "GOOGL", "JNJ", "JPM", "V", "PG",
                   "MA", "UNH", "HD", "DIS", "BAC", "VZ", "KO", "CSCO", "PFE", "CVX"]
        for symbol in symbols:
            self.AddEquity(symbol, Resolution.Daily)
            self.kerDict[symbol] = self.KER(symbol, self.kerPeriod)
            self.kerWinDict[symbol] = RollingWindow[float](self.kerPeriod)
            self.SetWarmup(self.kerPeriod * 2, Resolution.Daily)

        # Schedule rebalance every month for each symbol
        for symbol in symbols:
            self.Schedule.On(self.DateRules.MonthEnd(symbol), self.TimeRules.AfterMarketOpen(symbol), lambda: self.Rebalance(symbol))

    def OnData(self, data):
        for symbol in self.kerDict.keys():
            # Wait for the window to be ready for each symbol
            if not self.kerDict[symbol].IsReady:
                return

            # Add current KER value to the rolling window for each symbol
            self.kerWinDict[symbol].Add(self.kerDict[symbol].Current.Value)

    def Rebalance(self, symbol):
        # Wait for the window to be ready for each symbol
        if not self.kerWinDict[symbol].IsReady:
            return

        # Calculate the average KER for the last year for each symbol
        avgKER = np.mean(list(self.kerWinDict[symbol]))
        self.Log("Average Kaufman Efficiency Ratio for {} for last {} days: {:.4f}".format(symbol, self.kerPeriod, avgKER))

        # Reset the window for each symbol for the next period
        self.kerWinDict[symbol].Reset()

    def OnEndOfAlgorithm(self):
        for symbol in self.kerDict.keys():
            # Calculate the final average KER for each symbol
            avgKER = np.mean(list(self.kerWinDict[symbol]))
            self.Log("Final average Kaufman Efficiency Ratio for {} for last {} days: {:.4f}".format(symbol, self.kerPeriod, avgKER))