import pandas as pd
import numpy as np
from datetime import datetime, timedelta
from QuantConnect.Data.UniverseSelection import *
from scipy.stats import linregress

class Momentum(QCAlgorithm):

    def __init__(self):
        self.reb1 = 1 # set the flag for rebalance
        self.reb2 = 12 # set the flag for rebalance
        self.num_coarse = 20 # Number of stocks to pass CoarseSelection process
        self.num_fine = 20 # Number of stocks to long
        self.symbols = None
        
    def Initialize(self):
        self.SetStartDate(2012, 1, 1) # Set Start Date
        self.SetEndDate(2013,6, 1) # Set End Date
        self.SetCash(150000) # Set Strategy Cash
        
        self.btal = self.AddEquity("BTAL", Resolution.Minute).Symbol # BTAL hedge
        self.iei = self.AddEquity("IEI", Resolution.Minute).Symbol # bond hedge
        self.tlt = self.AddEquity("TLT", Resolution.Minute).Symbol # bond hedge
        self.list = [self.btal, self.iei, self.tlt]

        self.reb1 = 1 # set the flag for momentum stock rebalancement
        self.reb2 = 1 # set the flag for rebalance
        self.AddUniverse(self.CoarseSelectionFunction,self.FineSelectionFunction)
        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol
        self.Schedule.On(self.DateRules.MonthStart(self.spy), 
        self.TimeRules.AfterMarketOpen(self.spy,5), Action(self.rebalance))
        self.SetSecurityInitializer(self.CustomSecurityInitializer)
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Cash)
        
    def CustomSecurityInitializer(self, security):
        security.SetDataNormalizationMode(DataNormalizationMode.SplitAdjusted)

    def CoarseSelectionFunction(self, coarse):
    # if the rebalance flag is not 1, return null list to save time.
        if self.reb1 != 1:
            return Universe.Unchanged
        
        # make universe selection once a month
        sortedByDollarVolume = sorted(coarse, key=lambda x: x.DollarVolume, reverse=True)  
        filtered = [x.Symbol for x in sortedByDollarVolume if x.HasFundamentalData]
        
        # filtered down to the 500 most liquid stocks
        return filtered[:self.num_coarse]
        
    def FineSelectionFunction(self, fine):
        # return null list if it's not time to rebalance
        if self.reb1 != 1:
            return Universe.Unchanged 
            
        # drop counter (will update back to 1 after rebalancement has occurred)
        self.reb1 = 0
        
        #will not include this portion of code - essentially filters for top momentum stocks
        # self.long = sorted_symbolLong[:self.num_fine]
        # return self.long
        return [f.Symbol for f in fine]

    def OnData(self, data):
        pass

    def rebalance(self):
        long_list = [x for x in self.ActiveSecurities.Keys]
        
        # to update the momentum stocks only monthly
        for i in self.Portfolio.Values: 
            if i.Invested and i.Symbol not in long_list and i.Symbol not in self.list: 
                self.Liquidate(i.Symbol)
                
        # annual portfolio rebalancement
        if self.reb2 % 12 == 0:
            
            self.SetHoldings(self.iei, 0.15)
            self.SetHoldings(self.tlt, 0.15)
            self.SetHoldings(self.btal, 0.10)
            
            for i in long_list: 
                self.SetHoldings(i, 0.6/self.num_fine)
                
        else: 
        # monthly rebalancement of momentum stocks only, leaving bonds and market neutral etfs untouched
            for i in long_list: 
                if not self.Securities[i].Invested: 
                    self.SetHoldings(i, 0.6/self.num_fine)
            
        self.reb1 = 1
        self.reb2 += 1