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.UniverseSettings.Resolution = Resolution.Hour
        self.UniverseSettings.DataNormalizationMode = DataNormalizationMode.Raw
        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 CustomSecurityInitializer(self, security):
        security.SetDataNormalizationMode(DataNormalizationMode.Raw)

    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)
            for d in raw_er['trade_date'].unique():
                i = str(d)
                self.backtestSymbolsPerDay[i] = (raw_er[raw_er['trade_date']== d]).symbol.tolist()

        index = self.Time.strftime("%Y%m%d")

        try:
            self.current_universe = self.backtestSymbolsPerDay[index]
        except:
            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):
        self.Liquidate()

    def OnData(self, slice):
        if self.Time.hour == 13 and self.Time.minute == 1:
            for security in self.ActiveSecurities.Keys:
                self.Log(f'{self.Time} >> Security in Universe: {security.Value}')

    
    
        
    def OnSecuritiesChanged(self, changes):
        self.changes = changes
        if self.x == 1: 

            for x in changes.AddedSecurities:
                if x.Symbol.Value == 'SPY': continue
                if x.Symbol.SecurityType != SecurityType.Equity: continue
                
                
                option = self.AddOption(x.Symbol.Value, Resolution.Minute)
                option.PriceModel = OptionPriceModels.BinomialCoxRossRubinstein()
                option.SetFilter(self.UniverseFunc)
                self.symbols.append(option.Symbol)
                self.Log(f'{self.Time} >> {x.Symbol} and Options added to Universe')

            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.Log(f'{self.Time} >> {x.Symbol} and {symbol} removed from the Universe')
        self.x = 0 
        

    def UniverseFunc(self, universe):
        return universe.IncludeWeeklys().Strikes(-1, 1).Expiration(timedelta(2), timedelta(8))