from datetime import timedelta
from QuantConnect.Data.UniverseSelection import * 
from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel

class TheStockPicker(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2006, 1, 1)   #15 year backtest
        #self.SetEndDate(2020, 7, 15)
        self.SetCash(100000)
        self.SetBenchmark('SPY')
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage)
        self.UniverseSettings.Resolution = Resolution.Daily
        #self.UniverseSettings.Leverage = 1
        self.AddUniverseSelection(StockPickerUniverseSelectionModel())
        
        #spy = [ Symbol.Create("SPY", SecurityType.Equity, Market.USA) ]
        #self.AddUniverseSelection( ManualUniverseSelectionModel(spy) )
        
        self.SetAlpha(ROELongShortAlphaModel())
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        self.SetRiskManagement(TrailingStopRiskManagementModel(0.05)) #parameter n1 step 0.01 from 0.00 to 0.10
        self.SetExecution(VolumeWeightedAveragePriceExecutionModel())
        #self.SetExecution(ImmediateExecutionModel())

class StockPickerUniverseSelectionModel(FundamentalUniverseSelectionModel):
    
    def __init__(self):
        super().__init__(True, None, None)
        self.lastMonth = -1
        self.dollar_volume_count = 1000 #parameter n2 step 1 start from 500 to 1500
        self.by_mom_count = 200 #parameter n3 step 1 start from 100 to 300  
        self.by_roe_count = 60 #parameter n4 step 1 start from 0 to 20
        self.mom = {}
        
    def SelectCoarse(self, algorithm, coarse):
        selected = []
        if self.lastMonth == algorithm.Time.month:
            return Universe.Unchanged
        self.lastMonth = algorithm.Time.month

        universe = sorted([x for x in coarse if x.HasFundamentalData and x.Volume > 0 and x.Price > 0],
            key=lambda x: x.DollarVolume, reverse=True)
        universe = [x for x in universe[:self.dollar_volume_count]]
        if len(universe) == 0:
            return Universe.Unchanged
            
        for coarse in universe:
            symbol = coarse.Symbol
            if symbol not in self.mom:
                history = algorithm.History(symbol, 20, Resolution.Daily)
                self.mom[symbol] = SelectionData(history) 
            self.mom[symbol].update(algorithm.Time, coarse.AdjustedPrice)
            
            if  self.mom[symbol].is_ready():
                selected.append({"symbol":symbol, "indicator": self.mom[symbol].mom})
                
        ordered = sorted(selected, key=lambda kv: kv["indicator"], reverse=True)
        orderedlist = []
        for i in range(len(ordered)):
            orderedlist.append(ordered[i]['symbol'])

        return orderedlist[:self.by_mom_count]

    def SelectFine(self, algorithm, fine):
        '''Performs fine selection by Return On Capital Employed
        The company's headquarter must in the U.S.
        The stock must be traded on either the NYSE or NASDAQ
        At least half a year since its initial public offering
        The stock's market cap must be greater than 500 million'''
        sortedByRoe = sorted([x for x in fine if x.CompanyReference.CountryId == "USA"
                                        and x.CompanyReference.PrimaryExchangeID in ["NYS","NAS"]
                                        and (algorithm.Time - x.SecurityReference.IPODate).days > 180
                                        and x.MarketCap > 5e8],
                               key = lambda x: x.OperationRatios.ROE.Value, reverse=True)
        
        count = len(sortedByRoe)
        if count == 0:
            return Universe.Unchanged
        
        # Update self.lastMonth after all QC500 criteria checks passed
        self.lastMonth = algorithm.Time.month
        
        universe = sortedByRoe[:self.by_roe_count]
        return [f.Symbol for f in universe]

class SelectionData():
    def __init__(self, history):
        self.mom = Momentum(14) #parameter n5 step 1 start from 2 to 26
        for bar in history.itertuples():
            self.mom.Update(bar.Index[1], bar.close)
            
    def is_ready(self):
        return self.mom.IsReady
    
    def update(self, time, price):
        self.mom.Update(time, price)

# Define the ROELongShortAlphaModel class  
class ROELongShortAlphaModel(AlphaModel):

    def __init__(self):
        self.lastMonth = -1
        self.removed = []
        self.added = []
    
    def OnSecuritiesChanged(self, algorithm, changes):
        for security in changes.AddedSecurities:
            self.added.append(security.Symbol)
        
        for security in changes.RemovedSecurities:
            self.removed.append(security.Symbol)

    def Update(self, algorithm, data):
        insights = []
        
        if self.lastMonth == algorithm.Time.month:
            return insights
        if algorithm.Time.month==8:
            return insights
        self.lastMonth = algorithm.Time.month
        
        #emit insights with insight directions based on momentum
        for symbol in self.added:
            insights.append(Insight.Price(symbol, timedelta(28), 1)) 
        for symbol in self.removed:
            insights.append(Insight.Price(symbol, timedelta(28), 0))
            
        self.removed.clear()
        self.added.clear()
        
        return insights