import numpy as np
from QuantConnect.Securities.Option import OptionPriceModels ## Import Options Pricing Models
from datetime import datetime, timedelta

class LocalVolatilityAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2018, 6, 1)  # Set Start Date
        self.SetEndDate(2018, 9, 1)
        self.SetCash(100000)  # Set Strategy Cash
        
        option = self.AddOption("SPY");
        option.SetFilter(-2, 2, timedelta(0), timedelta(30))
        option.PriceModel = OptionPriceModels.CrankNicolsonFD()
        
        self.localVolatility = {}

    def OnData(self, slice):
        for chain in slice.OptionChains.Values:
            contracts = sorted(sorted(chain, \
                               key = lambda x: abs(chain.Underlying.Price - x.Strike)), \
                               key = lambda x: x.Expiry, reverse=True)
            for contract in contracts:
                self.Log(str(contract.Symbol))
                ## Calculate Local Volatility and store in dictionary, indexed by contract Symbol
                self.localVolatility[contract.Symbol] = self.CalculateLocalVolatility(contract)
                self.Log(f'Local Volatility: {self.localVolatility[contract.Symbol]}')
                

    def CalculateLocalVolatility(self, contract):
        ## Access relevant greeks and strike price
        delta = contract.Greeks.Delta
        strike = contract.Strike
        gamma = contract.Greeks.Gamma
        
        ## Calculate LV
        if (delta == 0) or (gamma == 0):
            localVolatility = 0.0
        else:
            localVolatility = np.sqrt( delta / ( 0.5*(strike**2)*gamma ) )
        return localVolatility