from datetime import datetime,timedelta
import pandas as pd
import numpy as np
import sys
import math
from QuantConnect.Algorithm.Framework.Risk import TrailingStopRiskManagementModel
from Execution.StandardDeviationExecutionModel import StandardDeviationExecutionModel
from QuantConnect.Securities.Option import OptionPriceModels
from datetime import timedelta
import decimal as d
from my_calendar import last_trading_day


'''
implements selling a 1 month call above price and using it to fund a put below price
'''

class  MomentumAlgorithm(QCAlgorithm):
    
    def Initialize(self):
        self.VERBOSEMODE = True
        self.VERBOSEMODE2 = False          
        self.SetStartDate(2019,5,10)
        self.SetEndDate(2019,5,12)
        self.SetCash(100000)
        self.benchmark = "SPY" 
        self.SetRiskManagement(TrailingStopRiskManagementModel(0.0172))
        self.symbols = ["SPY","TLT","GLD"]
        self.lastSymbol = None
        self.MAX_EXPIRY =  30 # max num of days to expiration => for uni selection
        self.MIN_EXPIRY = 3 # max num of days to expiration => for uni selection        
        self._no_K = 6       # no of strikes around ATM => for uni selection
        self.resol = Resolution.Minute   # Resolution.Minute .Hour  .Daily
        self.previous_delta, self.delta_threshhold = d.Decimal(0.0), d.Decimal(0.05) 
        self.Lev = d.Decimal(0.5)
        self._assignedOption = False   

        # add in globally accessible dictionaries that are keyed on symbol
        self.allOptions = {} 
        self.allEquities = {}
        self.allEquitySymbols= {}        
        self.allOptionSymbols = {} 
        self.putSymbols = {}   
        self.callSymbols = {}  
        self.orderCallIds= {} 
        self.orderPutIds= {} 
        self.last_call_days =  {}   
        self.last_put_days = {}   
        self.Deltas= {}
        self.atr15={}
        self.atrMultiplier = float(20) # factor by which to multiply average range to determine what is a "significant" move up or down
        self.stopPrices = {}  
        self.lastStandards= {}  
        self.protectedPositions = {}
        self.tradedContract = {}
                
        self.slice = None
        
        for s in self.symbols:
            #instantiate the equities and bind raw data to each needed for options populating the allEquities dictionary
            self.allEquities[s] = self.AddEquity(s,Resolution.Minute)
            self.allEquitySymbols[s] = self.allEquities[s].Symbol
            self.allEquities[s].SetDataNormalizationMode(DataNormalizationMode.Raw)
            # instantiate the options into a dictionary
            self.allOptions[s] = self.AddOption(s,Resolution.Minute)
            # get the options symbol as a handle into a dictionary
            self.allOptionSymbols[s] = self.allOptions[s].Symbol
            # set our strike/expiry filter for this option chain in list
            self.allOptions[s].SetFilter(self.UniverseFunc) # option.SetFilter(-2, +2, timedelta(0), timedelta(30))
            # for greeks and pricer (needs some warmup) - https://github.com/QuantConnect/Lean/blob/21cd972e99f70f007ce689bdaeeafe3cb4ea9c77/Common/Securities/Option/OptionPriceModels.cs#L81
            self.allOptions[s].PriceModel = OptionPriceModels.CrankNicolsonFD()  # both European & American, automatically
            self.callSymbols[s] = None
            self.putSymbols[s] = None
            self.Deltas[s] = None
            self.last_call_days[s] = None
            self.last_put_days[s] = None            
            self.stopPrices[s] = None 
            self.lastStandards[s] = None
            self.orderCallIds[s] = None
            self.orderPutIds[s] = None
            self.tradedContract[s] = False            
            self.protectedPositions[s] = False
            self.atr15[s] = self.ATR(s, 15, MovingAverageType.Simple,Resolution.Minute)        
            self.Securities[s].FeeModel = ConstantFeeModel(2.00)        
            

        # need to parse security from options in Keys to run history requests or it will bomb
        dict = self.parseKeysToTypes(self.Securities.Keys)
        self.optionKeys = dict["options"]
        self.securityKeys = dict["securities"]
            



        # this is needed for Greeks calcs
        self.SetWarmUp(TimeSpan.FromDays(15))   


      # add in execution model that only buys when price has dropped by 2 standard deviations
        #self.SetExecution(StandardDeviationExecutionModel(10, 2, Resolution.Daily))

        #schedule functions to run at repeated intervals
        #self.Schedule.On(self.DateRules.WeekStart("SPY"), self.TimeRules.AfterMarketOpen("SPY", 10), self.Rebalance)
        #self.Schedule.On(self.DateRules.EveryDay("SPY"),self.TimeRules.AfterMarketOpen("SPY", 10), self.Rebalance)
        
   
    def Rebalance(self):
        
        #self.EmitInsights(Insight.Price(self.allEquitySymbols["SPY"], timedelta(1), InsightDirection.Up))     
        if not self.Portfolio.Invested:
            self.SetHoldings(self.allEquitySymbols["SPY"],self.Lev)
            # need to get the options contracts  1 time
            self.get_contracts(self.slice)
            self.get_greeks(self.slice)        
            
            self.tradeOptions(self.allEquitySymbols["SPY"])            
        return

    def tradeOptions(self,symbol):
        breakout = False
        # make sure call and put options have already been defined 
        for c in self.callSymbols:
            if c is None: 
                breakout = True
        for c in self.putSymbols:
            if c is None: 
                breakout = True       
        if breakout: 
            self.Log("unable to setup options. Values have not yet been assigned")
            return
        symbolStr = symbol.Value
        if (self.callSymbols[symbolStr] is None) or (self.putSymbols[symbolStr] is None) : 
            self.Log("cannot trade options as none are defined yet")
            return
        qnty = self.calcOptionContracts(symbol)
        optionSymbol = self.callSymbols[symbolStr]
        #if not self.Portfolio[optionSymbol].Invested: 
        if not   self.tradedContract[symbolStr] : #self._assignedOption #self.Portfolio[optionSymbol].Invested: 
            self.Log("selling calls {} in tradeOptions " . format(optionSymbol))
            
            self.orderCallIds[symbolStr] = self.Sell(optionSymbol, qnty) 
            self.tradedContract[symbolStr] = True   
        
        optionSymbol = self.putSymbols[symbolStr]
        '''    
        if not self.Portfolio[optionSymbol].Invested: 
            self.Log("buying {} in tradeOptions " . format(optionSymbol))
            self.orderPutIds[symbolStr]  = self.Buy(optionSymbol, qnty)            
            
        '''
        
    
  
    def calcOptionContracts(self,symbol):
        symbolStr = symbol.Value
        quantity = float(self.Portfolio[symbolStr].Quantity)  
        optionContractCount = math.ceil(quantity / 100)
        if self.VERBOSEMODE: self.Log("contract count for symbol {} quantity {} is {} " . format(symbolStr,quantity, optionContractCount))
        return optionContractCount
        
  
        
    def closeOpenOptionContracts(self,symbol):
        symbolStr = symbol.Value

        if self.Securities.ContainsKey(self.callSymbols[symbolStr] ) and self.Portfolio[self.callSymbols[symbolStr] ].Invested:
            self.Log("closing {} in closeOptionsContracts " . format(self.callSymbols[symbolStr]))
            self.Liquidate(self.callSymbols[symbolStr] )    


        
    def OnData(self, slice):
        self.slice = slice
        if not self.Portfolio.Invested and self.HourMinuteIs(10, 1):
            self.Log("I should only be called every day at 10:01")
            self.SetHoldings(self.allEquitySymbols["SPY"],self.Lev)
            # need to get the options contracts  1 time
            self.get_contracts(self.slice)
            self.get_greeks(self.slice)        
            
            self.tradeOptions(self.allEquitySymbols["SPY"])                  

  

    
    def get_contracts(self, slice):
        """
        Get ATM call and put
        """
        for kvp in slice.OptionChains:
            for optionSymbol in self.allOptionSymbols:
                if not kvp.Key.Contains(optionSymbol): continue
    
                chain = kvp.Value   # option contracts for each 'subscribed' symbol/key 
    
                spot_price = chain.Underlying.Price
                #   self.Log("spot_price {}" .format(spot_price))
    
                # prefer to do in steps, rather than a nested sorted
                
                # 1. get furthest expiry            
                contracts_by_T = sorted(chain, key = lambda x: x.Expiry, reverse = True)
                if not contracts_by_T: return
                self.expiry = contracts_by_T[0].Expiry.date() # furthest expiry 
                self.last_call_days[optionSymbol] = last_trading_day(self.expiry)
 
                self.early_expiry = contracts_by_T[-1].Expiry.date() # earliest date from dates retrieved
                self.last_put_days[optionSymbol] = last_trading_day(self.early_expiry) 
 
                # get contracts with further expiry and sort them by strike
                slice_T = [i for i in chain if i.Expiry.date() == self.expiry]
                sorted_contracts = sorted(slice_T, key = lambda x: x.Strike, reverse = False)
            #   self.Log("Expiry used: {} and shortest {}" .format(self.expiry, contracts_by_T[0].Expiry.date()) )
    
                # 2a. get the ATM closest CALL to short
                calls = [i for i in sorted_contracts \
                         if i.Right == OptionRight.Call and i.Strike >= spot_price]
                self.callSymbols[optionSymbol] = calls[0].Symbol if calls else None
            #   self.Log("delta call {}, self.call type {}" .format(self.call.Greeks.Delta, type(self.call)))
            #   self.Log("implied vol {} " .format(self.call.ImpliedVolatility))

                # get contracts with further expiry and sort them by strike
                slice_T = [i for i in chain if i.Expiry.date() == self.early_expiry]
                sorted_contracts = sorted(slice_T, key = lambda x: x.Strike, reverse = False)
            #   self.Log("Expiry used: {} and shortest {}" .format(self.early_expiry, contracts_by_T[-1].Expiry.date()) )
    
        

                # 2b. get the ATM closest put to short
                puts = [i for i in sorted_contracts \
                         if i.Right == OptionRight.Put and i.Strike <= spot_price]
                self.putSymbols[optionSymbol] = puts[-1].Symbol if puts else None


    def get_greeks(self, slice):
        """
        Get greeks for invested option: self.call and self.put
        """
        breakout = False
        for c in self.callSymbols:
            if c is None: 
                breakout = True
        for c in self.putSymbols:
            if c is None: 
                breakout = True       
        if breakout: return

        
        for kvp in slice.OptionChains:
            for optionSymbol in self.allOptionSymbols:
                symbolStr = optionSymbol
                if not kvp.Key.Contains(optionSymbol): continue
                chain = kvp.Value   # option contracts for each 'subscribed' symbol/key 
                traded_contracts = filter(lambda x: x.Symbol == self.callSymbols[symbolStr] or 
                                             x.Symbol == self.putSymbols[symbolStr], chain)
                if not traded_contracts: self.Log("No traded cointracts"); return
            
                deltas = [i.Greeks.Delta for i in traded_contracts]
            #   self.Log("Delta: {}" .format(deltas))
                self.Deltas[symbolStr]=sum(deltas)
                # self.Log("Vega: " + str([i.Greeks.Vega for i in contracts]))
                # self.Log("Gamma: " + str([i.Greeks.Gamma for i in contracts]))

  
            
    def UniverseFunc(self, universe):
        return universe.IncludeWeeklys()\
                    .Strikes(-self._no_K, self._no_K)\
                    .Expiration(timedelta(self.MIN_EXPIRY), timedelta(self.MAX_EXPIRY))
    
    # ----------------------------------------------------------------------
    # Other ancillary fncts
    # ----------------------------------------------------------------------   
    def OnOrderEvent(self, orderEvent):
        order = self.Transactions.GetOrderById(orderEvent.OrderId)
        self.Log("{0} - {1}:TEST: {2}".format(self.Time, order.Type, orderEvent))        
        pass

    def OnAssignmentOrderEvent(self, assignmentEvent):
        self.Log("we got assigned" + str(assignmentEvent))
        self._assignedOption = True
    #   if self.isMarketOpen(self.equity_symbol):
    #       self.Liquidate(self.equity_symbol)
    
    def TimeIs(self, day, hour, minute):
        return self.Time.day == day and self.Time.hour == hour and self.Time.minute == minute
    
    def HourMinuteIs(self, hour, minute):
        return self.Time.hour == hour and self.Time.minute == minute
    

    def parseKeysToTypes(self,keys):
        '''
        # options must be fetched from history with difference mechanism see in research
        option_history = qb.GetOptionHistory(spy_option.Symbol, datetime(2017, 1, 11, 10, 10), datetime(2017, 1, 13, 12, 10))
        #print(dir(option_history.GetAllData()))
        this helper function allows you to separate out those types into sepparate arrays, stored in a dictionary.
        call with 
        d = self.parseKeysToTypes(self.Securities.Keys)
        you can then send following history call which will now work instead of bombing. but will not give you history of options
        optionKeys = d["options"]
        securityKeys = d["securities"]
        h1 = self.History(securityKeys, 20, Resolution.Minute)
        '''
        options = []
        securities = []
        forex  = []
        futures = []
        cfds = []
        cryptos = []
        unknowns  = []    
        for el in keys:
            if el.SecurityType == 1:
                securities.append(el)
            elif el.SecurityType == 2: 
                options.append(el)
            elif el.SecurityType == 4: 
                forex.append(el)
            elif el.SecurityType == 5: 
                futures.append(el)
            elif el.SecurityType == 6: 
                cfds.append(el)
            elif el.SecurityType == 7: 
                cryptos.append(el)
            else:
                unknowns.append(el)
        dict = {"securities":securities, "options":options,"forex":forex,"futures":futures,"cfds":cfds,"cryptos":cryptos,"unknowns":unknowns}
        return dict

# ------------------------------------------------------------------------------
# Business days
# ------------------------------------------------------------------------------
from datetime import timedelta #, date
from pandas.tseries.holiday import (AbstractHolidayCalendar,    # inherit from this to create your calendar
                                    Holiday, nearest_workday,   # to custom some holidays
                                    #
                                    USMartinLutherKingJr,       # already defined holidays
                                    USPresidentsDay,            # "     "   "   "   "   "
                                    GoodFriday,
                                    USMemorialDay,              # "     "   "   "   "   "
                                    USLaborDay,
                                    USThanksgivingDay           # "     "   "   "   "   "
                                    )


class USTradingCalendar(AbstractHolidayCalendar):
    rules = [
      Holiday('NewYearsDay', month=1, day=1, observance=nearest_workday),
      USMartinLutherKingJr,
      USPresidentsDay,
      GoodFriday,
      USMemorialDay,
      Holiday('USIndependenceDay', month=7, day=4, observance=nearest_workday),
      USLaborDay,
      USThanksgivingDay,
      Holiday('Christmas', month=12, day=25, observance=nearest_workday)
    ]

# TODO: to be tested
def last_trading_day(expiry):
    # American options cease trading on the third Friday, at the close of business 
    # - Weekly options expire the same day as their last trading day, which will usually be a Friday (PM-settled), [or Mondays? & Wednesdays?]
    # 
    # SPX cash index options (and other cash index options) expire on the Saturday following the third Friday of the expiration month. 
    # However, the last trading day is the Thursday before that third Friday. Settlement price Friday morning opening (AM-settled).
    # http://www.daytradingbias.com/?p=84847
    
    dd = expiry     # option.ID.Date.date()
    
    # if expiry on a Saturday (standard options), then last trading day is 1d earlier 
    if dd.weekday() == 5:
        dd -= timedelta(days=1)   # dd -= 1 * BDay()
        
    # check that Friday is not an holiday (e.g. Good Friday) and loop back
    while USTradingCalendar().holidays(dd, dd).tolist():    # if list empty (dd is not an holiday) -> False
        dd -= timedelta(days=1) 
        
    return dd