import pandas as pd

def oc_to_df(optionchain):
    """
    Converts OptionChain data to a Pandas DataFrame
    """
    return pd.DataFrame([[x.Right,
                          float(x.Strike),
                          x.Expiry,
                          float(x.BidPrice),
                          float(x.AskPrice), 
                          float(x.LastPrice),
                          float(x.UnderlyingLastPrice),
                          float(x.TheoreticalPrice),
                          float(x.OpenInterest),
                          float(x.Volume),
                          float(x.ImpliedVolatility),
                          float(x.Greeks.Delta),
                          float(x.Greeks.Gamma),
                          float(x.Greeks.Vega),
                          float(x.Greeks.Theta),
                          float(x.Greeks.Rho)
                          ] for x in optionchain],
                        index=[x.Symbol.Value for x in optionchain],
                        columns=['Type', 'Strike', 'Expiry', 
                                 'Bid', 'Ask', 'Last', 'UnderLast',
                                 'Theo,', 'OI', 'Volume', 'IV', 
                                 'Delta', 'Gamma', 
                                 'Vega', 'Theta', 'Rho' ])

"""
Notes:
 - Figure out why greeks are not working
"""

from QuantConnect.Securities.Option import OptionPriceModels
from datetime import timedelta
import pandas as pd
from optutils import oc_to_df

class ZeroDTE(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2019, 10, 7)
        self.SetEndDate(2019, 10, 14)
        self.SetCash(100000)
        self.count = 0       
 
        self.ticker = 'SPY'  # Note: They don't seem to have index option data. No SPX!

        under = self.AddEquity(self.ticker, Resolution.Minute)  
        under.SetDataNormalizationMode(DataNormalizationMode.Raw) # IMPORTANT: default
        
        option = self.AddOption(self.ticker, Resolution.Minute)

        # set our strike/expiry filter for this option chain
        option.SetFilter(lambda universe: universe\
            .IncludeWeeklys()\
            .Strikes(-50, 50)\
            .Expiration(timedelta(0), timedelta(0)))
        
        option.PriceModel = OptionPriceModels.CrankNicolsonFD() 
        # Warmup is needed for Greeks calcs
        self.SetWarmUp(timedelta(10))  
        
    def OnData(self, slice):

        if self.IsWarmingUp: return
    
        if self.HourMinuteIs(10, 00):
            self.Debug(f'{self.Time}, # of Chains: {len(slice.OptionChains)}')
            if len(slice.OptionChains) == 1:
                for kvp in slice.OptionChains:
                    df = oc_to_df(kvp.Value)
                    self.Debug(f'# of Contracts: {len(df.index)}')
                    self.Debug(f'{df.head()}')
                    self.Debug(f'Min Delta: {min(df.Delta)}, Max Delta: {max(df.Delta)}')
                    self.Debug(f'Min IV: {min(df.IV)}, Max IV: {max(df.IV)}')
        
    def HourMinuteIs(self, hour, minute):
        return self.Time.hour == hour and self.Time.minute == minute
        
    def OnOrderEvent(self, orderEvent):
       self.Debug(f'{self.time} Order Event: {orderEvent}')
        
    def OnAssignmentOrderEvent(self, assignmentEvent):
        self.Debug(f'{self.time} Assignment Event: {assignmentEvent}')