from clr import AddReference
AddReference("System")
AddReference("QuantConnect.Algorithm")
AddReference("QuantConnect.Common")
from decimal import Decimal
from System import *
from QuantConnect import *
from QuantConnect.Algorithm import QCAlgorithm
from QuantConnect.Data.UniverseSelection import *
from QuantConnect.Securities.Option import OptionPriceModels
import base64
import numpy as np
import pandas as pd
from datetime import timedelta, datetime, date

class EarningsAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2018,1,1)
        self.SetEndDate(2018,1,26)
        self.SetCash(2000000)
        self.backtestSymbolsPerDay = {}
        self.current_universe = []
        self.symbols = []
        self.syms = []
        # Set the security initializer with the characteristics defined in CustomSecurityInitializer
        self.SetSecurityInitializer(self.CustomSecurityInitializer)
        self.UniverseSettings.Resolution = Resolution.Hour
        #self.AddUniverse("my-dropbox-universe",Resolution.Daily,  self.selector)
      
        # time set here realises at 2pm US
        self.SetUniverseSelection(ScheduledUniverseSelectionModel(self.DateRules.Every(DayOfWeek.Monday, DayOfWeek.Tuesday, 
        DayOfWeek.Wednesday, DayOfWeek.Thursday, DayOfWeek.Friday),self.TimeRules.At(9, 00, 00), self.selector ))
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, 
                                AccountType.Margin)

        self.changes = None                   
        bench = self.AddEquity("SPY", Resolution.Minute)
        self.SetBenchmark(bench.Symbol)                        
        self.x = 0
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Monday, DayOfWeek.Tuesday, 
        DayOfWeek.Wednesday, DayOfWeek.Thursday, DayOfWeek.Friday), self.TimeRules.At(11, 00, 00), Action(self.reb))

    
    def selector(self, datetime):
        self.current_universe = []
        if len(self.backtestSymbolsPerDay) == 0:
            url = "https://www.dropbox.com/s/u5s56xlhk5ukcyo/er.csv?dl=1"
            raw_er = pd.read_csv(url)
            self.Debug(str(raw_er.head()))
            for d in raw_er['trade_date'].unique():
                i = str(d)
                self.backtestSymbolsPerDay[i] = (raw_er[raw_er['trade_date']== d]).symbol.tolist()
            self.Debug(self.backtestSymbolsPerDay)
            
        index = self.Time.strftime("%Y%m%d")
        self.Debug(index)
        
        try:
            self.current_universe = self.backtestSymbolsPerDay[index]
        except:
            self.current_universe = []

        self.Debug(self.Time) 
        self.Debug(self.current_universe)
        #
        equities = []
        self.symbols = []
        for ticker in self.current_universe:
            equities.append(Symbol.Create(ticker, SecurityType.Equity, Market.USA))
        
        #bench = self.AddEquity("SPY", Resolution.Minute)
        #self.SetBenchmark(bench.Symbol)  

        # set a marker to only run after daily universe creation, otherwise it runs at midnight too and liquidates on open.
        self.x = 1 
        return equities
        
    def reb(self):
        #sells all trades.
        self.Debug(str(self.Time) +'  '+'liquidating')
        self.Liquidate()

    def OnData(self, slice):
        if slice.Bars.Count == 0: return
        if self.changes is None: return
        if self.Portfolio.Invested: return
        if self.Time.hour < 15: return
        for security in self.ActiveSecurities.Keys:
            self.Log(f'Security in Universe: {security.Value}')
            
        self.option_data = slice
        self.TradeOptions(slice) 
        #self.TradeEquities(slice)

        self.Debug(str(self.Time) +'  '+'end OnData')
    
    def TradeOptions(self,slice):
        self.Debug('trading options')
        for symbol in self.symbols:
            if slice.OptionChains.Count == 0: return

            for i in slice.OptionChains:
                self.Debug('now chains')
                if i.Key == symbol:

                    chains = i.Value
                    # divide option chains into call and put options 
                    self.calls = list(filter(lambda x: x.Right == OptionRight.Call, chains))
                    self.puts = list(filter(lambda x: x.Right == OptionRight.Put, chains))
                    
                    # if lists are empty return
                    if not self.calls or not self.puts: return
                    
                    #expries should be same for put and calls so collecting only puts for now. (convinience)
                    expiries = [i.Expiry for i in self.puts]
                    
                    # determine expiration date
                    self.expiry = min(expiries, key=lambda x: abs((x.date()-self.Time.date()).days-90))
                    #self.Debug(str(expiries))
                    #self.Debug("we choose this expiry...     " + str(self.expiry))
                    #self.Debug(symbol)
                    #self.symbol = epic
                   
                    self.straddle()
                    self.trade_execution()
        return
    
    
    def straddle(self):
        #strikes
        put_strikes = [i.Strike for i in self.puts]
        call_strikes = [i.Strike for i in self.calls]
        underlying_price = self.calls[0].UnderlyingLastPrice 
        
        # determine at-the-money strike. Again doesnt matter call or put using put.
        strike = min(put_strikes, key=lambda x: abs(x-underlying_price))
       
        # Set up the legs.
        self.atm_call = [i for i in self.calls if i.Expiry == self.expiry and i.Strike == strike]
        self.atm_put = [i for i in self.puts if i.Expiry == self.expiry and i.Strike == strike]    
        #self.Debug(' ...underlying is...   ' + 
        #'  ...put strike...' + str(self.atm_put[0].Strike) +
        #'  ...put delta...' + str(self.atm_put[0].Greeks.Delta)+
        #'  ...put atm expiry...' + str(self.atm_put[0].Expiry) +
        #'  ...call atm strike...' + str(self.atm_call[0].Strike)+
        #'  ...call atm delta...' + str(self.atm_call[0].Greeks.Delta)+
        #'  ...call atm expiry...' + str(self.atm_call[0].Expiry))
    
    
    def trade_execution(self):
        #self.Debug("lets check the legs exist and are sensible")
        if self.atm_put and self.atm_call:
            
            #if self.atm_put[0].Greeks.Delta < -0.55 or self.atm_put[0].Greeks.Delta > -0.40:return
            #if self.atm_call[0].Greeks.Delta < 0.40 or self.atm_call[0].Greeks.Delta > 0.55:return
            if self.atm_put[0].AskPrice < 0.30 :return
            if self.atm_call[0].AskPrice < 0.30:return
            
            trade_cost = self.atm_put[0].AskPrice +self.atm_call[0].AskPrice
            #self.Debug('Cost is :')
            #self.Debug(trade_cost)
            if trade_cost == 0: return

            #set the size
            #self.Debug('  ...atm strike...' + str(self.atm_put[0].AskPrice)+
            #'  ...atm strike...' + str(self.atm_call[0].AskPrice))
            size = int(1000 / (trade_cost * 100))
            #self.Portfolio.TotalPortfolioValue 
            #deltas = 20000 / (self.Securities[self.symbol].Price * 100)
            #self.Debug(str(deltas))
            #size = int(deltas)
            
            self.Sell(self.atm_call[0].Symbol, size)
            self.Sell(self.atm_put[0].Symbol, size)

        else:
            # no options so try tomorrow
            #self.Debug(str(self.Time) + '   ...No optons...  ' )
            return
    
    
        
    def OnSecuritiesChanged(self, changes):
        self.changes = changes
        self.Log("{}: {}".format(self.Time, changes))
        if self.x == 1: 

            for x in changes.AddedSecurities:
                
                if x.Symbol.Value == 'SPY': continue
                if x.Symbol.SecurityType != SecurityType.Equity: continue
                x.SetDataNormalizationMode(DataNormalizationMode.Raw)
                
                option = self.AddOption(x.Symbol.Value, Resolution.Minute)
                option.PriceModel = OptionPriceModels.BinomialCoxRossRubinstein()
                option.SetFilter(self.UniverseFunc)
                self.symbols.append(option.Symbol)

            for x in changes.RemovedSecurities:
                    if x.Symbol.Value == 'SPY': continue   
                    
                    for symbol in self.Securities.Keys:
                        if symbol.SecurityType == SecurityType.Option and symbol.Underlying == x.Symbol:
                            self.RemoveSecurity(symbol)
                            self.RemoveSecurity(x.Symbol)
        self.x = 0 
        

    def UniverseFunc(self, universe):
        #the option chain
        return universe.IncludeWeeklys().Strikes(-1, 1).Expiration(timedelta(2), timedelta(8))
        
    def CustomSecurityInitializer(self, security):
        '''Initialize the security with raw prices and zero fees 
        Args:
            security: Security which characteristics we want to change'''
        security.SetDataNormalizationMode(DataNormalizationMode.Raw)
        #security.SetFeeModel(ConstantFeeModel(0))
        
    def OnEndOfAlgorithm(self):
        self.Log('End captured')