from QuantConnect.Securities.Option import OptionPriceModels
from datetime import timedelta
import pandas as pd
import numpy as np

class ImpliedVolatilityToDataFrame(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2017, 7, 1)
        self.SetEndDate(2017, 8, 25)
        self.SetCash(100000)
        
        self.symbols = []
        
        # Symbols in Dow Jones Industrial Average (DJIA)
        self.tickers = ["MMM", "AXP", "AAPL", "BA", "CAT", "CVX", "CSCO","KO",
                    "DIS","DD","XOM","GE","GS","HD","IBM","INTC","JPM","MCD",
                    "MRK","MSFT","NKE","PFE","PG","TRV","UTX","UNH","VZ","V","WMT"]    
                    
        self.IV = pd.DataFrame(columns=self.tickers)
        nRow = pd.Series(np.repeat(np.nan,len(self.tickers)),index=self.tickers)
        self.IV=self.IV.append(nRow,ignore_index=True)

        # Add options and set pricing model
        for ticker in self.tickers:
            option = self.AddOption(ticker,Resolution.Minute)
            self.symbols.append(option.Symbol)
            option.SetFilter(-1, 1, timedelta(10), timedelta(30))
            option.PriceModel = OptionPriceModels.CrankNicolsonFD()

        self.SetWarmUp(TimeSpan.FromDays(7))
        
        # Set iteration variable
        self.m = 0

    def OnData(self,slice):
        if self.Portfolio.Invested: return
        for symbol in self.symbols:
            if self.IsMarketOpen(symbol):
                # Loop through the option chain at the end of day
                if self.Time.hour == 15 and self.Time.minute == 59:
                    for kvp in slice.OptionChains:
                        if kvp.Key == symbol:
                            chain = kvp.Value
                            call = [i for i in chain if i.Right == OptionRight.Call]
                            if len(call) == 0: continue
                        
                            # Retrieve the call ATM option contract, at the furthest date within our filtered universe of options
                            contracts =  sorted(sorted(call, key=lambda x: x.Expiry, reverse=False), key=lambda x: abs(x.UnderlyingLastPrice - x.Strike))
                            atm_contract = contracts[0]
                            strSymbol=symbol.Value[1:len(symbol.Value)]
                            
                            # Populate the DataFrame IV, which includes implied volatility at the end of day for each stock
                            self.IV[strSymbol].iloc[self.m]=atm_contract.ImpliedVolatility
                            
                            ## Add new row for next iteration
                            if strSymbol==self.tickers[len(self.tickers)-1]:
                                # Retrieve implied volatility for all stocks, for seven last observations and print to log file.
                                if self.m >= 7:
                                    self.Log(str(self.IV.tail(7)))
                                    # Print Strike, Implied Volatility and Greeks to log file
                                    self.logAttributes(atm_contract)
                                    
                                self.m=self.m+1
                                nRow = pd.Series(np.repeat(np.nan,len(self.tickers)),index=self.tickers)
                                self.IV=self.IV.append(nRow,ignore_index=True)
                                


                            
    def logAttributes(self,contract):
        self.Log("Symbol: " + str(contract.Symbol) + " - Strike: " + str(contract.Strike))
        self.Log("Symbol: " + str(contract.Symbol) + " - Delta: " +  str(contract.Greeks.Delta))
        self.Log("Symbol: " + str(contract.Symbol) + " - Vega: " +  str(contract.Greeks.Vega))
        self.Log("Symbol: " + str(contract.Symbol) + " - Gamma: " +  str(contract.Greeks.Gamma))