# https://quantpedia.com/Screener/Details/15
import pandas as pd
import numpy as np
from statistics import mean
from datetime import datetime
from collections import defaultdict

class CountryEquityIndexesMomentumAlgorithm(QCAlgorithm):

    def Initialize(self):

        self.SetStartDate(2006, 1, 1)  
        self.SetEndDate(datetime.now())  
        self.SetCash(10000) 
        # create a dictionary to store momentum indicators for all symbols 
        self.CombinedMomentum = {}
        self.inv_vola = {}
        self.Weights = {}
        
        self.symbols = ["SPY","IWM", "QQQ", "EFA", "EEM","VNQ", "LQD", "GLD", "SHY", "IEF", "TLT", "AGG"] 

        # warm up the MOM indicator
        self.SetWarmUp(252)
        
        for symbol in self.symbols:
            self.AddEquity(symbol, Resolution.Daily)
            self.Securities[symbol].SetDataNormalizationMode(DataNormalizationMode.TotalReturn)
            self.CombinedMomentum[symbol] = CombinedMomentum(self, symbol)
        
        
        #schedule the function to fire at the month start 
        
        def getMonthLastTradingDay(self):
            month_last_day = DateTime(self.Time.year, self.Time.month, DateTime.DaysInMonth(self.Time.year, self.Time.month))
            tradingDays = self.TradingCalendar.GetDaysByType(TradingDayType.BusinessDay, DateTime(self.Time.year, self.Time.month, 1), month_last_day)
            tradingDays = [day.Date.date() for day in tradingDays]
            return tradingDays[-1]
        
        #self.Schedule.On(self.DateRules.MonthEnd("SPY"), self.TimeRules.AfterMarketOpen("SPY", 60), self.Rebalance)
            
    def OnData(self, data):
        if self.IsWarmingUp: return

    def Rebalance(self):
        #Determine the sec. to invest in at the last trading day of the month
        if (self.Time.date() == self.getMonthLastTradingDay()) and (self.Time.hour == 15):
            top = pd.Series(self.CombinedMomentum).sort_values(ascending = False)[:5]
            #Determine the weights of the portfolio
            # for symbol in top.index:
            #     self.DailyReturn[symbol] = self.LOGR(symbol, 1, Resolution.Daily)
            #     self.inv_vola[symbol] = 1/(self.STD(self.DailyReturn[symbol], 63, Resolution.Daily))
            #     self.Weights[symbol] = self.inv_vola[symbol].Current.Value/ self.SUM(self.inv_vola.Current.Values)
                
            for kvp in self.Portfolio:
                security_hold = kvp.Value
                # liquidate the security which is no longer in the top momentum list
                if security_hold.Invested and (security_hold.Symbol.Value not in top.index):
                    self.Liquidate(security_hold.Symbol)
            #Buy the new sucurities
            added_symbols = []        
            for symbol in top.index:
                if not self.Portfolio[symbol].Invested:
                    added_symbols.append(symbol)
            for added in added_symbols:
                self.SetHoldings(added, 1/5)
    
class CombinedMomentum():
    def __init__(self, algo, symbol):
        
        self.MOMP_1M = algo.MOMP(symbol, 21, Resolution.Daily)
        self.MOMP_3M = algo.MOMP(symbol, 63, Resolution.Daily)
        self.MOMP_6M = algo.MOMP(symbol, 126, Resolution.Daily)
        self.MOMP_12M = algo.MOMP(symbol, 252, Resolution.Daily)
        
    def getValue(self):
        value = (self.MOMP_1M.Current.Value + self.MOMP_3M.Current.Value + self.MOMP_6M.Current.Value + self.MOMP_12M.Current.Value)/4
        
        return value
        
# class InverseVolatility():
    
#     def __init__(self,algo,symbol):
#         self.DailyReturn = algo.LOGR(symbol, 1, Resolution.Daily)
        
#     def getValue(self):
#         value = 1/(self.STD(self.DailyReturn, 63, Resolution.Daily).Current.Value)
#         return value