import pandas as pd
from functools import partial
from QuantConnect.Securities.Option import OptionPriceModels

class ParticleCalibratedCoil(QCAlgorithm):

    def Initialize(self):
        
        '''
            Parameters for adjusting
        '''
        self.numberOfLiquidStocks = 100 # Controls the number of stocks in play
        
        
        '''
            Backtesting variables
        '''
        self.SetStartDate(2018, 1, 1)
        self.SetEndDate(2020, 1, 1)
        self.SetCash(1000000)
        
        '''
            Algorithm variables
        '''
        self.UniverseSettings.Resolution = Resolution.Minute
        self.AddUniverse(self.CoarseSelectionFilter)
        self.UniverseSettings.DataNormalizationMode = DataNormalizationMode.Raw
        
        self.indicators = {}
        self.rankingOfAbsoluteIV = []
        
        self.dayCounter = 0
        self.Schedule.On(self.DateRules.EveryDay(), \
                         self.TimeRules.At(9, 40), \
                         self.EveryDayAfterMarketOpen)
        
    def CoarseSelectionFilter(self, coarse):
        '''
            1. Sorts each element of the coarse object by dollar volume
            2. Returns a list of coarse object, limited to top 100 highest volume
            3. Returns a list of symbols we want to initialise into the algorithm
        '''
        
        self.sortedByDollarVolume = sorted(coarse, key=lambda c: c.DollarVolume, reverse=True)
        self.topHundredMostLiquid = self.sortedByDollarVolume[:self.numberOfLiquidStocks]
            
        return [stock.Symbol for stock in self.topHundredMostLiquid]
        
    def OnSecuritiesChanged (self,changes):
        
        for underlying in changes.AddedSecurities:
            if underlying.Symbol.SecurityType != SecurityType.Equity: continue
            option = self.AddOption(underlying.Symbol.Value, Resolution.Minute)
            option.SetFilter(-5, +2, timedelta(30), timedelta(60))
            option.PriceModel = OptionPriceModels.CrankNicolsonFD()
            
            
            if not underlying.Symbol.Value in self.indicators:
                self.indicators[underlying.Symbol.Value] = {"Volatility": self.STD(underlying.Symbol.Value, 240, Resolution.Daily)}
                self.indicators[underlying.Symbol.Value]["CloseWindow"] = self.SMA(underlying.Symbol.Value, 1, Resolution.Daily)
                
                # Warm up STD indicator
                history = self.History([underlying.Symbol], 240, Resolution.Daily).loc[underlying.Symbol]
                for idx, row in history.iterrows():
                    self.indicators[underlying.Symbol.Value]['Volatility'].Update(idx, row['close'])
    
        for underlying in changes.RemovedSecurities:
            self.RemoveSecurity(underlying.Symbol)
            for symbol in self.Securities.Keys:
                if symbol.SecurityType == SecurityType.Option and symbol.Underlying == underlying.Symbol:
                    self.RemoveSecurity(symbol)
        
        
    def OnData(self, slice):

        for chain in slice.OptionChains.Values:
            
            # Filter for the first ATM contract
            if chain.Contracts.Count < 1:
                continue
            
            atmContract = sorted(chain, key = lambda x: abs(x.UnderlyingLastPrice - x.Strike))[0]
            
            impliedVolatility = atmContract.ImpliedVolatility
            underlyingSymbol = atmContract.UnderlyingSymbol.Value
            if self.indicators[underlyingSymbol]["Volatility"].IsReady:
                self.Debug((underlyingSymbol, self.indicators[underlyingSymbol]["Volatility"].current.Value))
        
    def EveryDayAfterMarketOpen(self):
        self.dayCounter += 1
        self.Debug(f"This is day {self.dayCounter} at {self.Time}")