from AlgorithmImports import *
import pandas as pd
import numpy as np

class MSFTOptionAnalysis(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2024, 1, 18)  # Start date
        self.SetEndDate(2024, 1, 19)    # End date - data is collected only for the target date, avoiding unnecessary computation
        self.SetCash(100000)            # Starting cash
        
        self.equity = self.AddEquity("MSFT", Resolution.Hour)
        option = self.AddOption("MSFT", Resolution.Hour)
        option.SetFilter(lambda u: u.Strikes(-50, 50).Expiration(0, 90))    # Broad filter only for retrieving purposes
        
        self.target_date = datetime(2024, 1, 18, 15, 0, 0)  # Target date, time, and expiry
        self.expiry_date = datetime(2024, 3, 15)
        
        # Our strikes
        self.call_strike = 360
        self.put_strike = 320
        self.otc_put_strike = 328.5
        
        self.option_data = {}
        self.data_processed = False  #preventing OnData from running multiple times and flooding with logs
    
    def OnData(self, slice):
        #process data at Jan 18, 2024 at 3:00 PM
        if not self.ShouldProcessData():
            return
        
        chain = self.GetOptionChain(slice)
        if not chain:
            return
        
        self.StoreOptionPrices(chain)
        self.InterpolateOTCImpliedVol(chain)
        self.PrintResults()
        
        self.data_processed = True
    
    def ShouldProcessData(self):
        return (
            self.Time.date() == self.target_date.date() and
            self.Time.hour == self.target_date.hour and
            not self.data_processed
        )
    
    def GetOptionChain(self, slice):
        if slice.OptionChains.Count == 0:
            return None
        chains = list(slice.OptionChains.Values)
        return chains[0] if chains else None
    
    def StoreOptionPrices(self, chain):
        underlying_price = chain.Underlying.Price
        self.Debug(f"Underlying Price (MSFT): ${underlying_price}")

        #store all call options expiring on our target date
        put_contracts = [c for c in chain if c.Right == OptionRight.Put and c.Expiry.date() == self.expiry_date.date()]
        call_contracts = [c for c in chain if c.Right == OptionRight.Call and c.Expiry.date() == self.expiry_date.date()]
        
        self.StoreOptionData(call_contracts, self.call_strike, 'call_360')
        self.StoreOptionData(put_contracts, self.put_strike, 'put_320')
    
    def StoreOptionData(self, contracts, strike, key):
        option = next((c for c in contracts if c.Strike == strike), None)
        if option:
            price = (option.BidPrice + option.AskPrice) / 2
            self.option_data[key] = {
                'price': price,
                'strike': option.Strike,
                'expiry': option.Expiry,
                'days_to_expiry': (self.expiry_date - self.Time).days,
                'implied_vol': option.ImpliedVolatility  # Use contract's implied volatility
            }
    
    def InterpolateOTCImpliedVol(self, chain):

        ## Calls are not used because IV behavior can be diff between calls & Puts due to skew/smile effects

        put_contracts = [c for c in chain if c.Right == OptionRight.Put and c.Expiry.date() == self.expiry_date.date()]
        put_data = [{'strike': p.Strike, 'implied_vol': p.ImpliedVolatility} for p in put_contracts]
        
        if put_data:
            df = pd.DataFrame(put_data).sort_values('strike').set_index('strike')
            if len(df) >= 2 and min(df.index) < self.otc_put_strike < max(df.index):
                self.option_data['otc_put'] = {
                    'strike': self.otc_put_strike,
                    'implied_vol': np.interp(self.otc_put_strike, df.index, df['implied_vol'])
                }
    
    def PrintResults(self):
        self.Debug("Analysis Results:")
        for key in ['call_360', 'put_320']:
            if key in self.option_data:
                data = self.option_data[key]
                self.Debug(f"MSFT {key.split('_')[0].capitalize()} Option (Strike: {data['strike']}, Expiry: {data['expiry'].strftime('%Y-%m-%d')}):")
                self.Debug(f"Price: ${data['price']:.2f}")
                self.Debug(f"Implied Volatility: {data['implied_vol']:.4f}")
        
        if 'otc_put' in self.option_data:
            otc_data = self.option_data['otc_put']
            self.Debug(f"Interpolated IV for OTC Put Option (Strike: {otc_data['strike']}, Expiry: {self.expiry_date.strftime('%Y-%m-%d')}):")
            self.Debug(f"{otc_data['implied_vol']:.4f}")