# Taken from https://www.quantconnect.com/forum/discussion/3377/momentum-strategy-with-market-cap-and-ev-ebitda
# Created by Jing Wu
# Edited by Nathan Wells trying to mirror Original by: Christopher Cain, CMT & Larry Connors
#Posted here: https://www.quantopian.com/posts/new-strategy-presenting-the-quality-companies-in-an-uptrend-model-1



from clr import AddReference
AddReference("System.Core")
AddReference("System.Collections")
AddReference("QuantConnect.Common")
AddReference("QuantConnect.Algorithm")

from System import *
from System.Collections.Generic import List
from QuantConnect import *
from QuantConnect.Algorithm import QCAlgorithm
from QuantConnect.Data.UniverseSelection import *
from QuantConnect.Indicators import *

import math
import numpy as np
import pandas as pd
import scipy as sp
# import statsmodels.api as sm

class FundamentalFactorAlgorithm(QCAlgorithm):

    def Initialize(self):

        self.SetStartDate(2017, 1, 1)  #Set Start Date
        self.SetEndDate(2019, 12, 31)  #Set End Date       
        self.SetCash(10000)            #Set Strategy Cash
    
        #changed from Daily to Monthly
        self.UniverseSettings.Resolution = Resolution.Daily
        #self.AddUniverse(self.Universe.Index.QC500)
        #self.AddUniverse(self.Universe.Index.QC500, self.FineSelectionFunction)
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        #changed from Minuite to Daily
        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol 
        self.holding_months = 1
        self.num_screener = 100
        self.num_stocks = 5
        self.formation_days = 126
        self.lowmom = False
        self.month_count = self.holding_months
        self.Schedule.On(self.DateRules.MonthEnd("SPY"), self.TimeRules.AfterMarketOpen("SPY", 1), Action(self.monthly_rebalance))
        self.Schedule.On(self.DateRules.MonthEnd("SPY"), self.TimeRules.AfterMarketOpen("SPY", 1), Action(self.rebalance))
        # rebalance the universe selection once a month
        self.rebalence_flag = 0
        # make sure to run the universe selection at the start of the algorithm even it's not the manth start
        self.first_month_trade_flag = 1
        self.trade_flag = 0 
        self.symbols = None
        self.periodCheck = -1
        self.symboldict = {}
 
    def CoarseSelectionFunction(self, coarse):
        if self.rebalence_flag or self.first_month_trade_flag:
            # drop stocks which have no fundamental data or have too low prices
            #selected = [x for x in coarse if (x.HasFundamentalData) and (float(x.Price) > 10)]
            selected = [x for x in coarse if (x.HasFundamentalData) and (float(x.Price) > 10)]
            topDollarVolume = sorted(selected, key=lambda k : k.DollarVolume, reverse=True)[:1500]
            
            return [ x.Symbol for x in topDollarVolume]
        else:
            return self.symbols


    def FineSelectionFunction(self, fine):
        if self.rebalence_flag or self.first_month_trade_flag:
            
                
            #self.periodCheck = algorithm.Time.year
            # Filter by Market Capitalization and USA
            filtered = [f for f in fine if f.CompanyReference.CountryId == "USA"
                                        and f.CompanyReference.PrimaryExchangeID in ["NYS","NAS"]
                                        and f.MarketCap > 5e8]
            #filter_market_cap = [f for f in fine if f.MarketCap > 500000000]
            
            # Filter for top quality
            top_quality = sorted(filtered, key=lambda x: x.OperationRatios.ROIC.ThreeMonths + x.OperationRatios.LongTermDebtEquityRatio.ThreeMonths + (x.ValuationRatios.CashReturn + x.ValuationRatios.FCFYield), reverse=True)[:60]

            # When we get new symbols, we add them to the dict and warm up the indicator
            symbols = [x.Symbol for x in top_quality if x.Symbol not in self.symboldict]
            history = self.History(symbols, 146, Resolution.Daily)
            if not history.empty:
                history = history.close.unstack(0)
                for symbol in symbols:
                    if str(symbol) not in history:
                        continue
                    
                    df = history[symbol].dropna()
                    if not df.empty:
                        self.symboldict[symbol] = SymbolData(self, df)
                        
            # Now, we update the dictionary with the latest data
            for x in fine:
                symbol = x.Symbol
                if symbol in self.symboldict:
                    self.symboldict[symbol].Update(x.EndTime, x.Price)
            
            topMOM = sorted(self.symboldict.items(), key=lambda x: x[1].DeltaMOM, reverse=True)[:10]
            #return [x[0] for x in topMOM]
            
            #self.symbols = [x.Symbol for x in topMOM]
            self.symbols = [x[0] for x in topMOM]
            self.rebalence_flag = 0
            self.first_month_trade_flag = 0
            self.trade_flag = 1
            return self.symbols
        else:
            return self.symbols
    
    def OnData(self, data):
        pass
    
    def monthly_rebalance(self):
        self.rebalence_flag = 1

    def rebalance(self):
        #Looks like this sells if they drop below the mean of SPY, so I disabled it
        #spy_hist = self.History([self.spy], self.formation_days, Resolution.Daily).loc[str(self.spy)]['close']
        #if self.Securities[self.spy].Price < spy_hist.mean():
        #    for symbol in self.Portfolio.Keys:
        #        if symbol.Value != "TLT":
        #            self.Liquidate()
        #    self.AddEquity("TLT")
        #    self.SetHoldings("TLT", 1)
        #    return

        if self.symbols is None: return
        chosen_df = self.calc_return(self.symbols)
        chosen_df = chosen_df.iloc[:self.num_stocks]
        
        self.existing_pos = 0
        add_symbols = []
        for symbol in self.Portfolio.Keys:
            #if symbol.Value == 'SPY': continue
            if (symbol.Value not in chosen_df.index):  
                #self.SetHoldings(symbol, 0)
                self.Liquidate(symbol)
            elif (symbol.Value in chosen_df.index): 
                self.existing_pos += 1
            
        weight = 0.99/len(chosen_df)
        for symbol in chosen_df.index:
            self.AddEquity(symbol)
            self.SetHoldings(symbol, weight)    
                
    def calc_return(self, stocks):
        #Need to change this to just be an uptrend or downtrend...and buy bonds in downtrend.
        hist = self.History(stocks, self.formation_days, Resolution.Daily)
        current = self.History(stocks, 10, Resolution.Daily)
        
        self.price = {}
        ret = {}
     
        for symbol in stocks:
            if str(symbol) in hist.index.levels[0] and str(symbol) in current.index.levels[0]:
                self.price[symbol.Value] = list(hist.loc[str(symbol)]['close'])
                self.price[symbol.Value].append(current.loc[str(symbol)]['close'][0])
        
        for symbol in self.price.keys():
            ret[symbol] = (self.price[symbol][-1] - self.price[symbol][0]) / self.price[symbol][0]
        df_ret = pd.DataFrame.from_dict(ret, orient='index')
        df_ret.columns = ['return']
        sort_return = df_ret.sort_values(by = ['return'], ascending = self.lowmom)
        
        return sort_return

class SymbolData:
    def __init__(self, symbol, history):
        self.mom10 = Momentum(10)
        self.mom146 = Momentum(146)

        for time, close in history.iteritems():
            self.Update(time, close)

    def Update(self, time, close):
        self.mom10.Update(time, close)
        self.mom146.Update(time, close)


    @property
    def DeltaMOM(self):
        return self.mom10.Current.Value - self.mom146.Current.Value
    
    def __repr__(self):
        return f'{self.DeltaMOM}'