# QUANTCONNECT.COM - Democratizing Finance, Empowering Individuals.
# Lean Algorithmic Trading Engine v2.0. Copyright 2014 QuantConnect Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from clr import AddReference
AddReference("QuantConnect.Algorithm.Framework")
AddReference("QuantConnect.Indicators")
AddReference("QuantConnect.Common")

from QuantConnect import *
from QuantConnect.Indicators import *
from QuantConnect.Algorithm.Framework.Alphas import *
from datetime import timedelta
from enum import Enum

class PPOAlphaModel(AlphaModel):
    '''Uses PPO to create insights. 
    Using default settings, slow/fast is 12/26. Above/below -.5 will trigger a new insight.'''
    
    rebalance_date = None
    rebalance_complete = False

    def __init__(self,
                 fastPeriod = 12*5,
                 slowPeriod = 26*5,
                 trigger = 0.0,
                 consolidationPeriod = 7,
                 resolution = Resolution.Daily):
        self.lookback = (slowPeriod * 3)
        self.fastPeriod = fastPeriod
        self.slowPeriod = slowPeriod
        self.resolution = resolution
        self.consolidationPeriod = consolidationPeriod
        self.trigger = trigger
        self.predictionInterval = Extensions.ToTimeSpan(self.resolution)
        self.symbolDataBySymbol = {}

        resolutionString = Extensions.GetEnumString(resolution, Resolution)
        self.Name = '{}({},{},{})'.format(self.__class__.__name__, fastPeriod, slowPeriod, resolutionString)

    def Update(self, algorithm, data):
        '''Updates this alpha model with the latest data from the algorithm.
        This is called each time the algorithm receives data for subscribed securities
        Args:
            algorithm: The algorithm instance
            data: The new data available
        Returns:
            The new insights generated'''
            
        
        #algorithm.Debug("Update Called on PPOAlphaModel")
        # generate insights when scheduled
        insights = []
        #and algorithm.Time.hour == 8 and algorithm.Time.minute == 59
        if self.rebalance_complete or not (algorithm.Time.date() == self.rebalance_date):
            return insights
        
        # only get here to build new insights once on the scheduled trading date
        self.rebalance_complete = True
        #algorithm.Debug("Now Generating Insights in PPOAlpha")
        for symbol, symbolData in self.symbolDataBySymbol.items():
            
            if symbolData.CanEmit:
                direction = InsightDirection.Flat
                magnitude = symbolData.Return
                #algorithm.Debug("got a positive CanEmit with symbol and magnitude: " + str(symbol)  + str(magnitude))
                if magnitude > self.trigger: direction = InsightDirection.Up
                if magnitude < self.trigger: direction = InsightDirection.Flat

                insights.append(Insight.Price(symbol, self.predictionInterval, direction))
                algorithm.Log(str(symbol) + ": " + str(magnitude) + ": " + str(direction))
                
        return insights


    def OnSecuritiesChanged(self, algorithm, changes):
        '''Event fired each time the we add/remove securities from the data feed
        Args:
            algorithm: The algorithm instance that experienced the change in securities
            changes: The security additions and removals from the algorithm'''
            
        #algorithm.Debug("starting OnSecuritiesChanged")
        # clean up data for removed securities
        for removed in changes.RemovedSecurities:
            symbolData = self.symbolDataBySymbol.pop(removed.Symbol, None)
            if symbolData is not None:
                symbolData.RemoveConsolidators(algorithm)
         
        #algorithm.Debug("in OnSecuritiesChanged past removed list")       
        # initialize data for added securities
        symbols = [ x.Symbol for x in changes.AddedSecurities ]
        #algorithm.Debug("before history retrieval")
        history = algorithm.History(symbols, self.lookback, self.resolution)
        #algorithm.Debug("after history retrieval")
        if history.empty: return

        tickers = history.index.levels[0]
        for ticker in tickers:
            symbol = SymbolCache.GetSymbol(ticker)

            if symbol not in self.symbolDataBySymbol:
                symbolData = SymbolData(symbol, self.fastPeriod, self.slowPeriod, self.consolidationPeriod)
                self.symbolDataBySymbol[symbol] = symbolData
                symbolData.RegisterIndicators(algorithm, self.resolution)
                symbolData.WarmUpIndicators(history.loc[ticker])
        

class SymbolData:
    '''Contains data specific to a symbol required by this model'''

    
    def __init__(self, symbol, fastPeriod, slowPeriod, consolidationPeriod):
        self.Symbol = symbol
        self.Fast = fastPeriod
        self.Slow = slowPeriod
        self.PPO = PercentagePriceOscillator(fastPeriod, slowPeriod, MovingAverageType.Exponential)
        #self.Consolidator = TradeBarConsolidator(TimeSpan.FromDays(consolidationPeriod))
        self.Consolidator = None
        self.Previous = None

        # True if the fast is above the slow, otherwise false.
        # This is used to prevent emitting the same signal repeatedly
        #self.FastIsOverSlow = False
        #self.abovePPOTrigger = False
        
    #def setPPOTrigger (self, value):
        #self.abovePPOTrigger = value

    @property
    def belowPPOTrigger(self):
        return (not self.abovePPOTrigger)
        
    @property
    def CanEmit(self):
        if self.Previous == self.PPO.Samples:
            return False

        self.Previous = self.PPO.Samples
        return self.PPO.IsReady

    @property
    def Return(self):
        return float(self.PPO.Current.Value)
        
    #def OnWeeklyData(self, sender, bar):
        #self.PPO.Update(bar.EndTime, bar.Close)
        
    def RegisterIndicators(self, algorithm, resolution):
        #self.Consolidator = algorithm.ResolveConsolidator(self.Symbol, resolution)
        algorithm.RegisterIndicator(self.Symbol, self.PPO, self.Consolidator)

    def RemoveConsolidators(self, algorithm):
        if self.Consolidator is not None:
            algorithm.SubscriptionManager.RemoveConsolidator(self.Symbol, self.Consolidator)

    def WarmUpIndicators(self, history):
        for tuple in history.itertuples():
            self.PPO.Update(tuple.Index, tuple.close)
            #self.Consolidator.Update(tuple)

class StockSelectionStrategyBasedOnFundamentalFactorsAlgorithm(QCAlgorithm):

    def Initialize(self):

        self.SetStartDate(2021, 11, 1)  # Set Start Date
        #self.SetEndDate(2017, 5, 2)    # Set End Date
        self.SetCash(100000)          # Set Strategy Cash

        self.current_month = -1
        self.coarse_count = 300
        self.fine_count = 10

        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)

        self.SetAlpha(ConstantAlphaModel(InsightType.Price, InsightDirection.Up, timedelta(30)))

        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel(lambda time:None))
        
        #self.AddEquity("SPY", Resolution.Daily)
        
        self._list = [ ]
        

    def CoarseSelectionFunction(self, coarse):

        if self.current_month == self.Time.month:
            return Universe.Unchanged

        self.current_month = self.Time.month

        selected = [x for x in coarse if (x.HasFundamentalData) and x.DollarVolume > 1000000]

        return [x.Symbol for x in selected]

    def FineSelectionFunction(self, fine):

        
        fine = [x for x in fine if x.FinancialStatements.CashFlowStatement.FreeCashFlow.ThreeMonths > 0
                        #and ((x.OperationRatios.RevenueGrowth.OneYear) - (x.FinancialStatements.IncomeStatement.TotalRevenue.OneMonth)) / (x.FinancialStatements.IncomeStatement.TotalRevenue.OneMonth) > 0.25
                        and x.OperationRatios.NormalizedNetProfitMargin.OneYear > 0.07
                        and x.FinancialStatements.IncomeStatement.TotalRevenue.TwelveMonths < 900000000
                        and x.FinancialStatements.CashFlowStatement.OperatingCashFlow.TwelveMonths > 0
                        and x.OperationRatios.ROE.OneYear > 0.15
                        and x.SecurityReference.IsPrimaryShare
                        and x.SecurityReference.SecurityType == "ST00000001"
                        and x.SecurityReference.IsDepositaryReceipt == 0
                        and x.CompanyReference.IsLimitedPartnership == 0]
                        #and x.AssetClassification.MorningstarSectorCode == MorningstarSectorCode.Technology]


        return [x.Symbol for x in fine]