'''
Left to do  - 
    --1. 2 Day close rule- not closing, also make it close at the end of the 2nd day
    2. Stoploss didnt work 
        2020-08-05 15:59:00 :    Entering SRNE VL1GER10SC4L on {self.Time}...Entry Price: 13.74, Take Profit: 11.49, StopLoss: 14.865
            Shoud've hit stoploss @ 14.87 on 8/7/2020 but holds position entire time 
    ---3. Stoploss should be
                latest_daily_bar.High + 0.5 * atr
            vs 
                self.entry_price + 0.5 * atr
    ---4. Allow for 2 positions at a time @ 15% equity each 
'''
from SymbolData import SymbolData
from TradeManagement import TradeManagement

class CryingBlueFlamingo(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 7, 1)  # Set Start Date
        self.SetEndDate(2021, 5, 18)  # Set End Date
        self.SetCash(1000000)  # Set Strategy Cash
        
        self.max_positions = 2
        self.total_short_margin_allocation = 0.30
        self.universe_size = 30
        
        self.benchmark = "SPY"
        self.AddEquity(self.benchmark)
        
        # fast way of modelling margin of shorting
        # SetHoldings(0.5) -> 50% of the 30% allowed is allocated aka 15% is allocated
        # self.Settings.FreePortfolioValuePercentage = 0.30
        
        self.AddUniverse(self.CoarseSelection)
        self.UniverseSettings.Resolution = Resolution.Minute
        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)
    
    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
            
            #if symbol.Value == 'SRNE':
            
                #self.Debug(f"SRNE: {self.Portfolio[symbol].Invested}, {self.Portfolio[symbol].Quantity}, {self.trade_managers[symbol].days_active}")
            
            signal = self.CalculateSignal(symbol_data)
            '''
            hammer_signal = self.CalculateSignal(symbol_data)
            '''
            trade_manager = self.trade_managers[symbol]
            
            number_of_open_positions = len([symbol for symbol in self.Portfolio.Keys if self.Portfolio[symbol].Invested])
            
            # Go short if there is a Hanging Man Signal 
            if signal and not self.Portfolio[symbol].Invested and number_of_open_positions < self.max_positions:
                
                # trade_manager = self.trade_managers[symbol]
                
                #trade_manager.CreateEntry(-10)
                portfolio_allocation = self.total_short_margin_allocation / self.max_positions
                trade_manager.CreateEntry(-portfolio_allocation)
                
                self.Debug(f"{symbol}, {self.Portfolio[symbol].Invested}")
                
            if self.Portfolio[symbol].Invested:
                
                trade_manager.days_active += 1
                self.Debug(f"{symbol} : days active: {trade_manager.days_active}")
            
            '''
            # 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
            days_active = trade_manager.days_active
            
            if current_price > stop_loss or current_price < take_profit:
                trade_manager.Liquidate()
            
            if days_active is 3:
                trade_manager.Liquidate()
                #trade_manager.OnMarketCloseLiquidate()
                self.Debug(f"{symbol} -- held for {days_active}...Liquidating at market close")
    
    #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):
        
        # Daily bars
        bars = symbol_data.bar_window
        
        # Minute bars
        symbol_data.CalculateOHLC()
        max_price = max([x.High for x in symbol_data.minute_bar_window])
        low_price = min([x.Low for x in symbol_data.minute_bar_window])
        latest_daily_bar = symbol_data.summary_bar
        latest_consolidator = symbol_data.todays_minute_bars[0]
        first_consolidator = symbol_data.todays_minute_bars[-1]
        number_of_bars_today = len(symbol_data.todays_minute_bars)
        
        # 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)[:6]])
        downtrend = all([latest_daily_bar.Low < bar.Low for bar in list(bars)[:6]]) 
        
        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)
            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)
        '''
        sma = (sum([b.Close for b in list(bars)[:-1]]) + latest_daily_bar.Close) / 10
        
        # 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 self.Portfolio[symbol].Invested:
                self.trade_managers[symbol].Liquidate()    
                
            # if not self.Portfolio[symbol].Invested:
            self.RemoveSecurityFromDictionaries(symbol)
            
    def RemoveSecurityFromDictionaries(self, symbol):
        
        if symbol in self.symbol_data:
                symbol_data_object = self.symbol_data.pop(symbol, None)
                symbol_data_object.KillDailyConsolidator()
                symbol_data_object.KillMinuteConsolidator()
            
        if symbol in self.trade_managers:
                self.trade_managers.pop(symbol, None)

from SymbolData import *

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.SetHoldings(self.symbol, quantity)
        
        current_price = self.algorithm.Securities[self.symbol].Price
        
        symbol_data = self.algorithm.symbol_data[self.symbol]
        
        # Update our 1 period ATR with latest bar, so we have today's range
        symbol_data.atr.Update(symbol_data.summary_bar)
        
        atr = symbol_data.atr.Current.Value
        
        self.entry_price = current_price
        
        summary_bar = self.algorithm.symbol_data[self.symbol].summary_bar
        
        self.stop_loss = summary_bar.High + 0.5 * atr 
        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
    
        # in_universe = False
        # # checking if symbol shows up in any of the defined universes
        # for universe in self.algorithm.UniverseManager.Values:
            
        #     if self.symbol in universe.Members.Keys:
        #         in_universe = True
                
        # # if that symbol does not exist in any universe,
        # # we remove subscriptions and remove from list after liquidation
        # if not in_universe:
        #     self.algorithm.RemoveSecurityFromDictionaries(self.symbol)
                

class SymbolData:
    '''Containers to hold relevant data for each symbol'''
    def __init__(self, algorithm, symbol):
                
        self.algorithm = algorithm
        self.symbol = symbol
        
        # self.minute_consolidator = self.algorithm.SubscriptionManager.ResolveConsolidator(Resolution.Minute)
        self.minute_consolidator = TradeBarConsolidator(timedelta(minutes=1))
        self.algorithm.SubscriptionManager.AddConsolidator(self.symbol, self.minute_consolidator)
        self.minute_consolidator.DataConsolidated += self.OnMinuteBar
        
        # 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](10)
        self.minute_bar_window = RollingWindow[TradeBar](500) 
        
        self.summary_bar = None
        
        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)
        
    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 and self.minute_bar_window.IsReady
        
    def CalculateOHLC(self):
        
        # Rolling window open
        bars = list(self.minute_bar_window)
        todays_bars = [bar for bar in bars if bar.Time.day == self.algorithm.Time.day]
        
        # desecending in time, larger indices -> further in past
        todays_bars_sorted = sorted(todays_bars, key=lambda b:b.Time, reverse=True)
        
        if len(todays_bars_sorted) == 0:
            self.algorithm.Debug(f"{self.symbol} -- {len(todays_bars_sorted)}")
        
        opening_bar = todays_bars_sorted[-1]
        open = opening_bar.Open
        
        # Rolling window close
        closing_bar = todays_bars_sorted[0]
        close = closing_bar.Close
        
        # High and low over period
        high = max([x.High for x in todays_bars_sorted])
        low = min([x.Low for x in todays_bars_sorted])
        
        # Calculate volume
        volume = sum([x.Volume for x in todays_bars_sorted])
        
        # Time
        time = opening_bar.Time
        
        period = TimeSpan.FromMinutes((self.algorithm.Time - time).seconds // 60)
        
        # Create a summary trade bar
        self.summary_bar = TradeBar(time, self.symbol, open, high, low, close, volume, period)
    
    @property
    def todays_minute_bars(self):
        
        bars = list(self.minute_bar_window)
        
        # self.Debug(f"Filtering bars for {self.symbol} ON....{self.algorithm.Time.day}")
        todays_bars = [bar for bar in bars if bar.Time.day == self.algorithm.Time.day]
        
        
        # desecending in time, larger indices -> further in past
        todays_bars_sorted = sorted(todays_bars, key=lambda b:b.Time, reverse=True)
        
        return todays_bars_sorted