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 SimpleMovingAverage
from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel

class CostumUniverseSelectionModel(FundamentalUniverseSelectionModel):

    def __init__(self, filterFineData = True, universeSettings = None, securityInitializer = True):
        super().__init__(filterFineData, universeSettings, securityInitializer)
        self.NumberOfSymbolsCoarse = 1000
        self.NumberOfSymbolsFine = 500
        self.dollarVolumeBySymbol = {}
        self.symbols = []
    
    def SelectCoarse(self, algorithm, coarse):
        coarse = list(coarse)
        
        if len(coarse) == 0:
            return self.symbols
        
        filtered = [x for x in coarse if x.HasFundamentalData and x.Volume > 100000]
        top = sorted(filtered, key = lambda x: x.DollarVolume, reverse = True)[:self.NumberOfSymbolsCoarse]
        self.dollarVolumeBySymbol = {i.Symbol: i.DollarVolume for i in top}
        self.symbols = [self.dollarVolumeBySymbol.keys()]
        
        return self.symbols
    
    def SelectFine(self, algorithm, fine):
        filteredFine = [x for x in fine if x.ValuationRatios.FCFRatio < 20
                                        and x.OperationRatios.ROIC.OneYear > 0
                                        and x.OperationRatios.ROE.OneYear > 0
                                        and x.OperationRatios.TotalDebtEquityRatio.OneYear < 1]
        sortedByROE = sorted(filteredFine, key = lambda x: x.OperationRatios.ROE.OneYear, reverse = True)[:self.NumberOfSymbolsFine]
        self.symbols = [f.Symbol for f in sortedByROE]
        
        return self.symbols
        
class SMAUniverseSelectionModel(FundamentalUniverseSelectionModel):
    
    def __init__(self, fastPeriod = 20, intermediatePeriod = 50, slowPeriod = 100, universeCount = 500, universeSettings = None, securityInitializer = None):
        super().__init__(False, universeSettings, securityInitializer)
        self.fastPeriod = fastPeriod
        self.intermediatePeriod = intermediatePeriod
        self.slowPeriod = slowPeriod
        self.universeCount = universeCount
        self.tolerance = 0.01
        self.averages = {}
    
    def SelectCoarse(self, algorithm, coarse):
        filtered = []
        
        for cf in coarse:
            if cf.Symbol not in self.averages:
                self.averages[cf.Symbol] = self.SelectionData(cf.Symbol, self.fastPeriod, self.intermediatePeriod, self.slowPeriod)
            
            avg = self.averages.get(cf.Symbol)
            
            if avg.Update(cf.EndTime, cf.AdjustedPrice) and self.Securities[cf].Price > avg.fast * (1 + self.tolerance) and avg.fast > avg.intermediate * (1 + self.tolerance) and avg.intermediate > avg.slow * (1 + self.tolerance):
                filtered.append(avg)
        
        filtered = sorted(filtered, key = lambda avg: avg.ScaledDelta, reverse = True)
        
        return [x.Symbol for x in filtered[:self.universeCount]]

class SelectionData:
    def __init__(self, symbol, fastPeriod, intermediatePeriod, slowPeriod):
        self.Symbol = symbol
        self.fastSma = SimpleMovingAverage(fastPeriod)
        self.intermediateSma = SimpleMovingAverage(intermediatePeriod)
        self.slowSma = SimpleMovingAverage(slowPeriod)
    
    @property
    def Fast(self):
        return float(self.fastSma.Current.Value)
    
    @property
    def intermediate(self):
        return float(self.intermediateSma.Current.Value)
    
    @property
    def slow(self):
        return float(self.slowSma.Current.Value)
    
    @property
    def ScaledDelta(self):
        return (self.fast - self.intermediate) / ((self.fast + self.intermediate) / 2)
    
    def Update(self, time, value):
        return self.fastSma.Update(time, value) & self.intermediateSma.Update(time, value) & self.slowSma.Update(time, value)