from UniverseSelection import UniverseSelectionModel
from Portfolio.EqualWeightingPortfolioConstructionModel import EqualWeightingPortfolioConstructionModel
from Alphas.ConstantAlphaModel import ConstantAlphaModel
from Execution.ImmediateExecutionModel import ImmediateExecutionModel
from Risk.NullRiskManagementModel import NullRiskManagementModel


class BasicTemplateFrameworkAlgorithm(QCAlgorithmFramework):

    def Initialize(self):

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

        self.SetStartDate(2018, 1, 1)   #Set Start Date
        self.SetEndDate(2018, 3, 1)    #Set End Date
        self.SetCash(100000)           #Set Strategy Cash

        count = 10
        self.SetUniverseSelection(UniverseSelectionModel())
        self.SetAlpha(ConstantAlphaModel(InsightType.Price, InsightDirection.Up, TimeSpan.FromMinutes(20), 0.025, None))
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        self.SetExecution(ImmediateExecutionModel())
        self.SetRiskManagement(NullRiskManagementModel())

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

class UniverseSelectionModel(FundamentalUniverseSelectionModel):

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

        super().__init__(False, universeSettings, securityInitializer)
        self.universeCount = universeCount
        self.averages = {}

    def SelectCoarse(self, algorithm, coarse):
        # We are going to use a dictionary to refer the object that will keep the moving averages
        for cf in coarse:
            if cf.Symbol not in self.averages:
                self.averages[cf.Symbol] = SymbolData(cf.Symbol)

            # Updates the SymbolData object with price and volume
            avg = self.averages[cf.Symbol]
            avg.update(cf.EndTime, cf.Price, cf.Volume)
        
        # Filter the values of the dict: we only want up-trending securities 
        values = list(filter(lambda x: x.is_uptrend, self.averages.values()))
        return [ x.symbol for x in values[:self.universeCount] ]

# class used to improve readability of the coarse selection function
class SymbolData:
    def __init__(self, symbol):
        self.symbol = symbol
        self.priceSMA = SimpleMovingAverage(10)
        self.volSMA = SimpleMovingAverage(10)
        self.priceWin = RollingWindow[float](10)
        self.is_uptrend = False

    def update(self, time, price, volume):
        self.priceWin.Add(price)
        self.priceSMA.Update(time, price)
        self.volSMA.Update(time, volume)
        if self.priceSMA.IsReady and self.volSMA.IsReady and self.priceWin.IsReady:
            MAprice = self.priceSMA.Current.Value
            MAvol = self.volSMA.Current.Value
            # current price > 10-days moving average price
            # current volume > 10-days moving average volume
            # current price > yesterday's close
            # current price > 10-days maximum price 
            self.is_uptrend =  price > MAprice and volume > MAvol and price > self.priceWin[1] and price > max(self.priceWin)