#region imports
from AlgorithmImports import *
import pandas as pd

#endregion
class ParticleTransdimensionalAutosequencers(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2023, 1, 15)
        self.SetEndDate(2023, 1, 20)
        self.SetCash(50000)
        self.cash = 50000
        self.daily_stop = 0.02
        # List of equities that would cause trouble when running the algo and needed to be manually added at initialization
        self.equityls = ['SPY','HCM','MWGP','EXPD','PLAY','PRGS','ADPT','CDW','SIX','TWOU','DHR']
        for symbol in self.equityls:
            self.AddEquity(symbol, Resolution.Minute, extendedMarketHours=True)
        self.AddUniverseSelection(CoarseFundamentalUniverseSelectionModel(self.CoarseSelectionFunction))
        self.UniverseSettings.Resolution = Resolution.Minute
        self.UniverseSettings.extendedMarketHours = True

        # Track open price of the day
        self.Schedule.On(
            self.DateRules.EveryDay("SPY"),
            self.TimeRules.At(9, 31),
            self.TrackOpen
        )
        # Select stocks with highest RVOL in the first 15 mins of trading
        self.Schedule.On(
            self.DateRules.EveryDay("SPY"),
            self.TimeRules.At(9, 35),
            self.SelectUniverse
        )

        # Liquidate all position before market close
        self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.BeforeMarketClose("SPY", 1), self.ClosePositions)

        # List or dictionary to keep track of values for the universe selection
        self.universe = []
        self.volume_by_symbol = {}
        self.open_by_symbol = {}
        self.gap = {}
        self.vol_after_gap = {}
        self.rvol_by_symbol = {}
        self.logged = False
        # Essential to have this automatic warmup, otherwise indicators are all zero and indicator.is_ready = False
        self.EnableAutomaticIndicatorWarmUp = True
        
        # List of dictionary to keep track of indicators that determines entry and exit criteria
        self.green_vol = {}
        self.red_vol = {}
        self.green_red_df = {}
        self.vwap_dict = {}
        self.vwap_df = {}
        self.vwap_cross_flag = {}
        self.vwap_above_flag = {}
        self.high_vol_flag = {}
        self.green_red_flag = {}
        self.can_trade_flag = False
        
    # Coarse Selection Filter executed at 00:00 on every day (once only) based on data of the previous trading day
    def CoarseSelectionFunction(self, coarse):
        self.volume_by_symbol = {c.Symbol: 0 for c in coarse if c.Price > 30 and c.Price < 400 and c.Volume > 2000000 and c.HasFundamentalData}
        self.Debug(f"Universe size before volume filter: {len(self.volume_by_symbol)} at {self.Time}")
        self.green_vol = self.volume_by_symbol.copy()
        self.red_vol = self.volume_by_symbol.copy()
        # initialize df and various flags
        self.can_trade_flag = False
        for symbol in self.volume_by_symbol.keys():
            self.green_red_df[symbol] = pd.DataFrame([],columns = ['volume', 'green_red'])
            self.vwap_df[symbol] = pd.DataFrame([],columns = ['vwap', 'close','high'])
            self.vwap_cross_flag[symbol] = False
            self.green_red_flag[symbol] = 0
            self.high_vol_flag[symbol] = False
            self.vwap_above_flag[symbol] = 0.0
        #for symbol, volume in self.volume_by_symbol.items():
        #    try:
        #        self.vwap_dict[symbol] = SymbolData(symbol, self)
        #    except:
        #        self.Debug(f'{symbol} cannot log vwap data')
        return list(self.volume_by_symbol.keys())

    # OnData executed whenever there is new data arriving. In this case, new data arrive at 9:30 every day at minute interval (due to Resolution.Minute)
    def OnData(self, data):
        
        # register VWAP indicator at 9:30
        if self.Time.hour == 9 and self.Time.minute == 30:
            for symbol in self.volume_by_symbol.keys():
                self.vwap_dict[symbol] = SymbolData(symbol, self)
        
        # Every minute, data got pumped in for many securities. We record the minute volume data for stocks that passes our coarse selection filter
        for symbol in self.volume_by_symbol.keys():
            if data.ContainsKey(symbol):
                try: 
                    #self.Log(f'{symbol} has minute volume of {data[symbol].Volume}')
                    self.volume_by_symbol[symbol] += data[symbol].Volume
                    
                    # if the bar is green, record the volume in a seperate dictionary
                    if data[symbol].Close > data[symbol].Open:
                        self.green_vol[symbol] += data[symbol].Volume
                        self.green_red_df[symbol].loc[self.Time] = [data[symbol].Volume, "Green"]
                    if data[symbol].Close < data[symbol].Open:
                        self.red_vol[symbol] += data[symbol].Volume
                        self.green_red_df[symbol].loc[self.Time] = [data[symbol].Volume, "Red"]

                    self.vwap_dict[symbol].Update(data[symbol].EndTime, data[symbol].Close)
                except:
                    self.Log(f'{symbol} has data[symbol] as non-type')
        
        # When volume_by_symbol is cleared (it is cleared at 9:35 by a scheduled event), it will print the length of the final universe to be traded
        if len(self.volume_by_symbol) == 0:
            if not self.logged:
                self.logged = True
                self.Debug(f"Universe size after volume filter: {len(self.universe)} at {self.Time}")
            #return 
        
        # from 9:35 onwards, only keep track of indicators for selected universe
        # Track green vs red volume every min
        for symbol in self.universe:
            if data.ContainsKey(symbol):
                try:
                    # Track green and red bar volume
                    if data[symbol].Close > data[symbol].Open:
                        #self.green_vol[symbol] += data[symbol].Volume
                        self.green_red_df[symbol].loc[self.Time] = [data[symbol].Volume, "Green"]
                    if data[symbol].Close < data[symbol].Open:
                        #self.red_vol[symbol] += data[symbol].Volume
                        self.green_red_df[symbol].loc[self.Time] = [data[symbol].Volume, "Red"]
                except:
                    self.Log(f'{symbol} has data[symbol] as non-type')


        # Track green vs red volume in last 10 bars
        for symbol in self.universe:
            if data.ContainsKey(symbol):
                data_symbol = self.green_red_df[symbol] 
                length_of_data = data_symbol.shape[0]
                if length_of_data <= 10:
                    last10barvol = data_symbol[-length_of_data:]
                else:
                    last10barvol = data_symbol[-10:]
                green_vol_last10 = last10barvol.loc[last10barvol['green_red'] == "Green", "volume"].sum()
                red_vol_last10 = last10barvol.loc[last10barvol['green_red'] == "Red", "volume"].sum()
                try:
                    bullishness_last10 = green_vol_last10 / (green_vol_last10 + red_vol_last10)
                except:
                    self.Debug(f'{symbol} has green vol of {green_vol_last10} and red vol of {red_vol_last10}')
                    bullishness_last10 = 0.5
                if bullishness_last10 >= 0.75:
                    self.green_red_flag[symbol] = 1
                elif bullishness_last10 <= 0.25:
                    self.green_red_flag[symbol] = -1
                else:
                    self.green_red_flag[symbol] = 0

        # Track VWAP and whether price crosses VWAP
        for symbol in self.universe:
            self.vwap_cross_flag[symbol] = False # reset flag to 0 every bar
            self.vwap_above_flag[symbol] = 0.0 # reset flag to 0 every bar
            if data.ContainsKey(symbol):
                # Track vwap
                #self.vwap_dict[symbol].Update(data[symbol].EndTime, data[symbol].Close)
                try:
                    vwap = self.vwap_dict[symbol].vwap.Current.Value
                    # vwap cross flag
                    if (data[symbol].High > vwap and data[symbol].Low < vwap) and (data[symbol].Close > data[symbol].Open):
                        self.vwap_cross_flag[symbol] = 1
                    elif (data[symbol].High > vwap and data[symbol].Low < vwap) and (data[symbol].Close < data[symbol].Open):
                        self.vwap_cross_flag[symbol] = -1
                    else:
                        self.vwap_cross_flag[symbol] = False
                    # vwap df tracking
                    self.vwap_df[symbol].loc[self.Time] = [vwap, data[symbol].Close, data[symbol].High]
                    no_of_bars = self.vwap_df[symbol].shape[0]
                    no_of_bars_above_vwap = len(self.vwap_df[symbol][self.vwap_df[symbol]['close'] > self.vwap_df[symbol]['vwap']])
                    perc_of_bars_above_vwap = no_of_bars_above_vwap /  no_of_bars
                    self.vwap_above_flag[symbol] = perc_of_bars_above_vwap
                except:
                    self.vwap_cross_flag[symbol] = False
                    self.vwap_above_flag[symbol] = 0.0
            
            #if self.vwap_cross_flag[symbol] == 1: #and self.Time.hour >= 13 and self.Time.minute >= 25:
            #    self.Debug(f"{symbol} has vwap of {self.vwap_dict[symbol].vwap.Current.Value} vs open price of {data[symbol].Open} and close price of {data[symbol].Close} and vwap flag of {self.vwap_cross_flag[symbol]} at {self.Time}")
            #    self.Debug(f"{symbol} has cross vwap at {self.Time}")

        # Track volume of current bar vs the past 20 bars
        for symbol in self.universe:
            self.high_vol_flag[symbol] = False
            try:
                if data.ContainsKey(symbol):
                    vol_last_20_bar = self.SMA(symbol, 20, Resolution.Minute, Field.Volume)
                    if data[symbol].Volume > 2.0 * vol_last_20_bar.Current.Value:
                        self.high_vol_flag[symbol] = True
            except:
                self.high_vol_flag[symbol] = False
            #if self.high_vol_flag[symbol] == True: #and self.Time.hour >= 13 and self.Time.minute >= 25:
            #    self.Debug(f"{symbol} has volume of {data[symbol].Volume} vs avg of last 10 bar of {vol_last_10_bar.Current.Value} and high_vol_flag of {self.high_vol_flag[symbol]} at {self.Time}")            

        # Buying criteria
        for symbol in self.universe:
            if data.ContainsKey(symbol):
                invested = [x.Symbol.Value for x in self.Portfolio.Values if x.Invested]
                if (self.Time.hour == 15 and self.Time.minute >=50) or (self.Time.hour == 16):
                    self.can_trade_flag == False
                if self.high_vol_flag[symbol] == True and self.vwap_cross_flag[symbol] == 1 and self.green_red_flag[symbol] == 1 and symbol not in invested and self.can_trade_flag == True and self.vwap_above_flag[symbol] >= 0.5:
                    risk_quantum = self.daily_stop * self.cash 
                    stop_loss = self.vwap_dict[symbol].vwap.Current.Value - self.ATR(symbol, 20, Resolution.Daily).Current.Value*2
                    risk = self.ATR(symbol, 20, Resolution.Daily).Current.Value*2
                    take_profit = self.vwap_dict[symbol].vwap.Current.Value + risk * 3
                    quantity = 1000
                    self.MarketOrder(symbol, quantity)
                    self.StopMarketOrder(symbol, -quantity, stop_loss)
                    self.LimitOrder(symbol, -quantity, take_profit)
                    #self.SetHoldings(symbol, 0.33)
                    self.Debug(f"Buy {symbol} at {self.Time} with market entry, stop loss at {stop_loss}, risk at {risk} atr at {self.ATR(symbol, 20, Resolution.Daily).Current.Value} and take profit at {take_profit} because its vol flag is {self.high_vol_flag[symbol]}, vwap_cross_flag is {self.vwap_cross_flag[symbol]}, vwap_above_flag is {self.vwap_above_flag[symbol]} and green_red flag is {self.green_red_flag[symbol]}")

        """
        for symbol in self.universe:
            if self.vwap_cross_flag[symbol] == 1 and self.high_vol_flag[symbol] == True:
                self.Debug(f"{symbol} meet both vwap cross and high vol criteria at {self.Time}.")

        for symbol in self.universe:
            if self.green_red_flag[symbol] == 1:
                self.Debug(f"{symbol} has more green volume than red in the past 10 bars at {self.Time}.")
            if self.green_red_flag[symbol] == -1:
                self.Debug(f"{symbol} has more red volume than green in the past 10 bars at {self.Time}.")         
        """

        # Check if variables are tracked throughout the session by looking at various dict at 15:49 every day
        #if self.Time.hour == 9 and self.Time.minute > 34 and self.Time.minute <= 45:
        if self.Time.hour == 15 and self.Time.minute == 59:
            self.Debug(f"it is {self.Time} now.")
            for symbol in self.universe:
                #bullishness = self.green_vol[symbol] / (self.green_vol[symbol] + self.red_vol[symbol])
                data = self.green_red_df[symbol]
                df_green_vol = data.loc[data['green_red'] == "Green", "volume"].sum()
                df_red_vol = data.loc[data['green_red'] == "Red", "volume"].sum()
                self.Debug(f"According to dataframe, {symbol} has green volume of {df_green_vol} and red volume of {df_red_vol} at {self.Time}")
                #self.Debug(f"{symbol} has vwap of {self.vwap_dict[symbol].vwap.Current.Value} vs open price of {data[symbol].Open} and close price of {data[symbol].Close} at {self.Time}") 

    # TrackOpen is executed everyday at 9:31 and track the opening price of the current trading day. Slight descrepency with actual opening price but close enough
    def TrackOpen(self):
        for symbol, volume in self.volume_by_symbol.items():
            historyDataMin = self.History(symbol,1,Resolution.Minute)
            try:
                open_price_sym = historyDataMin['open'][-1]
            except:
                self.Debug(f"Opening price data for current day unavailable for {symbol.Value}")
            self.open_by_symbol[symbol] = open_price_sym
            #self.Debug(f'{symbol} has open price of {open_price_sym} at {self.Time}')
        # Calculate gap%
        for symbol, volume in self.volume_by_symbol.items():
            historyData = self.History(symbol,2,Resolution.Daily)
            try:
                #openDayAfterEarnings = historyData['open'][0]
                openDayAfterEarnings = self.open_by_symbol[symbol]
                closeDayBeforeEarnings = historyData['close'][-1]
            except:
                self.Debug(f"History data unavailable for {symbol.Value}")
                continue
            priceGap = openDayAfterEarnings - closeDayBeforeEarnings
            percentGap = priceGap / closeDayBeforeEarnings
            if (percentGap > 0.03):
                self.gap[symbol] = percentGap
                self.Debug(f'{symbol} closes at {closeDayBeforeEarnings} previously and opens at {openDayAfterEarnings} now, it has gap percentage of {percentGap} at {self.Time}')
                self.vol_after_gap[symbol] = volume
                #self.vwap_dict[symbol] = SymbolData(symbol, self)
        self.Debug(f'Universe size after gap filter: {len(self.gap)}')
     
    # SelectUniverse is executed at 9:35 everyday to calculate gap %, filter by +-5%, then calculate rvol, sort the narrowed list after gap filter by rvol, then output the top 10    
    def SelectUniverse(self):
        self.universe = []
        # Calculate rvol and sort that. Get top 10 stocks
        for symbol, gap in self.gap.items():
            volume = self.vol_after_gap[symbol]
            symbol_sma = self.SMA(symbol, 20, Resolution.Daily, Field.Volume)
            symbol_rvol = self.volume_by_symbol[symbol] / self.SMA(symbol, 20, Resolution.Daily, Field.Volume).Current.Value
            self.rvol_by_symbol[symbol] = symbol_rvol        
        sorted_rvol_ls = sorted(self.rvol_by_symbol.items(), key=lambda x:x[1], reverse=True)
        temp_ls = []
        for tuple in sorted_rvol_ls:
            if tuple[1] > 0:
                temp_ls.append(tuple[0])
        self.universe = temp_ls[:10]
        for symbol in self.universe:
            self.Debug(f'{symbol} has avg daily ATR of {self.ATR(symbol, 20, Resolution.Daily).Current.Value}, avg daily volume of {self.SMA(symbol, 20, Resolution.Daily, Field.Volume)} and intraday volume of {self.volume_by_symbol[symbol]} and rvol of {self.rvol_by_symbol[symbol]} and gap of {self.gap[symbol]} and green vol of {self.green_vol[symbol]} and red vol of {self.red_vol[symbol]} at {self.Time}')
            #self.vwap_dict[symbol] = SymbolData(symbol, self)
            
        # Clear all dictionary to start fresh the next day
        self.volume_by_symbol.clear()
        self.rvol_by_symbol.clear()
        self.open_by_symbol.clear()
        self.gap.clear()
        self.vol_after_gap.clear()
        self.logged = False
        self.can_trade_flag = True

    def ClosePositions(self):        
        if self.Portfolio.Invested:
            self.Liquidate()

class SymbolData:
    def __init__(self,symbol,algo):
        self.algo = algo
        self.symbol = symbol
        self.vwap = algo.VWAP(self.symbol)

    def Update(self,time,close):
        self.vwap.Update(time,close)