# region imports
from AlgorithmImports import *
# endregion

class FormalBlackAnguilline(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2022, 9, 3)
        self.SetEndDate(2022, 10, 4)
        self.SetCash(100000)  # Set Strategy Cash
        spy = self.AddEquity("SPY", Resolution.Daily)
        VolatilityModel = StandardDeviationOfReturnsVolatilityModel(30, Resolution.Daily)
        self.symbol = spy.Symbol

        contracts = [
            Symbol.CreateOption(self.symbol, Market.USA, OptionStyle.American, OptionRight.Put, 378, datetime(2022, 10, 5)),
            Symbol.CreateOption(self.symbol, Market.USA, OptionStyle.American, OptionRight.Call, 377, datetime(2022, 10, 5))
        ]
        for contract in contracts:
            option = self.AddOptionContract(contract, Resolution.Daily)
            option.PriceModel = OptionPriceModels.BjerksundStensland()

        self.SetWarmUp(31)

        self.df = pd.DataFrame()


    def OnData(self, data: Slice):
        if self.IsWarmingUp: return

        equity = self.Securities[self.symbol]
        for canonical_symbol, chain in data.OptionChains.items():
            for contract in chain:
                greeks = contract.Greeks
                data = {
                    "IV" : contract.ImpliedVolatility,
                    "Delta": greeks.Delta,
                    "Gamma": greeks.Gamma,
                    "Vega": greeks.Vega,
                    "Rho": greeks.Rho,
                    "Theta": greeks.Theta,
                    "LastPrice": contract.LastPrice,
                    "Close": self.Securities[contract.Symbol].Close,
                    "theoreticalPrice" : contract.TheoreticalPrice,
                    "underlyingPrice": equity.Close,
                    "underlyingVolatility": equity.VolatilityModel.Volatility
                }
                symbol = contract.Symbol
                right = "Put" if symbol.ID.OptionRight == 1 else "Call"
                index = pd.MultiIndex.from_tuples([(symbol.ID.Date, symbol.ID.StrikePrice, right, symbol.Value, self.Time)], names=["expiry", "strike", "type", "symbol", "endTime"])
                self.df = pd.concat([self.df, pd.DataFrame(data, index=index)])
    

    def OnEndOfAlgorithm(self):

        self.ObjectStore.Save("price-models/backtest-df", self.df.sort_index().to_csv())