import decimal as d
import numpy as np
import pandas as pd
import math
import datetime
import json

class DropboxBaseDataUniverseSelectionAlgorithm(QCAlgorithm):

    def Initialize(self):

        self.SetStartDate(2019,1,8)
        self.SetEndDate(2019,1,15)
        self.SetCash(100000)
        
        # add equity tickers to the universe (before the market open of each trading day)
        self.UniverseSettings.Resolution = Resolution.Minute;
        self.AddUniverse(StockDataSource, "my-stock-data-source", self.stockDataSource)
        self.UniverseSettings.DataNormalizationMode = DataNormalizationMode.Raw
        
        # set schedule to liquidate at 10 minutes prior to the market close of each trading day
        spy = self.AddEquity("SPY", Resolution.Minute)
        spy.SetDataNormalizationMode(DataNormalizationMode.Raw)
        self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.BeforeMarketClose("SPY", 10), self.EveryDayBeforeMarketClose)

    def stockDataSource(self, data): # This will grab for each date the different tickers in the csv and add them to the universe
        list = []
        for item in data:
            for symbol in item["Symbols"]:
                list.append(symbol)
        #self.Debug(str(self.Time))
        #self.Debug(str(list))
        return list

    def EveryDayBeforeMarketClose(self):
        #self.Debug("############## Closing Position " + str(self.Time.date()) + " " + str(self.Time) + "############## ")
        self.Liquidate()
        
        for equity in self.Portfolio:
            self.RemoveSecurity(equity.Key)
        
        spy = self.AddEquity("SPY", Resolution.Minute)
        spy.SetDataNormalizationMode(DataNormalizationMode.Raw)
        
        self.Debug("Positions closed")
        
        
    def OnData(self, data):
        option_invested = [x.Key for x in self.Portfolio if x.Value.Invested and x.Value.Type==SecurityType.Option]
        
        if (self.Time.time() < datetime.time(9, 40, 0)) & (self.Time.time() > datetime.time(9, 31, 0)): # to avoid OnData to trade again just before the end of the day (after the liquidation)
            for symbol in data.Keys:
                if symbol.SecurityType == 1: #Selecting only stocks
    
                    invested = [option for option in option_invested if option.Underlying == symbol]
                    if len(invested) > 0: 
                        self.Debug("Already invested in "+ str(symbol))
                        return 
                    
                    if symbol.Value == "SPY":
                        return
                    
                    self.Log(self.Time.time())
                    self.Log(symbol)
                    #self.SetHoldings(key, 0.1)            
                    stk = self.AddEquity(symbol.Value, Resolution.Minute)
                    stk.SetDataNormalizationMode(DataNormalizationMode.Raw)
                    
                    contracts = self.OptionChainProvider.GetOptionContractList(symbol, self.Time.date()) # Get list of strikes and expiries
                    self.TradeOptions(contracts, symbol.Value) # Select the right strikes/expiries and trade
            
    def TradeOptions(self, contracts, ticker):
        # run CoarseSelection method and get a list of contracts expire within 15 days from now on
        # and the strike price between rank -1 to rank 1, rank being the step of the contract
        
        filtered_contracts = self.CoarseSelection(ticker, contracts, -1, 1, 0, 15) # set min_expiry as 1 would avoid trading the contract that expires on the same day
        
        if len(filtered_contracts) >0:
                
            expiry = sorted(filtered_contracts,key = lambda x: x.ID.Date, reverse=False)[0].ID.Date # Take the closest expiry
            
            # filter the call options from the contracts expire on that date
            call = [i for i in filtered_contracts if i.ID.Date == expiry and i.ID.OptionRight == 0]
            # sorted the contracts according to their strike prices
            call_contracts = sorted(call,key = lambda x: x.ID.StrikePrice)
            self.call = call_contracts[0]
            
            for i in filtered_contracts:
                if i.ID.Date == expiry and i.ID.OptionRight == 1 and i.ID.StrikePrice ==call_contracts[0].ID.StrikePrice:
                    self.put = i
    
            ''' Before trading the specific contract, you need to add this option contract
                AddOptionContract starts a subscription for the requested contract symbol '''
        
            # self.call is the symbol of a contract 
            self.AddOptionContract(self.call, Resolution.Minute)
            self.AddOptionContract(self.put, Resolution.Minute)
            
            self.SetHoldings(self.call.Value, -0.01)
            self.SetHoldings(self.put.Value, -0.01)

            # Some Logging  
            #self.Debug("Strike Price : "+str(self.call.ID.StrikePrice))
            #self.Debug("Expiry : "+str(self.call.ID.Date))
            self.Debug("Call Mid-Point : "+str(self.Securities[self.call].Price))
            
            #self.Debug("IV : "+str(self.call.ImpliedVolatility))
            
        else:
            pass
        
        
    def CoarseSelection(self, underlyingsymbol, symbol_list, min_strike_rank, max_strike_rank, min_expiry, max_expiry):

        ''' This method implements the coarse selection of option contracts
            according to the range of strike price and the expiration date,
            this function will help you better choose the options of different moneyness '''
        
        # filter the contracts based on the expiry range
        contract_list = [i for i in symbol_list if min_expiry <= (i.ID.Date.date() - self.Time.date()).days < max_expiry]
        #self.Log("Ticker Und : " + str(underlyingsymbol))
        #self.Log("Nb of contract found : " + str(len(contract_list)))
        #self.Log("Underlying price : "+str(self.Securities[underlyingsymbol].Price))
        
        # find the strike price of ATM option
        # It seems like sometimes OptionChainProvider.GetOptionContractList is bugging and returns nothing, so let's try/except
        try:
            atm_strike = sorted(contract_list,
                                key = lambda x: abs(x.ID.StrikePrice - self.Securities[underlyingsymbol].Price))[0].ID.StrikePrice
            strike_list = sorted(set([i.ID.StrikePrice for i in contract_list]))
            
            # find the index of ATM strike in the sorted strike list
            atm_strike_rank = strike_list.index(atm_strike)
            try: 
                min_strike = strike_list[atm_strike_rank + min_strike_rank]
                max_strike = strike_list[atm_strike_rank + max_strike_rank]
            except:
                min_strike = strike_list[0]
                max_strike = strike_list[-1]
            # filter the contracts based on the range of the strike price rank
            filtered_contracts = [i for i in contract_list if i.ID.StrikePrice >= min_strike and i.ID.StrikePrice <= max_strike]
        except:
            self.Debug("Problem")

        return filtered_contracts 


class StockDataSource(PythonData):
    
    def GetSource(self, config, date, isLiveMode):
        url = "https://www.dropbox.com/s/2az14r5xbx4w5j6/daily-stock-picker-live.csv?dl=1" if isLiveMode else \
            "https://www.dropbox.com/s/ofzgxsp2b27pkri/quantconnect_triggers.csv?dl=1"

        return SubscriptionDataSource(url, SubscriptionTransportMedium.RemoteFile)
    
    def Reader(self, config, line, date, isLiveMode):
        #if not (line.strip() and line[0].isdigit()): return None
        
        stocks = StockDataSource()
        stocks.Symbol = config.Symbol
        
        csv = line.rstrip(',').split(',') # rstrip is essential because quantconnect throws an empty element error (extra commas at the end of the csv)
        if isLiveMode:
            stocks.Time = date
            stocks["Symbols"] = csv
        else:
            stocks.Time = datetime.datetime.strptime(csv[0], "%Y%m%d")
            stocks["Symbols"] = csv[1:]
        return stocks