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

from QuantConnect.Data.UniverseSelection import *
from QuantConnect.Indicators import ExponentialMovingAverage
from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel

import decimal as d

class UniverseSelectionModel(FundamentalUniverseSelectionModel):

    def __init__(self,
                 universeCount = 20,
                 filterFineData = True,
                 universeSettings = None,
                 securityInitializer = None):

        super().__init__(filterFineData, universeSettings, securityInitializer)
        self.universeCount = universeCount
        self.resolution = Resolution.Daily
        
        # Holds our coarse fundamental indicators by symbol
        self.calculations = {}

    def SelectCoarse(self, algorithm, coarse):
        
        '''Performs coarse selection based on price and volume'''
        
        # We are going to use a dictionary to refer the object that will keep the calculations
        for cf in coarse:
            if cf.Symbol not in self.calculations:
                self.calculations[cf.Symbol] = SymbolData(algorithm, cf.Symbol, self.resolution)


            #self.calculations[cf.Symbol].RegisterIndicators(algorithm, self.resolution)
            self.calculations[cf.Symbol].update(cf.EndTime, cf.Price, cf.Volume, cf.DollarVolume, cf.HasFundamentalData)
            
            # Updates the SymbolData object with price, volume, dollar volume, etc
            #avg = self.calculations[cf.Symbol]
            #avg.RegisterIndicators(algorithm)
            #avg.update(cf.EndTime, cf.Price, cf.Volume, cf.DollarVolume, cf.HasFundamentalData)
        
        # Filter the values of the dict based on chosen criteria.
        filtered_coarse = list(filter(lambda x: (x.is_uptrend and x.fundamentaldata), self.calculations.values()))
        
        # Sort the values of the dict: we want those with greater DollarVolume
        sorted_coarse = sorted(filtered_coarse, key = lambda x: x.dollarvolume, reverse = True)
        
        #for x in sorted_values[:self.universeCount]:
        #    algorithm.Log(str(algorithm.Time) + 'symbol: ' + str(x.symbol.Value)
        #    + ' close price: ' + str(x.price)
        #    + ' volume: ' + str(x.volume)
        #    + ' dollar volume: ' + str(x.dollarvolume)
        #    + '  mean price: ' + str(x.priceSMA.Current.Value)
        #    + '  mean vol: ' + str(x.volSMA.Current.Value))
        
        # Select Top Stocks (based on universeCount)
        return [ x.Symbol for x in sorted_coarse[:self.universeCount] ]
        
    def SelectFine(self, algorithm, fine):
        
        '''Performs fine selection:
        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'''
        
        filtered_fine = list(filter(lambda x: (x.CompanyReference.CountryId == "USA"
        and (x.CompanyReference.PrimaryExchangeID == "NYS" or x.CompanyReference.PrimaryExchangeID == "NAS")
        and (algorithm.Time - x.SecurityReference.IPODate).days > 180
        and x.EarningReports.BasicEPS.TwelveMonths > 0
        and x.ValuationRatios.PERatio > 0
        and x.EarningReports.BasicAverageShares.ThreeMonths > 0), fine))
        #and x.EarningReports.BasicAverageShares.ThreeMonths * (x.EarningReports.BasicEPS.TwelveMonths * x.ValuationRatios.PERatio) > 5e8), fine))
        
        # Sort the values of the dict by market cap
        sorted_fine = sorted(filtered_fine, key = lambda x: x.EarningReports.BasicAverageShares.ThreeMonths * (x.EarningReports.BasicEPS.TwelveMonths * x.ValuationRatios.PERatio), reverse = True)
        
        # Select Top Stocks (up to UniverseCount)
        return [ x.Symbol for x in filtered_fine[:20] ]

# class used to improve readability of the coarse selection function
class SymbolData:
    
    def __init__(self, algorithm, symbol, resolution):
        
        self.Symbol = symbol
        #self.Consolidator = None
        
        self.priceWin = RollingWindow[float](10)
        
        self.priceSMA10 = SimpleMovingAverage(10)
        self.priceSMA50 = SimpleMovingAverage(50)
        self.priceSMA200 = SimpleMovingAverage(200)
        
        self.volSMA = SimpleMovingAverage(50)
        
        self.is_uptrend = False
        
        #self.Consolidator = algorithm.ResolveConsolidator(self.Symbol, resolution)
        #algorithm.RegisterIndicator(self.Symbol, self.priceSMA10, self.Consolidator)
        #smaName = algorithm.CreateIndicatorName(self.Symbol, "SimpleMovingAverage{}".format(10), resolution)
        #self.SMA = SimpleMovingAverage(smaName, 10)
        #algorithm.RegisterIndicator(self.Symbol, self.SMA, self.Consolidator)

    #def RegisterIndicators(self, algorithm, resolution):

        #self.Consolidator = algorithm.ResolveConsolidator(self.Symbol, resolution)
        #algorithm.RegisterIndicator(self.Symbol, self.priceSMA10, self.Consolidator)
        #algorithm.RegisterIndicator(self.symbol, self.priceSMA50, self.Consolidator)
        #algorithm.RegisterIndicator(self.symbol, self.priceSMA200, self.Consolidator)
        #algorithm.RegisterIndicator(self.symbol, self.volSMA, self.Consolidator)

    def update(self, time, price, volume, dollarvolume, fundamentaldata):
        
        self.price = price
        self.volume = volume
        self.dollarvolume = dollarvolume
        self.fundamentaldata = fundamentaldata
        
        self.priceWin.Add(price)
        
        self.priceSMA10.Update(time, price)
        self.priceSMA50.Update(time, price)
        self.priceSMA200.Update(time, price)
        
        self.volSMA.Update(time, volume)
        
        # self.priceSMA200.IsReady
        
        if self.priceSMA10.IsReady and self.priceSMA50.IsReady and self.volSMA.IsReady and self.priceWin.IsReady:
            
            SMA10price = self.priceSMA10.Current.Value
            SMA50price = self.priceSMA50.Current.Value
            SMA200price = self.priceSMA200.Current.Value
            SMAvol = self.volSMA.Current.Value
            
            # Here we can add the criteria for our universe
            # > SMA200price
            
            self.is_uptrend =  (SMA10price > SMA50price
            and price > 20 and price > SMA10price and price > max(self.priceWin)
            and price > (self.priceWin[1] + float(1))
            and SMAvol > 500 and volume > (d.Decimal(1.5) * SMAvol))

from UniverseSelection import UniverseSelectionModel
from AlphaCreation import ConstantAlphaModel
from Portfolio.EqualWeightingPortfolioConstructionModel import EqualWeightingPortfolioConstructionModel
from Execution.ImmediateExecutionModel import ImmediateExecutionModel
from Risk.NullRiskManagementModel import NullRiskManagementModel
from RiskManagement import PnLManagementModel

class BasicTemplateFrameworkAlgorithm(QCAlgorithmFramework):

    def Initialize(self):

        # Set requested data resolution
        self.UniverseSettings.Resolution = Resolution.Minute

        self.SetStartDate(2018, 5, 1)   # Set Start Date
        self.SetEndDate(2018, 10, 1)    # Set End Date
        self.SetCash(10000)            # Set Strategy Cash

        # Adding SPY Market to use in Schedule Function
        #self.AddEquity("SPY")
        # Scheduled Function to liquidate all positions 15 minutes before market close
        #self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.BeforeMarketClose("SPY", 15), self.rebalance)
        
        self.SetUniverseSelection(UniverseSelectionModel(universeCount = 5, filterFineData = False))
        self.SetAlpha(ConstantAlphaModel())
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        self.SetExecution(ImmediateExecutionModel())
        #self.SetRiskManagement(NullRiskManagementModel())
        self.SetRiskManagement(PnLManagementModel(maximumDrawdownPercent = 0.03, TakeProfitPercent = 0.1))
            
    #def rebalance(self):
        #algorithm.Log(str(algorithm.Time) + 'Liquidating End Of Day: ' + str(security))
        #self.Liquidate()

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

from QuantConnect import *
from QuantConnect.Algorithm import *
from QuantConnect.Algorithm.Framework import *
from QuantConnect.Algorithm.Framework.Portfolio import PortfolioTarget
from QuantConnect.Algorithm.Framework.Risk import RiskManagementModel

class PnLManagementModel(RiskManagementModel):
    '''Provides an implementation of IRiskManagementModel that limits the drawdown per holding to the specified percentage'''

    def __init__(self, maximumDrawdownPercent = 0.05, TakeProfitPercent = 0.05):
        '''Initializes a new instance of the MaximumDrawdownPercentPerSecurity class
        Args:
            maximumDrawdownPercent: The maximum percentage drawdown allowed for any single security holding
            TakeProfitPercent: The take profit percentage for any single security holding'''
        self.maximumDrawdownPercent = -abs(maximumDrawdownPercent)
        self.TakeProfitPercent = abs(TakeProfitPercent)

    def ManageRisk(self, algorithm, targets):
        '''Manages the algorithm's risk at each time step
        Args:
            algorithm: The algorithm instance'''
        targets = []
        for kvp in algorithm.Securities:
            security = kvp.Value

            if not security.Invested:
                continue

            pnl = security.Holdings.UnrealizedProfitPercent
            
            if pnl <= self.maximumDrawdownPercent:
                # Liquidate
                algorithm.Log(str(algorithm.Time) + 'Liquidating Due To MaxDrawdown: ' + str(security))
                targets.append(PortfolioTarget(security.Symbol, 0))
            elif pnl >= self.TakeProfitPercent:
                # Liquidate
                algorithm.Log(str(algorithm.Time) + 'Taking Profit: ' + str(security))
                targets.append(PortfolioTarget(security.Symbol, 0))

        return targets

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

from QuantConnect import *
from QuantConnect.Algorithm import *
from QuantConnect.Algorithm.Framework import *
from datetime import timedelta
from QuantConnect.Algorithm.Framework.Alphas import AlphaModel, Insight, InsightType, InsightDirection

class ConstantAlphaModel(AlphaModel):

    def __init__(self):
        self.period = timedelta(hours=5,minutes=30)
        self.securities = []
        self.insightsTimeBySymbol = {}
        self.ShouldEmitInsight = True
        self.Name = 'ConstantAlphaModel'


    def Update(self, algorithm, data):
        insights = []
        if self.ShouldEmitInsight:
            for security in self.securities:
                if security.Symbol.Value != "SPY":
                    insights.append(Insight.Price(security.Symbol, self.period, InsightDirection.Up))
                    self.ShouldEmitInsight = False
        return insights


    def OnSecuritiesChanged(self, algorithm, changes):
        self.removedSecurities = []
        for added in changes.AddedSecurities:
            self.securities.append(added)
            self.ShouldEmitInsight = True

        # this will allow the insight to be re-sent when the security re-joins the universe
        for removed in changes.RemovedSecurities:
            self.removedSecurities.append(removed)
            self.ShouldEmitInsight = True
            if removed in self.securities:
                self.securities.remove(removed)
            if removed.Symbol in self.insightsTimeBySymbol:
                self.insightsTimeBySymbol.pop(removed.Symbol)