from QuantConnect.Securities.Option import OptionPriceModels
from datetime import timedelta
import decimal as d
#import pandas as pd

class PutMarginCalc(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2017, 1, 6)
        self.SetEndDate(2017, 1, 24)
        self.SetCash(1000000)
        
        self.Log("PERIOD: JAN 2017")
        
        # ----------------------------------------------------------------------
        # Algo params
        # ----------------------------------------------------------------------
        self.MIN_EXPIRY = 40 # min num of days to expiration => for uni selection
        self.MAX_EXPIRY = 55 # max num of days to expiration => for uni selection
        self.MIN_DELTA = .21 #lower delta = lower risk = less returns
        self.MAX_DELTA = .40 #higher delta = higher risk = more returns
        self._no_K = 25 # no of strikes around ATM => for uni selection
        self.resol = Resolution.Minute   # Resolution.Minute #Options only supports minute currently
        self.tkr = "SPY"
        self.Lev = d.Decimal(1.0)
        self.single_trade_portfolio = .2 #max portfolio value to place one single trade
        self.single_trade_margin = .5 #max margin value to place one single trade
        # add underlying Equity 
        equity = self.AddEquity(self.tkr, self.resol)  
        equity.SetDataNormalizationMode(DataNormalizationMode.Raw) # IMPORTANT: default
        self.equity_symbol = equity.Symbol
        
        # Add options
        option = self.AddOption(self.tkr, self.resol)
        option.SetDataNormalizationMode(DataNormalizationMode.Raw) # IMPORTANT: default
        self.option_symbol = option.Symbol

        # filter options
        option.SetFilter(self.UniverseFunc) # option.SetFilter(-2, +2, timedelta(0), timedelta(30))

        #seed greeks
        option.PriceModel = OptionPriceModels.CrankNicolsonFD()  # both European & American, automatically
        self.SetWarmUp(TimeSpan.FromDays(7))    # timedelta(7)
        
        self.call, self.put = None, None
   
            
    def OnData(self, slice):
        if self.IsWarmingUp: return
        if not self.HourMinuteIs(10, 1): return ##trade once a day max
        
        
        self._assignedOption = False

        if not self.Portfolio.Invested or self.Portfolio.GetBuyingPower(self.tkr) > self.Portfolio.TotalPortfolioValue * self.single_trade_portfolio:
            
            # select contract
            #self.Log("get contracts")
            self.get_contracts(slice)
            if (not self.put): return
    
            #unit_price = self.put.UnderlyingLastPrice * d.Decimal(100.0)
            #unit_price =  self.Securities[self.equity_symbol].Price * d.Decimal(100.0) 

            unit_price = self.put.Strike * d.Decimal(100.0)
            qty = int((self.Portfolio.GetBuyingPower(self.tkr) * self.single_trade_margin) / unit_price)
            bp = str(self.Portfolio.GetBuyingPower(self.tkr))
            tv = str(self.Portfolio.TotalPortfolioValue)
            mr = str(self.Portfolio.MarginRemaining)
            mu = str(self.Portfolio.TotalMarginUsed)
            cash = str(self.Portfolio.Cash) #settled only
            u_cash = str(self.Portfolio.UnsettledCash) #settled only
            
            #blah = str(self.CalculateOrderQuantity(self.tkr, .2)) #Calculates for buying shares, not contracts

           
            self.Debug("Starting BP: " + bp + " Margin Remaining: " + mr + " Margin Used: " + mu)
            self.Debug("Selling puts: " + str(qty) + "@" + str(self.put.Strike))
            if self.put is not None:  
                order = self.MarketOrder(self.put.Symbol, -qty)
                
            bp = str(self.Portfolio.GetBuyingPower(self.tkr))
            mr = str(self.Portfolio.MarginRemaining)
            mu = str(self.Portfolio.TotalMarginUsed)
            self.Debug("Ending BP: " + bp + " Margin Remaining: " + mr + " Margin Used: " + mu)



    def get_contracts(self, slice):
        
        for kvp in slice.OptionChains:
            if kvp.Key != self.option_symbol: continue
            optionchain = kvp.Value   # option contracts for each 'subscribed' symbol/key 
            puts = [x for x in optionchain if x.Right == 1]
            contracts_by_T = sorted(puts, key = lambda x: x.Expiry, reverse = True)
            if not contracts_by_T: return
            self.expiry = contracts_by_T[0].Expiry.date() # furthest expiry 
            slice_T = [i for i in puts if i.Expiry.date() == self.expiry]
            sorted_contracts = sorted(slice_T, key = lambda x: x.Strike, reverse = False)

            self.put = puts[0] if puts else None 
            #self.Log("found contract: " + str(self.put))


       
    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):
    #   self.Log("Order Event -> {}" .format(orderEvent))
        pass

    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