from Execution.ImmediateExecutionModel import ImmediateExecutionModel
from Portfolio.EqualWeightingPortfolioConstructionModel import EqualWeightingPortfolioConstructionModel
from QuantConnect.Securities.Option import OptionPriceModels
from QuantConnect.Securities.Option import IOptionPriceModel 
import optionsalpha

class OptionDataGenerator(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 8, 28)  # Set Start Date
        self.SetEndDate(2020,9,18)
        self.SetCash(100000)  # Set Strategy Cash
        self.SetTimeZone(TimeZones.Chicago)
        self.SetSecurityInitializer(self.CustomSecurityInitializer)
        
        self.variable=False
        self.endOfDay=False
         
        self.SetExecution(NullExecutionModel())
        self.SetPortfolioConstruction(NullPortfolioConstructionModel())
        self.SetRiskManagement(NullRiskManagementModel())
        
        self.AddAlpha(optionsalpha.alpha(self))

        symbols = [ Symbol.Create("SPY", SecurityType.Equity, Market.USA) ]
        self.SetUniverseSelection( ManualUniverseSelectionModel(symbols) )
        self.UniverseSettings.Resolution = Resolution.Minute #Minute resolution for options
        self.UniverseSettings.FillForward = True #Fill in empty data will next price
        self.UniverseSettings.ExtendedMarketHours = False #Does not takes in account after hours data
        self.UniverseSettings.MinimumTimeInUniverse = 1 # each equity has to spend at least 1 hour in universe
        self.UniverseSettings.Leverage=2 #Set's 2 times leverage
        self.Settings.FreePortfolioValuePercentage = 0.5
  
      
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.At(14,59), self.checktwo)
        
        
        if self.ObjectStore.ContainsKey("MyObject"):
            self.ObjectStore.Delete("MyObject")
            
    def check(self):
        self.variable=True
    def checktwo(self):
        self.endOfDay=True
        
    def OnData(self, slice):
        if self.IsWarmingUp: return
    
    def CustomSecurityInitializer(self, security):
        #Intialize each security's prices 
        security.SetDataNormalizationMode(DataNormalizationMode.Raw)
        
        if security.Type == SecurityType.Equity:
            #if equity intialize volatility model and perform history call
            security.VolatilityModel = StandardDeviationOfReturnsVolatilityModel(2)
            history = self.History(security.Symbol, 3, Resolution.Daily)
            if history.empty or 'close' not in history.columns:
                return
            
            for time, row in history.loc[security.Symbol].iterrows():
                trade_bar = TradeBar(time, security.Symbol, row.open, row.high, row.low, row.close, row.volume)    
                security.VolatilityModel.Update(security, trade_bar)
        
        elif security.Type == SecurityType.Option:
            #Intialize Option Pricing model here 
            security.PriceModel = OptionPriceModels.BaroneAdesiWhaley()

from datetime import datetime, timedelta
import json

class alpha(AlphaModel):   
    def __init__(self,algorithm):
        self.algorithm=algorithm
        self.symbolDataBySymbol ={} 
        self.day=None
        self.optionDataBySymbol={} #Keyed by symbol, valued by timestamp/greeks
    def Update(self, algorithm, slice):
        
        insights=[]
        for symbol, symbolData in self.symbolDataBySymbol.items():
            if not slice.Bars.ContainsKey(symbol) :
                continue
            
            
            option = algorithm.AddOption(symbolData.Symbol)
            option.SetFilter(lambda universe: universe.Strikes(-2, 2).Expiration(timedelta(60), timedelta(90)))
            for chain in slice.OptionChains:
                volatility = algorithm.Securities[chain.Key.Underlying].VolatilityModel.Volatility
                for contract in chain.Value:
                    if contract.TheoreticalPrice<0.001: #TOl for NPV
                        continue
                    
                    
                    
                    #Deal with NoneTypes here
                    
                    
                    if contract.Symbol.Value not in self.optionDataBySymbol:
                        
                        #First instance of contract append key-value pair
                        self.optionDataBySymbol[contract.Symbol.Value]=[[
                        str(algorithm.Time),
                         #contract.LastPrice,
                        contract.TheoreticalPrice,
                        
                        #Point of time
                        contract.BidPrice, #Bid's
                        contract.AskPrice #Ask's
                        #Last Price
                        #contract.OpenInterest, #Open Interest
                        #volatility, #Underlying Volatility 
                    
                         #calculated from our option price model,
                        #contract.Greeks.Delta,#All of our greeks are calculated from our pricing method as well
                        #contract.Greeks.Gamma,
                        #contract.Greeks.Vega,
                        #contract.Greeks.Rho,
                        #contract.Greeks.Theta / 365,
                        #contract.ImpliedVolatility]]
                        ]]
                    else:
                        
                        self.optionDataBySymbol[contract.Symbol.Value].append([ #already had a data point append additional datapoint 
                          str(algorithm.Time),
                         #contract.LastPrice,
                        contract.TheoreticalPrice,
                        contract.BidPrice,
                       contract.AskPrice
                   
                       # contract.OpenInterest,
                       # volatility,
                    
                       
                        #contract.Greeks.Delta,
                       # contract.Greeks.Gamma,
                        #contract.Greeks.Vega,
                        #contract.Greeks.Rho,
                        #contract.Greeks.Theta / 365,
                        #contract.ImpliedVolatility])
                        
                        ])
                        
                        
        if self.algorithm.endOfDay==True:
            dump = json.dumps(self.optionDataBySymbol)
        
            algorithm.ObjectStore.Save("MyObject", dump)    
            
            #End of day log all of the contracts traded along with respective datapoints 
            #for contract in self.optionDataBySymbol.keys():
                #algorithm.Log(contract)
                #for datapoint in self.optionDataBySymbol[contract]:
                    #algorithm.Log("Time={0} Bid={1} Ask={2} Last={3} OI={4} sigma={5:.3f} NPV={6:.3f} delta={7:.3f} gamma={8:.3f} vega={9:.3f} beta={10:.2f} theta={11:.2f} IV={12: .2f}".format(datapoint[0]
                   # ,datapoint[1],datapoint[2],datapoint[3],datapoint[4],datapoint[5],datapoint[6],datapoint[7],datapoint[8],datapoint[9],datapoint[10],datapoint[11],datapoint[12]))
                             
            self.algorithm.endOfDay=False
            
        return insights
        
    def OnSecuritiesChanged(self, algorithm, changes):
        addedSymbols = [ x.Symbol for x in changes.AddedSecurities if (x.Symbol not in self.symbolDataBySymbol and x.Symbol.SecurityType ==SecurityType.Equity)]
        if len(addedSymbols) == 0: return
        #if no new symbols we do not need to generate any new instances
        for symbol in addedSymbols:
            #can pass all indicators and consolidators through symbol data class
            self.symbolDataBySymbol[symbol] = SymbolData(symbol, algorithm)
            
class SymbolData:  #Here store all data associated with symbol
     def __init__(self, symbol, algorithm):
         self.Symbol = symbol