import pandas as pd
import math

class Coin(QCAlgorithm):
    
    def __init__(self):
         
        # Global variables
        self.initialized = False
        
        # symbol data
        self.ovd = {}                   # overvalued data
        self.uvd = {}                   # undervalued data
        self.ovs = []                   # holds overvalued symbols
        self.uvs = []                   # holds undervalued symbols
        
        # Coarse Filter
        self.top = 450  # Only take the top equities from coarse raw
        
        # Fine filter
        self.hitop = 25     # Only take top hitop equities from overvalued
        self.lotop = 25     # Only take top lotop equities from undervalued
        self.himax = 0.99   # Upper limit for overvalued percentile
        self.himin = 0.70   # lower limit for overvalued percentile
        self.lomax = 0.30   # upper limit for undervalued percentile
        self.lomin = 0.00   # lower limit for undervalued percentile
        self.ratios = ["per", "pbr", "psr", "pgr", "tdr", "evr"] # filter through these ratios
        
    def Initialize(self):
        
        self.SetStartDate(2003, 1, 1)
        self.SetEndDate(datetime.now().date() - timedelta(1))
        self.SetCash(20000)
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage)
        self.UniverseSettings.Resolution = Resolution.Minute
        self.AddUniverse(self.coarseFilter, self.fineFilter)
        
    def initData(self):
        aes = self.ovs + self.uvs
        for symbol in aes:
             if (symbol in self.ovs):
                 self.ovd[symbol] = UniverseSymbolData(
                     symbol, self.SubscriptionManager, self.RegisterIndicator)
             else:
                 self.uvd[symbol] = UniverseSymbolData(
                     symbol, self.SubscriptionManager, self.RegisterIndicator)
        self.SetWarmup(timedelta(1))
            
    def coarseFilter(self, coarse):
        raw = {}                                                    # raw data holding all stocks
        for eq in coarse:                                           # iterate through equitties
            if (eq.HasFundamentalData):                             # check for fundamental data
                raw[eq] = CoarseSymbolData(eq.Symbol, self.Debug)   # raw: key=eq, value=CoarseSymbolData
                candidate = raw[eq]                                 # candidate = get eq in raw
                candidate.update(eq.Volume, eq.DollarVolume,        # give candidate eq data
                        eq.AdjustedPrice, eq.EndTime, self.History)
        selected = list(filter(lambda sd: (sd.selected == True), raw.values())) # get only desired eq
        selected.sort(key=lambda sd: sd.dvolume, reverse=True)      # sort descending
        if (self.top > len(selected)):
            self.top = len(selected)
        return [sd.symbol for sd in selected[:self.top]] # return symbols

    def fineFilter(self, fine):
        raw = {}
        for eq in fine:                             # raw: key=equity, value=FineFilterData(equity)
            raw[eq] = FineSymbolData(eq, eq.Symbol)
        for ratio in self.ratios:                   # iterate through ratios
            raw = self.getPercentile(ratio, raw)    # filter using ratio in list of ratios
        selected = list(raw.values())
        selected.sort(key=lambda sd: sd.PERatio)    # sort by PE ratio
        if ((self.hitop + self.lotop) > len(selected)):
            self.hitop = math.ceil(len(selected) / 2)
            self.lotop = math.floor(len(selected) / 2)
        self.ovs = selected[len(selected) - self.hitop:]
        self.uvs = selected[:self.lotop]
        self.uvs = [sd.symbol for sd in self.uvs]
        self.ovs = [sd.symbol for sd in self.ovs]
        selected = selected[:self.lotop] + selected[len(selected) - self.hitop:] # take top 25 overvalued and undervalued
        #self.avs
        return [sd.symbol for sd in selected]

    def getPercentile(self, ratio, pool):
        eqr = {}            # holds symbols mapped to given ratio (in decimal)
        rts = {"ratio":[]}  # list of ratios (in decimal) of pool
        for eq in pool:
            eqr[eq] = pool[eq].getRatio(ratio)
            rts["ratio"].append(pool[eq].getRatio(ratio))
        daf = pd.DataFrame(rts)                             # dataframe to calculate percentiles below
        hmin = daf.quantile(self.himin)[0]
        hmax = daf.quantile(self.himax)[0]
        lmin = daf.quantile(self.lomin)[0]
        lmax = daf.quantile(self.lomax)[0]
        for eq in eqr:
            if ((eqr[eq] < hmin) | (eqr[eq] > hmax)) & (    # if ratio does not fall in overvalue range
                    (eqr[eq] < lmin) | (eqr[eq] > lmax)):   # and ratio does not fall in undervalue range
                pool.pop(eq)                                # remove it from pool
        return pool
 
    def OnData(self, data):
        if (self.initialized != True):
            self.initialized = True
            self.initData()
        
        
        
        CloseTimeString = "12:00:00"  # Set up pre-market close time
        CloseOutTime = datetime.strptime(CloseTimeString, "%H:%M:%S")
        CloseTime = datetime.time(CloseOutTime)

        
        OpenTimeString = "10:30:00" # Set up post-market open time
        OpenTradesTime = datetime.strptime(OpenTimeString, "%H:%M:%S")
        OpenTime = datetime.time(OpenTradesTime)
    
        #long comparisons
        CloseOversma = 1.02
        vwmaOverema = 1.04
            
        #short comparisons
        CloseUndersma = 0.98
        vwmaUnderema = 0.96
        
            
        # undervalued stocks
        for symbol in self.uvd:                         #loop through undervalued
            symboldata = self.uvd[symbol]
            if not self.Portfolio[symbol].Invested: #if we do not own this symbol
                
   
                # Get symbol trade time
                DataSymbolTime = data[symbol].Time
                DataTime = datetime.time(DataSymbolTime)

                if DataTime < CloseTime:    # Check to see if we are near market close
                    
                        self.Debug(str(symbol) + " is long ready on " + str(DataSymbolTime))
                        posSize = self.CalculateOrderQuantity(symbol, 0.4) # Set Position Size
                        self.MarketOrder(symbol, posSize) # Market Order to buy
                    
                    
        # overvalued stocks
        for symbol in self.ovd:                         #loop through undervalued
            symboldata = self.ovd[symbol]
            if not self.Portfolio[symbol].Invested: #if we do not own this symbol
                
   
                # Get symbol trade time
                DataSymbolTime = data[symbol].Time 
                DataTime = datetime.time(DataSymbolTime)

                if DataTime < CloseTime:  # Check to see if we are near market close
                    
                        self.Debug(str(symbol) + " is short ready on " + str(DataSymbolTime))
                        posSize = self.CalculateOrderQuantity(symbol, -0.4) # Set Position Size
                        self.MarketOrder(symbol, posSize) # Market Order to buy
                        
    def OnEndOfDay(self): #triggered before end of day 
        
        #Liquidate all symbol
        self.Liquidate()
        self.Debug("Liquidating at end of day...")
        
class UniverseSymbolData(object):
    
    def __init__(self, symbol, manager, register):
        self.register = register
        self.manager = manager
        self.symbol = symbol
        self.PRD = 10       # period is 10 minutes

        # Periods
        self.SMAP = 30      # Simple Moving Average periods
        self.RSIP = 30      # Relative Strength Index periods
        self.ATRP = 30      # Average True Range periods
        self.EMAP = 5       # Exponential moving Average periods
        self.CLSP = 1      # "close value" simple moving average
        self.WMAP = 5       # Volume Weighted Moving Average periods (custom)
        self.PVPD = 144     # 1 day Pivot Point                     (custom)

        # indicators
        self.sma = SimpleMovingAverage(self.SMAP)
        self.rsi = RelativeStrengthIndex(self.RSIP)
        self.atr = AverageTrueRange(self.ATRP)
        self.ema = ExponentialMovingAverage(self.EMAP)
        self.cls = SimpleMovingAverage(self.CLSP)
        self.wma = WeightedMovingAverage(self.WMAP) # custom
        #self.pvp = PivotPoint(self.PVPD)            # custom

        # create consolidator and subscribe
        PRDConsolidator = TradeBarConsolidator(timedelta(minutes = self.PRD))
        self.manager.AddConsolidator(self.symbol, PRDConsolidator)

        # register indicators
        self.register(self.symbol, self.sma, PRDConsolidator)
        self.register(self.symbol, self.rsi, PRDConsolidator)
        self.register(self.symbol, self.atr, PRDConsolidator)
        self.register(self.symbol, self.ema, PRDConsolidator)
        self.register(self.symbol, self.cls, PRDConsolidator)
        self.register(self.symbol, self.wma, PRDConsolidator) #custom

class CoarseSymbolData(object): # this class for coarse filter equities (raw) 

    def __init__(self, symbol, debug):
        self.p = debug
        self.symbol = symbol
        self.MINDvl = 5000000                           # Minimum dollar volume
        self.SMVMul = 1                               # SMV 30 day multiplier
        self.SMAMin = 1                               # SMA min price 
        self.SMAVDS = 30                                # simple moving average volume days
        self.SMAVD = 1                                  # simple moving average volume day
        self.SMAD = 1                                   # simple moving average day
        self.SMAVA = SimpleMovingAverage(self.SMAVDS)   # SMAV alternate indicator
        self.SMAV = SimpleMovingAverage(self.SMAVD)     # SMAV indicator
        self.SMA = SimpleMovingAverage(self.SMAD)       # SMA indicator
        self.selected = False
        self.dvolume = None
        self.volume = None
        self.price = None
        self.time = None

    def update(self, volume, dvolume, price, time, history):
        self.dvolume = dvolume
        self.volume = volume
        self.price = price
        self.time = time
        hst = history(self.symbol, self.SMAVDS, Resolution.Daily)
        for t, r in hst.iterrows():
            self.SMAVA.Update(t[1], r["volume"])
        self.checkReqs()

    def checkReqs(self):
        if ((self.SMA.Update(self.time, self.price)) &          # check if indicators ready
                (self.SMAV.Update(self.time, self.volume)) &
                (self.SMAVA.IsReady)):
            smava = self.SMAVA.Current.Value
            smav = self.SMAV.Current.Value
            sma = self.SMA.Current.Value
            if ((sma >= self.SMAMin) &                          # sma has to be bigger than 10
                    (smav >= (self.SMVMul * smava)) &           # 1 day smav >= 1.5 * 30 day smav
                    (self.dvolume > self.MINDvl)):              # min 1 mill dollar volume
                self.selected = True
                
class FineSymbolData(object):
    
    def __init__(self, eq, symbol):
        self.eq = eq
        self.symbol = symbol
        self.PERatio = self.eq.ValuationRatios.PERatio
        self.PBRatio = self.eq.ValuationRatios.PBRatio
        self.PSRatio = self.eq.ValuationRatios.PSRatio
        self.PGRatio = self.eq.ValuationRatios.PEGRatio
        self.TDRatio = self.eq.OperationRatios.TotalDebtEquityRatio.OneMonth
        self.EVRatio = self.eq.ValuationRatios.EVToEBITDA1YearGrowth
        
    def getRatio(self, ratio):
        if (ratio == "per"):
            return self.PERatio
        elif (ratio == "pbr"):
            return self.PBRatio
        elif (ratio == "psr"):
            return self.PSRatio
        elif (ratio == "pgr"):
            return self.PGRatio
        elif (ratio == "tdr"):
            return self.TDRatio
        elif (ratio == "evr"):
            return self.EVRatio
            
class WeightedMovingAverage:

    def __init__(self, period):
        self.lastdate = datetime.min
        self.IsReady = False
        self.period = period
        self.pcount = 0
        self.ptvsum = 0.0       # sum of price * volume
        self.Value = 0.0        # indicator value
        self.vsum = 0.0         # sum of volume

    def IsReady(self):
        if (self.pcount == self.period):
            self.Value = (self.ptvsum / vsum)
            self.pcount = 0
            self.ptvsum = 0.0
            self.vsum = 0.0
            return True
        return False

    def Update(self, input): # input is a TradeBar object
        success, volume, close = self.getData(input)
        if not success:
            return
        self.pcount += 1
        self.ptvsum += close * volume
        self.vsum += volume
        
    def getData(self, input):
        if not input.IsFillForward:
            return True, float(input.Volume), float(input.Close)
        return False, 0.0, 0.0


#---------------------------------------------- END ALGORITHM ------------------------------------#