"""
This file contains QuantConnect order codes for easy conversion and more 
intuitive custom order handling

References:
    https://github.com/QuantConnect/Lean/blob/master/Common/Orders/OrderTypes.cs
    https://github.com/QuantConnect/Lean/blob/master/Common/Orders/OrderRequestStatus.cs
"""

OrderTypeKeys = [
    'Market', 'Limit', 'StopMarket', 'StopLimit', 'MarketOnOpen',
    'MarketOnClose', 'OptionExercise',
]

OrderTypeCodes = dict(zip(range(len(OrderTypeKeys)), OrderTypeKeys))

OrderDirectionKeys = ['Buy', 'Sell', 'Hold']
OrderDirectionCodes = dict(zip(range(len(OrderDirectionKeys)), OrderDirectionKeys))

## NOTE ORDERSTATUS IS NOT IN SIMPLE NUMERICAL ORDER

OrderStatusCodes = {
    0:'New', # new order pre-submission to the order processor
    1:'Submitted', # order submitted to the market
    2:'PartiallyFilled', # partially filled, in market order
    3:'Filled', # completed, filled, in market order
    5:'Canceled', # order cancelled before filled
    6:'None', # no order state yet
    7:'Invalid', # order invalidated before it hit the market (e.g. insufficient capital)
    8:'CancelPending', # order waiting for confirmation of cancellation
}

import pandas as pd
import numpy as np
import decimal as d
from datetime import datetime, timedelta, time
from order_codes import (OrderTypeCodes, OrderDirectionCodes, OrderStatusCodes)

#--------------------------#
# Globals                  |
#--------------------------#

PT = 0.005 # percent
SL = 0.005 # percent
EXPIRY_MIN_DAYS = 30 # days
EXPIRY_LIQUIDATE_DAYS = 10 # days

#--------------------------#
# ALGORITHM                |
#--------------------------#

class BasicTemplateAlgorithm(QCAlgorithm):
    '''Basic template algorithm simply initializes the date range and cash'''

    def Initialize(self):

        self.SetStartDate(2016,7, 7)  #Set Start Date
        self.SetEndDate(2016,7,30)    #Set End Date
        self.SetCash(500000)          #Set Strategy Cash
        
        # brokerage model
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage,
                               AccountType.Margin)   
        # Find more symbols here: http://quantconnect.com/data
        futureES = self.AddFuture(Futures.Indices.SP500EMini, Resolution.Minute)
        futureES.SetFilter(TimeSpan.Zero, TimeSpan.FromDays(185))         
        futureNQ = self.AddFuture(Futures.Indices.NASDAQ100EMini, Resolution.Minute)
        futureNQ.SetFilter(TimeSpan.Zero, TimeSpan.FromDays(185)) 
        
        # init empty dict for order data
        self.order_tickets = dict()
        
        # must add contract multiple for accurate order handling
        self.contract_multiple = {}
        self.contract_multiple['ES'] = 50
        self.contract_multiple['NQ'] = 20        
    
    def OnData(self, slice):
        
        # only run during market hours                         
        #------------------------------------------------------#
        start_time = time(hour=9, minute=35)
        end_time = time(hour=16, minute=0)
        if not start_time < self.Time.time() < end_time: return
        
        # check if holdings need to be sold based on expiry
        #------------------------------------------------------#
        self._check_holdings_expiry()
        
        # check if bracket order has been sent
        # sometimes OnOrderEvent doesn't send the bracket order
        #------------------------------------------------------#
        self._confirm_bracket_order()
        
        base_symbol_invested = [self._get_base_symbol(x.Symbol) \
                                for x in self.Portfolio.Values if x.Invested]
        
        for chain in slice.FutureChains:
            
            # check if algo is already in invested in symbol
            #------------------------------------------------------#
            chain_base_sym = chain.Value.Symbol.Value
            if chain_base_sym in base_symbol_invested: continue
            
            # find the front contract expiring no earlier than in N days
            #------------------------------------------------------#
            contracts = [i for i in chain.Value \
                         if i.Expiry > self.Time+timedelta(EXPIRY_MIN_DAYS)]
            
            # if there is more than one contract, 
            # trade the one with the closest expiry date
            #------------------------------------------------------#
            if len(contracts) > 0:
                front = sorted(contracts, key=lambda x: x.Expiry, reverse=True)[0]
                sym = front.Symbol
               
                self.Debug('front contract: {}'.format(sym))
                
                # compute order quantity based on portfolio percentage
                #------------------------------------------------------#
                qty = self._calc_order_quantity(sym, pct=0.4)
                if qty is None:
                    self.Debug('price is zero for: {}'.format(sym.Value))
                    return
                
                # if quantity is good we place a market order and 
                # add the order data to the symbol data dict
                #------------------------------------------------------#
                self.Debug('front contract qty: {}'.format(qty))
                newTicket = self.MarketOrder(sym, qty) 
                self.order_tickets[sym.Value] = \
                    symbolOrderData(sym.Value, newTicket, 'Market', contract=front)
                
    def OnOrderEvent(self, orderEvent):
        """
        This function is triggered automatically every time an order event occurs.
        """
        
        self.Log(str(orderEvent))  
        
        if OrderStatusCodes[orderEvent.Status] == 'Submitted':
            return

        k = str(orderEvent.Symbol.Value)
        symbol = str(orderEvent.Symbol)
        
        if (not k in self.order_tickets.keys()): 
            self.Log('missing key in order tickets: {}'.format(k))
            self.Log('order tickets keys: {}'.format(self.order_tickets.keys()))
            return
        
        elif (k in self.order_tickets.keys()):
        
            orderData = self.order_tickets[k]
            orig_order_id = orderData.ticket.OrderId
            order = self.Transactions.GetOrderTicket(orig_order_id)
            
            # sometimes order is nonetype due to security price
            # is equal to zero
            #------------------------------------------------------#
            if not order: 
                self.Log('order is nonetype: {}'.format(k))
                del self.order_tickets[k] # delete order ticket data  
                return
                
            # if order is filled but bracket order not submitted,
            # submit bracket order
            #------------------------------------------------------#
            if (OrderStatusCodes[order.Status]=='Filled') and \
                (not orderData.bracket_submit):
                    
                if (orderEvent.OrderId == orig_order_id):
            
                    price = orderEvent.FillPrice
                    qty = orderEvent.FillQuantity
                    #qty = self.CalculateOrderQuantity(k, 0.0)
                    
                    profit_target = price * d.Decimal(1+PT)
                    limitTicket = self.LimitOrder(symbol, -1*qty, profit_target)
                    self.order_tickets[k].add_limit_order(limitTicket)
                    
                    stop_loss = price * d.Decimal(1-SL)
                    stopTicket = self.StopMarketOrder(symbol, -1*qty, stop_loss)
                    self.order_tickets[k].add_stop_market_order(stopTicket)
                    
                    self.order_tickets[k].is_bracket()
                    
                    self.Log('bracket order submitted: {}'.format(self.Time, k))
                    
            # Otherwise, one of the exit orders was filled,
            # so cancel the open orders
            #------------------------------------------------------#
            elif (orderData.bracket_submit) and \
                (OrderStatusCodes[orderEvent.Status]=='Filled'):
                self.Log('cancelling bracket orders for: {}'.format(self.Time, k))
                self.Transactions.CancelOpenOrders(symbol)
                self.order_tickets[k].bracket_submit=False
                del self.order_tickets[k]    
                
    #--------------------------#
    # ALGORITHM HELPER FUNCS   |
    #--------------------------#                
                
    def _calc_order_quantity(self, sym, pct):
        """
        Compute order quantity based on percentage of portfolio value.
        This is required because SetHoldings doesn't return an order ticket.
        
        This function returns None when security price is less than 1.
        """
        price = float(self.Securities[sym].Price)
        self.Log('{} calc order quantity price: {}'.format(sym, price))
        if price < 1: return None
        numr = pct * float(self.Portfolio.TotalPortfolioValue)
        denm = price * self.contract_multiple[self._get_base_symbol(sym)]
        qty =  numr/denm 
        return int(qty)
    
    def _get_base_symbol(self, symbol):
        """
        Get minimum symbol for comparisons.
            Example: 'ES XXXXXXXX' will return 'ES'
        """
        sym = str(symbol).split(' ')[0]
        return sym
        
    def _check_holdings_expiry(self):
        """
        Check if current holdings are expiring within certain number
        of days. If so liquidate the contract, cancel open orders and delete
        order data from dict.
        """
        
        invested = [x.Symbol for x in self.Portfolio.Values if x.Invested]
        #self.Log('[{}] invested:\n{}'.format(self.Time, invested))
        
        for sym in invested:
            
            k = sym.Value

            if (k in self.order_tickets.keys()):
                
                expiry = self.order_tickets[k].contract.Expiry
                
                if (expiry < self.Time + timedelta(EXPIRY_LIQUIDATE_DAYS)):
                    self.Debug('expiry too close liquidating: {}'.format(k))
                    self.Liquidate(sym)
                    self.Transactions.CancelOpenOrders(sym)
                    del self.order_tickets[k]     
        return
    
    def _confirm_bracket_order(self):
        """
        confirm bracket orders have been submitted.
        sometimes OnOrderEvent is not sending orders
        """
        
        invested = [x.Symbol for x in self.Portfolio.Values if x.Invested]
        #self.Log('[{}] invested:\n{}'.format(self.Time, invested))
        
        for sym in invested:
            
            k = sym.Value

            if (k in self.order_tickets.keys()): 
                
                orderData = self.order_tickets[k]
                orig_order_id = orderData.ticket.OrderId
                order = self.Transactions.GetOrderTicket(orig_order_id)

                if (not orderData.bracket_submit):
                    
                    price = order.AverageFillPrice
                    qty = order.QuantityFilled
                    
                    profit_target = price * d.Decimal(1+PT)
                    limitTicket = self.LimitOrder(sym, -1*qty, profit_target)
                    self.order_tickets[k].add_limit_order(limitTicket)
                    
                    stop_loss = price * d.Decimal(1-SL)
                    stopTicket = self.StopMarketOrder(sym, -1*qty, stop_loss)
                    self.order_tickets[k].add_stop_market_order(stopTicket)
                    
                    self.order_tickets[k].is_bracket()
                    
                    self.Log('check bracket >> bracket order submitted: {}'.format(k))
        return
            
                
            
class symbolOrderData:
    """
    Class object to hold symbol bracket order data
    
    Attributes:
    -----------
        symbol: str 
        ticket: ticket object 
            does not work with setHoldings
        order_type: str
            must match QC order types in order_codes.py
        bracket_submit: bool
            toggled when bracket order is submitted
        limit_order: None or ticket object
        stop_market_order: None or ticket object
        contract: None or futures contract object
    
    Methods:
    --------
        add_limit_order: set limit_order attr to limitTicket
        add_stop_market_order: set stop_market_order attr to stopMarketTicket
        is_bracket: if limit and stopMarket orders are submitted toggle
            bracket_submit attr
    """
    
    def __init__(self, symbol, ticket, order_type, contract=None):
        
        self.symbol = symbol
        self.ticket = ticket
        self.order_type = order_type
        if contract:
            self.contract = contract
        self.bracket_submit = False  
        
        self.limit_order = None
        self.stop_market_order = None
        
    def add_limit_order(self, limitTicket):
        self.limit_order = limitTicket
        return
    
    def add_stop_market_order(self, stopMarketTicket):
        self.stop_market_order = stopMarketTicket
        return
    
    def is_bracket(self):
        if (self.limit_order is not None) and (self.stop_market_order is not None):
            self.bracket_submit = True
        return