from SymbolData import SymbolData
from TradeManagement import TradeManagement

class CryingBlueFlamingo(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021, 2, 1)  # Set Start Date
        self.SetEndDate(2021, 2, 10)  # Set End Date
        self.SetCash(1000000)  # Set Strategy Cash
        
        self.benchmark = "SPY"
        self.AddEquity(self.benchmark)
        
        self.AddUniverse(self.CoarseSelection)
        self.UniverseSettings.Resolution = Resolution.Minute
        
        self.universe_size = 3
        
        self.symbol_data = {}
        self.trade_managers = {}
        
        # Want to open a short position before the market closes
        self.Schedule.On(self.DateRules.EveryDay(self.benchmark), self.TimeRules.BeforeMarketClose(self.benchmark, 1), self.Rebalance)
        
        # Every morning create new rolling windows
        self.Schedule.On(self.DateRules.EveryDay(self.benchmark), self.TimeRules.AfterMarketOpen(self.benchmark, 0), self.InitializeRollingWindows)
    
    def Rebalance(self):
        '''Fires everyday 1 minute before market close'''
        
        for symbol, symbol_data in self.symbol_data.items():
            
            if not symbol_data.IsReady:
                continue
            
            signal = self.CalculateSignal(symbol_data)
            #hammer_signal = self.CalculateSignal(symbol_data)
            
            # Go short if there is a Hanging Man Signal 
            if signal:
                
                trade_manager = self.trade_managers[symbol]
                
                #trade_manager.CreateEntry(-10)
                trade_manager.CreateEntry(-.1)
            '''
            # Go long if there is a hammer signal 
            if hammer_signal:
                
                trade_manager = self.trade_managers[symbol]
                
                trade_manager.CreateEntry(10)
            '''
    def OnData(self, data):
        
        for symbol, trade_manager in self.trade_managers.items():
            
            if not self.Portfolio[symbol].Invested:
                continue
            
            current_price = self.Securities[symbol].Price
            stop_loss = trade_manager.stop_loss
            take_profit = trade_manager.take_profit
            
            if current_price > stop_loss or current_price < take_profit:
                trade_manager.Liquidate()
            
        
    #Finds a Red Hanging Man candle whose High is higher than the prior 5 days highs, and above the 20SMA.
    def CalculateSignal(self, symbol_data):
        
        sma = symbol_data.sma.Current.Value
        atr = symbol_data.atr.Current.Value
        
        # Daily bars
        bars = symbol_data.bar_window
        
        # Minute bars
        #self.Debug(f"Rolling window for {symbol_data.symbol} on {self.Time} is size {symbol_data.minute_bar_window.Count}")
        symbol_data.CalculateOHLC()
        max_price = max([x.High for x in symbol_data.minute_bar_window])
        #self.Debug(f"Max price over minute consolidators for {symbol_data.symbol} on {self.Time} is {max_price}")
        low_price = min([x.Low for x in symbol_data.minute_bar_window])
        #self.Debug(f"Low price over minute consolidators for {symbol_data.symbol} on {self.Time} is {low_price}")
        latest_daily_bar = symbol_data.summary_bar
        latest_consolidator = symbol_data.minute_bar_window[0]
        first_consolidator = symbol_data.minute_bar_window[symbol_data.minute_bar_window.Count-1]

        self.Debug(f"EARLIEST Minute Bar Consolidator OHLC for {symbol_data.symbol} on {self.Time} has total period of {first_consolidator.Period} and start time of {first_consolidator.Time} ending at {first_consolidator.EndTime}")
        self.Debug(f"LATEST Minute Bar Consolidator OHLC for {symbol_data.symbol} on {self.Time} has total period of {latest_consolidator.Period} and start atime of {latest_consolidator.Time} ending at {latest_consolidator.EndTime}")
        self.Debug(f"Minute Bar Consolidator OHLC for {symbol_data.symbol} on {self.Time} is {latest_daily_bar}")
        self.Debug(f"Daily Bar Consolidator OHLC for {symbol_data.symbol} on {self.Time} is {bars[0]}")
        symbol_data.KillMinuteConsolidator()
        self.Debug(f"Minute Bar Consolidator for {symbol_data.symbol} has terminated")
        
        # checking if the high of the latest daily bar is greater than the high of all following bars
        uptrend = all([latest_daily_bar.High > bar.High for bar in list(bars)[:5]])
        downtrend = all([latest_daily_bar.Low < bar.Low for bar in list(bars)[:5]]) 
        
        
        red_bar = latest_daily_bar.Close < latest_daily_bar.Open
        #green_bar = latest_daily_bar.Close > latest_daily_bar.Open
        
        if red_bar:
        
            body = abs(latest_daily_bar.Open - latest_daily_bar.Close)
            shadow = abs(latest_daily_bar.Close - latest_daily_bar.Low)
            wick = abs(latest_daily_bar.High - latest_daily_bar.Open)
            #dayATR = abs(latest_daily_bar.High - latest_daily_bar.Low)
            hanging_man = (shadow > 2 * body) and  (wick < 0.3 * body)
        '''
        if  green_bar:
        
            body = abs(latest_daily_bar.Close - latest_daily_bar.Open)
            shadow = abs(latest_daily_bar.Open - latest_daily_bar.Low)
            wick = abs(latest_daily_bar.High - latest_daily_bar.Close)
            dayATR = abs(latest_daily_bar.High - latest_daily_bar.Low)
            green_hammer = (shadow > 2 * body) and  (wick < 0.3 * body)
        '''
        
        # latest_market_price 
        price = self.Securities[symbol_data.symbol].Price
        above_sma = latest_daily_bar.Close > sma
        below_sma = latest_daily_bar.Close < sma
        
        #Hanging Man Signal 
        signal = red_bar and uptrend and hanging_man and above_sma 
        
        #Hammer Signal
        #hammer_signal = green_bar and downtrend and green_hammer and below_sma 
        
        if signal:
            self.Debug(f" Signal Candle for {symbol_data.symbol} on {self.Time} is - Body: {body} , Wick: {wick} , shadow: {shadow}")
            self.Debug(f"Minute Bar Consolidator OHLC for Signal Day {symbol_data.symbol} on {self.Time} is {latest_daily_bar}")
            return signal
        '''
        if hammer_signal:
            return hammer_signal
        '''
    
    def CoarseSelection(self, coarse):
        
        # list of ~8500 stocks (coarse data)
        
        # coarse is a list of CoarseFundamental objects
        # Descending order
        sorted_by_liquidity = sorted(coarse, key=lambda c:c.DollarVolume, reverse=True)
        
        most_liquid_coarse = sorted_by_liquidity[:self.universe_size]
        
        # needs to return a list of Symbol object
        most_liquid_symbols = [c.Symbol for c in most_liquid_coarse]
        
        return most_liquid_symbols
        
    
    def OnSecuritiesChanged(self, changes):
        '''Fires after universe selection if there are any changes'''
        for security in changes.AddedSecurities:
            symbol = security.Symbol
            
            if symbol not in self.symbol_data and symbol.Value != self.benchmark:
                self.symbol_data[symbol] = SymbolData(self, symbol)
                self.trade_managers[symbol] = TradeManagement(self, symbol)
                
        
        for security in changes.RemovedSecurities:
            symbol = security.Symbol
            
            if symbol in self.symbol_data:
                symbol_data_object = self.symbol_data.pop(symbol, None)
                symbol_data_object.KillDailyConsolidator()
            
            if symbol in self.trade_managers:
                self.trade_managers.pop(symbol, None)
                
    def InitializeRollingWindows(self):
        
        for symbol, symbol_data in self.symbol_data.items():
            
            symbol_data.InitializeMinuteConsolidator()
            self.Debug(f"Minute Bar Consolidator for {symbol_data.symbol} has started")

from SymbolData import SymbolData
class TradeManagement:
    
    def __init__(self, algorithm, symbol):
        
        self.algorithm = algorithm
        self.symbol = symbol
        
        self.days_active = 0
        
        self.entry_price = None
        self.stop_loss = None
        self.take_profit = None
        
    def CreateEntry(self, quantity):
        
        # initial entry market order
        #self.algorithm.MarketOrder(self.symbol, quantity)
        self.algorithm.SetHoldings(self.symbol, quantity)
        
        current_price = self.algorithm.Securities[self.symbol].Price
        atr = self.algorithm.symbol_data[self.symbol].atr.Current.Value
        
        self.entry_price = current_price
        
        self.stop_loss = self.entry_price + 0.5 * atr # Use High from the current day 
        self.take_profit = self.entry_price - 1 * atr
        
        self.algorithm.Debug(f"Entering {self.symbol} on {{self.Time}}...Entry Price: {current_price}, Take Profit: {self.take_profit}, StopLoss: {self.stop_loss}")
        
        
    def Liquidate(self):
        
        self.algorithm.Debug(f"Liquidating.. {self.symbol}....{self.algorithm.Securities[self.symbol].Price}")
        
        self.algorithm.Liquidate(self.symbol)
        
        self.entry_price = None
        self.stop_loss = None
        self.take_profit = None
        
        self.days_active = 0

class SymbolData:
    '''Containers to hold relevant data for each symbol'''
    def __init__(self, algorithm, symbol):
                
        self.algorithm = algorithm
        self.symbol = symbol
        
        # defines daily consolidator and then registers to receive data
        self.daily_consolidator = TradeBarConsolidator(timedelta(days=1))
        self.algorithm.SubscriptionManager.AddConsolidator(symbol, self.daily_consolidator)
        self.daily_consolidator.DataConsolidated += self.OnDailyBar
        
        # 1. instantiantes a SimpleMovingAverage object
        # 2. subscribes it to receive data
        self.sma = SimpleMovingAverage(10) # Test 10 vs 20 
        self.algorithm.RegisterIndicator(symbol, self.sma, self.daily_consolidator)
        
        self.atr = AverageTrueRange(1)
        self.algorithm.RegisterIndicator(symbol, self.atr, self.daily_consolidator)
        
        # holds recent bars
        self.bar_window = RollingWindow[TradeBar](5)
        self.minute_bar_window = RollingWindow[TradeBar](388) 
        
        self.WarmUpIndicators()
    
    def WarmUpIndicators(self):
        
        # returns a dataframe
        history = self.algorithm.History(self.symbol, 20, Resolution.Daily)
        
        for bar in history.itertuples():
            
            time = bar.Index[1]
            
            open = bar.open
            high = bar.high
            low = bar.low
            close = bar.close
            volume = bar.volume
            
            trade_bar = TradeBar(time, self.symbol, open, high, low, close, volume)
            
            self.sma.Update(time, close)
            self.atr.Update(trade_bar)
            self.bar_window.Add(trade_bar)
        
    def OnDailyBar(self, sender, bar):
        '''Fires each time our daily_consolidator produces a bar
        that bar is passed in through the bar parameter'''
        
        # save that bar to our rolling window
        self.bar_window.Add(bar)

        
    def OnMinuteBar(self, sender, bar):
        '''Fires each time our minute_consolidator produces a bar
        that bar is passed in through the bar parameter'''
        
        # save that bar to our rolling window
        self.minute_bar_window.Add(bar)
        
        
        sorted(self.minute_bar_window, key = lambda thing: thing.Time)  
        
    def KillDailyConsolidator(self):
        self.algorithm.SubscriptionManager.RemoveConsolidator(self.symbol, self.daily_consolidator)
        
    def KillMinuteConsolidator(self):
        self.algorithm.SubscriptionManager.RemoveConsolidator(self.symbol, self.minute_consolidator)
    
    def IsReady(self):
        
        return self.sma.IsReady and self.atr.IsReady and self.bar_window.IsReady
        
    def InitializeMinuteConsolidator(self):
        """Initialize a minute consolidator for this symbol"""
        
        self.minute_consolidator = TradeBarConsolidator(timedelta(minutes=1))#5
        self.algorithm.SubscriptionManager.AddConsolidator(self.symbol, self.minute_consolidator)
        self.minute_consolidator.DataConsolidated += self.OnMinuteBar        
        
        
    def CalculateOHLC(self):
        
        # Rolling window open
        opening_bar = self.minute_bar_window[self.minute_bar_window.Count-1]
        open = opening_bar.Open
        
        # Rolling window close
        closing_bar = self.minute_bar_window[0]
        close = closing_bar.Close
        
        # High and low over period
        high = max([x.High for x in self.minute_bar_window])
        low = min([x.Low for x in self.minute_bar_window])
        
        # Calculate volume
        volume = sum([x.Volume for x in self.minute_bar_window])
        
        # Time
        time = self.minute_bar_window[self.minute_bar_window.Count-1].Time
        
        # Create a summary trade bar
        self.summary_bar = TradeBar(time, self.symbol, open, high, low, close, volume)