''' TickerState
    States a ticker can take while being processed against a brokerage.
'''
from enum import Enum, auto


class TickerState(Enum):
    NEW = auto()
    BUY = auto()
    SHORT = auto()
    LIQUIDATE = auto()

from datetime import datetime, timedelta
from QuantConnect.Data.UniverseSelection import * 
from config import Config
from universe import ScheduledUniverse
from alpha import Alpha
from portfolio import PortfolioConstruction
from execute import OrderExecution
from stop import StopPrices


class OpeningRangeBreakoutDL(QCAlgorithm):
    
    ''' OpeningRangeBreakoutDL
        v1.0
        Opening Range Breakout Using Downloaded Basket of Equities
    '''

    def Initialize(self):
        # Configuration parameters
        self.config = Config(self)

        # Backtesting
        self.SetStartDate(2021, 2, 25) # Set backtest Start Date
        self.SetEndDate(2021, 7, 24)   # Set backtest End Date
        self.SetCash(35000)            # Set backtest Strategy Cash

        # Brokerage
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)

        # Scheduled universe construction
        self.UniverseSettings.Resolution = Resolution.Minute
        self.scheduledUniverse = ScheduledUniverse(self, self.config)
        self.AddUniverseSelection(ScheduledUniverseSelectionModel(
            self.DateRules.WeekEnd(), self.TimeRules.At(23, 50), self.scheduledUniverse.SelectSymbols))

        # Alpha Model
        self.stopPrices = StopPrices()
        self.SetAlpha(Alpha(self.config, self.stopPrices))
        
        # Portfolio Construction
        self.SetPortfolioConstruction(PortfolioConstruction(self.config))
        
        # Execution Model
        self.SetExecution(OrderExecution(self.config, self.stopPrices))

from datetime import timedelta
from ticker import TickerState


class Alpha(AlphaModel):
    '''Alpha - Determine trading signals
       Insights are evaluated every X minutes, e.g. 10,
       but this method is called every minute.
       The indicators are likewise updated every minute.
    '''

    def __init__(self, config, stopPrices):
        self.config = config
        self.stopPrices = stopPrices
        self.symbolDataBySymbol = {}

    def Update(self, algorithm, data):
        ''' Update - Scan symbol data for insights
        
            Called every minute.
        '''

        def GenerateInsight(_sd, _direction, _state):
            # Create and return an Insight.
            # Set the SymbolData direction and state.
            insight = Insight.Price(_sd.symbol, timedelta(minutes=self.config.consolidationPeriod), _direction)
            _sd.direction = _direction
            _sd.state = _state
            return insight
            
        def CalcMidpoint(_high, _low):
            mid = abs((_high - _low) / 2)
            base = _low if _low < _high else _high
            return mid + base
        
        # List of Insights to emit
        insights = []    
        
        #
        # Scan for insights among our universe of symbols.
        #
        removedSymbolData = []
        for s, sd in self.symbolDataBySymbol.items():
            if sd.IsReady:
                
                # The opening bar status for this symbol is valid.
                price = data[sd.symbol].Close if data.Bars.ContainsKey(sd.symbol) else None
                if price is None: 
                    # Tradebar is not ready.
                    continue
                
                if algorithm.Portfolio[sd.symbol].Invested:
                    #
                    # Invested in the ticker.
                    #
                    if sd.state == TickerState.LIQUIDATE:
                        # End Of Day or ticker was removed.
                        insights.append(GenerateInsight(sd, InsightDirection.Flat, TickerState.LIQUIDATE))
                        # Reset opening bar.
                        sd.Clear()
                    if sd.remove:
                        # The ticker has been removed. Portfolio will handle closeout.
                        removedSymbolData.append(s)
                    # If invested, but TickerState is not LIQUIDATE, do nothing;
                    # We are waiting for the market to act.
                else:
                    #
                    # We are not invested in this ticker.
                    #
                    # Only trade during the specified opening period.
                    #
                    if (data.Bars[sd.symbol].Time.hour == self.config.exitTime.hour \
                        and data.Bars[sd.symbol].Time.minute > self.config.exitTime.minute) \
                        or (data.Bars[sd.symbol].Time.hour > self.config.exitTime.hour):
                            # Non-trading time period.
                            if sd.state == TickerState.BUY or sd.state == TickerState.SHORT:
                                # Past time for exit on this ticker, but it remains Ready and we are not invested.
                                # Either timeout, error, or StopMarketOrder executed.
                                algorithm.Log("StopMarketOrder, timeout, or error for {}".format(sd.symbol.Value))
                                # Treat as LIQUIDATE to clear the state in Portfolio.
                                insights.append(GenerateInsight(sd, InsightDirection.Flat, TickerState.LIQUIDATE))
                                # Reset opening bar.
                                sd.Clear()
                            # Always continue to the next ticker.
                            continue
                    #
                    # We are in the trading range time, we are not invested, the opening bar is valid.
                    # Check to see if we've already issued an insight to BUY or SHORT.
                    # If so, continue to wait.
                    # N.B.: TickerState.LIQUIDATE is treated the same as TickerState.NEW
                    # This is because the LIQUIDATE state should have resulted in clearing
                    # the opening bar, which should prevent entry to this part of the state
                    # engine unless the opening bar has been reacquired, and we are still in the 
                    # trading window.
                    #
                    if sd.state == TickerState.BUY or sd.state == TickerState.SHORT:
                        continue
                    # No insight issued yet.
                    # Determine if there is an Opening Range Breakout.
                    if price > sd.openingBar.High:
                        #
                        # Opening Range Breakout HIGH.
                        #
                        # Calculate appropriate StopMarketOrder price.
                        ## Small percentage lower than the high bar.
                        ## self.stopPrices[sd.symbol] = sd.openingBar.High * self.config.initialMargin
                        self.stopPrices[sd.symbol] = CalcMidpoint(sd.openingBar.High, sd.openingBar.Low)
                        # BUY the ticker.
                        insights.append(GenerateInsight(sd, InsightDirection.Up, TickerState.BUY))
                    elif price < sd.openingBar.Low:
                        #
                        # Opening Range Breakout LOW.
                        #
                        # Calculate appropriate StopMarketOrder price.
                        ## Small percentage higher than the low bar.
                        ## self.stopPrices[sd.symbol] = sd.openingBar.Low * (1 + (1 - self.config.initialMargin))
                        self.stopPrices[sd.symbol] = CalcMidpoint(sd.openingBar.High, sd.openingBar.Low)
                        # SHORT the ticker.
                        insights.append(GenerateInsight(sd, InsightDirection.Down, TickerState.SHORT))
        # Remove any deleted tickers.
        for s in removedSymbolData:
            del self.symbolDataBySymbol[s]

        return insights


    def OnSecuritiesChanged(self, algorithm, changes):
        for added in changes.AddedSecurities:
            symbolData = self.symbolDataBySymbol.get(added.Symbol)
            if symbolData is None:
                algorithm.Log("Adding ticker {}".format(added.Symbol.Value))
                algorithm.Securities[added.Symbol].MarginModel = PatternDayTradingMarginModel()
                self.symbolDataBySymbol[added.Symbol] = SymbolData(self.config, algorithm, added.Symbol)
            else:
                algorithm.Log("Resetting ticker {}".format(added.Symbol.Value))
        for removed in changes.RemovedSecurities:
            algorithm.Log("Removing ticker {}".format(removed.Symbol.Value))
            symbolData = self.symbolDataBySymbol.get(removed.Symbol)
            if symbolData is not None:
                if algorithm.Portfolio[removed.Symbol].Invested:
                    # We are invested in this ticker, so we must liquidate it.
                    symbolData.state = TickerState.LIQUIDATE
                symbolData.remove = True
    
    
class SymbolData:
    
    ''' Indicator for one symbol of our universe.
        Each indicator is updated automatically.
    '''
    
    def __init__(self, config, algorithm, symbol):
        self.config = config
        self.symbol = symbol
        self.openingBar = None
        self.direction = None
        self.state = TickerState.NEW
        self.remove = False
        self.consolidator = algorithm.Consolidate(self.symbol, timedelta(minutes=self.config.consolidationPeriod), self.OnDataConsolidated)

    @property    
    def IsReady(self):
        ''' The symbol data is ready when the 
            openingBar has been collected.
        '''
        return self.openingBar is not None
        
    def OnDataConsolidated(self, bar):
        if bar.Time.hour == self.config.entryTime.hour and bar.Time.minute == self.config.entryTime.minute:
            self.openingBar = bar
        elif bar.Time.hour == self.config.closeTime.hour and bar.Time.minute == self.config.closeTime.minute:
            self.ClosePosition()
            
    def ClosePosition(self):
        self.state = TickerState.LIQUIDATE
        
    def Clear(self):
        self.openingBar = None

from datetime import *
from decimal import *

class Config():
    
    ''' Instantiate QuantConnect Parameters '''
    
    def __init__(self, algorithm):
        #
        # Universe parameters
        #
        self.downloadFile = "https://www.dropbox.com/s/o0td3or2jrl4wz0/universe.txt?dl=1"
        # Calculated in universe.py/ScheduledUniverse.SelectSymbols()
        self.universeSize = 0
        #
        # Alpha parameters
        #
        self.entryTime = datetime.strptime("09:30", "%H:%M")
        self.exitTime = datetime.strptime("10:30", "%H:%M")
        self.closeTime = datetime.strptime("13:30", "%H:%M")
        self.consolidationPeriod = int(algorithm.GetParameter("consolidationPeriod"))
        #
        # Portfolio Construction parameters
        #
        self.initialMargin = float(algorithm.GetParameter("initialMargin"))

#
# Generic Simple Execution Model
# 

from clr import AddReference
AddReference("System")
AddReference("QuantConnect.Common")
AddReference("QuantConnect.Algorithm")
AddReference("QuantConnect.Algorithm.Framework")

from System import *
from QuantConnect import *
from QuantConnect.Orders import *
from QuantConnect.Algorithm import *
from QuantConnect.Algorithm.Framework import *
from QuantConnect.Algorithm.Framework.Execution import *
from QuantConnect.Algorithm.Framework.Portfolio import *


class OrderExecution(ExecutionModel):
    
    ''' Immediately submits market orders to achieve the desired portfolio targets,
        with stop loss and end-of-day liquidation.
    '''

    def __init__(self, config, stopPrices):
        '''Initializes a new instance of the simple Execution class'''
        self.config = config
        self.stopPrices = stopPrices

    def Execute(self, algorithm, targets):
        ''' Immediately submits orders for the specified portfolio targets.
            Implicit Args:
                stopPrices  Stop prices for symbols being traded
            Args:
                algorithm:  The algorithm instance
                targets:    The portfolio targets to be ordered
                            Positive quantity denotes an equity purchase
                            Negative quantity denotes an equity sale
                            Zero quantity denotes an equity liquidation
        '''

        for target in targets:
            # Get quantity to be ordered.
            quantity = target.Quantity
            if quantity != 0:
                # BUY or SHORT the equity.
                algorithm.Log("Market order for {} shares of {}".format(quantity, target.Symbol.Value))
                algorithm.MarketOrder(target.Symbol, quantity)
                # Create StopMarketOrder with price calculated by the AlphaModel.
                # algorithm.Log("Stop Market order for {} shares of {} at ${:.2f}".format(quantity, target.Symbol.Value, self.stopPrices[target.Symbol]))
                # algorithm.StopMarketOrder(target.Symbol, quantity, self.stopPrices[target.Symbol])
                _direction = OrderDirection.Buy if quantity > 0 else OrderDirection.Sell
                algorithm.Log("Available cash ${:.2f}, remaining margin ${:.2f}".format(algorithm.Portfolio.Cash, 
                    algorithm.Portfolio.GetMarginRemaining(target.Symbol, _direction)))
                # Remove the Stop Market Order Price, which is matched to each BUY or SHORT.
                # If we do not maintain queue discipline, this will eventually provoke a run time exception.
                del self.stopPrices[target.Symbol]
            else:
                # Liquidating position. Will also cancel any StopMarket order.
                algorithm.Log("Liquidating {}".format(target.Symbol.Value))
                algorithm.Liquidate(target.Symbol)
                algorithm.Log("Available cash now ${:.2f}".format(algorithm.Portfolio.Cash))

    def OnSecuritiesChanged(self, algorithm, changes):
        # Event handler for securities added and removed.
        # changes.AddedSecurities:
        #   We do not need to do anything for AddedSecurities.
        # changes.RemovedSecurities:
        #   Liquidate removed securities.
        for removed in changes.RemovedSecurities:
            algorithm.Log("Liquidating removed security {}".format(removed.Symbol.Value))
            algorithm.Liquidate(removed.Symbol)
            algorithm.Log("Available cash now ${:.2f}".format(algorithm.Portfolio.Cash))

#
# Portfolio Construction
#

from clr import AddReference
AddReference("QuantConnect.Common")
AddReference("QuantConnect.Algorithm.Framework")

from QuantConnect import Resolution, Extensions
from QuantConnect.Algorithm.Framework.Alphas import *
from QuantConnect.Algorithm.Framework.Portfolio import *
from itertools import groupby
from datetime import datetime, timedelta
from math import floor
from ticker import TickerState

class PortfolioConstruction(PortfolioConstructionModel):
    ''' PortfolioConstruction - Construct a portfolio based
        on Insights generated by the AlphaModel.
        
        Liquidation means a ticker has been delisted, or is otherwise 
        unavailable.
    '''
    
    def __init__(self, config):
        self.config = config
        self.symbolDataBySymbol = {}
    
    def CreateTargets(self, algorithm, insights):
        ''' CreateTargets for Portfolio Construction
            Scan insights looking for tickets we can purchase,
            hold, or sell. It is assumed that symbols can only
            appear once in any list of insights.
        '''
        targets = []         # Portfolio targets to be returned for execution
        deletedSymbols = []  # We will delete these symbols.
        # Account for having possibly already expended some number of slices of our cash.
        cashSlicesLeft = self.config.universeSize - len(self.symbolDataBySymbol)
        for insight in insights:
            sd = self.symbolDataBySymbol.get(insight.Symbol)
            if sd is None:
                #
                # We do not own this ticker.
                #
                if cashSlicesLeft == 0:
                    # No cash / margin slices left; cannot BUY or SHORT this ticker.
                    continue
                if insight.Direction == InsightDirection.Up:
                    # Buy this ticker.
                    # Calculate the number of shares to buy.
                    marginRemaining = algorithm.Portfolio.GetMarginRemaining(insight.Symbol, OrderDirection.Buy)
                    cash = algorithm.Portfolio.Cash
                    margin = marginRemaining if marginRemaining <= cash else cash
                    sharePrice = algorithm.Securities[insight.Symbol].Close
                    shares = floor((margin / cashSlicesLeft) / sharePrice)
                    if margin < 0 or shares == 0:
                        algorithm.Log("Not enough buying power to buy long ticker {} {} shares at ${:.2f}".format(insight.Symbol.Value, shares, sharePrice))
                        continue
                    # Create target for appropriate number of shares.
                    algorithm.Log("Buying long ticker {} {} shares at ${:.2f}".format(insight.Symbol.Value, shares, sharePrice))
                    targets.append(PortfolioTarget(insight.Symbol, shares))
                    # Record pending ticker purchase.
                    self.symbolDataBySymbol[insight.Symbol] = SymbolData(self.config, 
                        insight.Symbol, sharePrice, shares, TickerState.BUY)
                elif insight.Direction == InsightDirection.Down:
                    # Short the ticker.
                    # Calculate the number of shares to sell.
                    margin = algorithm.Portfolio.GetMarginRemaining(insight.Symbol, OrderDirection.Sell) * self.config.initialMargin
                    sharePrice = algorithm.Securities[insight.Symbol].Close
                    shares = floor((margin / cashSlicesLeft) / sharePrice)
                    if margin < 0 or shares == 0:
                        algorithm.Log("Not enough buying power to sell short ticker {} {} shares at ${:.2f}".format(insight.Symbol.Value, -shares, sharePrice))
                        continue
                    # Create target for appropriate number of shares.
                    algorithm.Log("Short selling ticker {} {} shares at ${:.2f}".format(insight.Symbol.Value, -shares, sharePrice))
                    targets.append(PortfolioTarget(insight.Symbol, -shares))
                    # Record pending ticker purchase.
                    self.symbolDataBySymbol[insight.Symbol] = SymbolData(self.config, 
                        insight.Symbol, sharePrice, -shares, TickerState.SHORT)
            else:
                #
                # We own this ticker.
                #
                if insight.Direction == InsightDirection.Flat \
                    or sd.state == TickerState.LIQUIDATE:
                    # Liquidate the ticker.
                    # We expect shares to be zero for this insight, 
                    # but in fact we don't care; this might be a delisted stock.
                    algorithm.Log("Liquidating ticker {} {} shares at ${:.2f}".format(sd.symbol.Value, sd.shares, sd.price))
                    # Liquidate the position.
                    targets.append(PortfolioTarget(insight.Symbol, 0))
                    sd.state = TickerState.LIQUIDATE
                    deletedSymbols.append(sd)
                else:
                    algorithm.Log("ERROR: Unexpected insight direction or ticker state, should only be Flat or LIQUIDATE, \ninsight direction {}, ticker state {}".format(insight.Direction, sd.state))
                    deletedSymbols.append(sd)
        # Remove the deleted symbols.
        for sd in deletedSymbols:
            self.symbolDataBySymbol.pop(sd.symbol, None)

        # Return the portfolio targets.
        return targets


    def OnSecuritiesChanged(self, algorithm, changes):
        # Event handler for securities added and removed.
        # changes.AddedSecurities:
        #   We do not need to do anything for AddedSecurities.
        # changes.RemovedSecurities:
        for removed in changes.RemovedSecurities:
            sd = self.symbolDataBySymbol.get(removed.Symbol)
            if sd is not None:
                sd.state = TickerState.LIQUIDATE


class SymbolData:
    
    ''' Symbol data for each ticker held in our portfolio '''
    
    def __init__(self, config, symbol, price, shares, state):
        self.config = config
        self.symbol = symbol
        self.price = price
        self.shares = shares
        self.state = state

class StopPrices(dict):
    
    """ StopPrices
        Dictionary holding Stop Market prices for equities being traded.
    """
    
    def __init__(self, **kwargs):
        # noinspection PyTypeChecker
        dict.__init__(self, kwargs)
        self.__dict__=self

class ScheduledUniverse():
    
    '''Select a new universe on a schedule by downloading a file.
    
       The file must have one ticker per line.
       Each line has three string fields, separated by commas ",".

       1 - Ticker
       2 - Security Type
       3 - Market
       
       If the Security Type or Market cannot be parsed, they are
       replaced with SecurityType.Equity and Market.USA, respectively.

       Severly malformed files may result in an abort of the trading
       algorithm.
    '''
    
    def __init__(self, algorithm, config):
        self.algorithm = algorithm
        self.config = config
        
    def SelectSymbols(self, DateTime):
        symbols = []
        self.algorithm.Log("Reading symbol file {}".format(self.config.downloadFile))
        file = self.algorithm.Download(self.config.downloadFile)
        for line in file.splitlines():
            (symbolStr, securityTypeStr, marketStr) = line.split(sep=",")
            self.algorithm.Log("Adding symbol {} of type {} in market {}".format(symbolStr, securityTypeStr, marketStr))
            securityType = self.ParseSecurityType(securityTypeStr)
            market = self.ParseMarket(marketStr)
            symbols.append(Symbol.Create(symbolStr, securityType, market))
        self.config.universeSize = len(symbols)
        self.algorithm.Log("Universe size {} symbols".format(self.config.universeSize))
        return symbols
        
    securityStr2Type = {
        "Equity" : SecurityType.Equity,
        "Option" : SecurityType.Option,
        "Future" : SecurityType.Future,
        "Forex"  : SecurityType.Forex,
        "Crypto" : SecurityType.Crypto,
        "Cfd"    : SecurityType.Cfd,
    }
    
    def ParseSecurityType(self, securityTypeStr):
        securityType = SecurityType.Equity
        for s in ScheduledUniverse.securityStr2Type.keys():
            if securityTypeStr == s:
                securityType = ScheduledUniverse.securityStr2Type[s]
                break
        return securityType
        
    marketStr2Type = {
        "USA"      : Market.USA,
        "FXCM"     : Market.FXCM,
        "GDAX"     : Market.GDAX,
        "Bitfinex" : Market.Bitfinex,
        "Oanda"    : Market.Oanda,
    }
    
    def ParseMarket(self, marketTypeStr):
        market = Market.USA
        for s in ScheduledUniverse.marketStr2Type.keys():
            if marketTypeStr == s:
                market = ScheduledUniverse.marketStr2Type[s]
                break
        return market