#region imports
from AlgorithmImports import *

from datetime import datetime
#endregion

class OptionsUniverseSelectionModel(UniverseSelectionModel):
    def __init__(self, universeSelectionModel: UniverseSelectionModel, optionFilter: Optional[Callable[[OptionFilterUniverse], OptionFilterUniverse]] = None) -> None:
        self.__nextRefreshTimeUtc: datetime = datetime.max
        self.__universeSelectionModel: UniverseSelectionModel = universeSelectionModel
        self.__optionFilter: Optional[Callable[[OptionFilterUniverse], OptionFilterUniverse]] = optionFilter
        self.__fundamentalUniverses: Dict[Universe, List[Symbol]] = {}

    def GetNextRefreshTimeUtc(self) -> datetime:
        parentRefreshTime = self.__universeSelectionModel.GetNextRefreshTimeUtc()
        if parentRefreshTime <= self.__nextRefreshTimeUtc:
            self.__fundamentalUniverses = {}
            self.__nextRefreshTimeUtc = parentRefreshTime
        return self.__nextRefreshTimeUtc
    
    def CreateUniverses(self, algorithm: QCAlgorithm) -> List[Universe]:
        self.__nextRefreshTimeUtc = datetime.max

        if len(self.__fundamentalUniverses) <= 0:
            universes: List[Universe] = self.__universeSelectionModel.CreateUniverses(algorithm)

            for universe in universes:
                self.__fundamentalUniverses[universe] : List[Symbol] = []
                universe.SelectionChanged += self.UniverseSelectionChanged
        
        allUniverses = []
        for universe in self.__fundamentalUniverses:
            allUniverses.append(universe)

            for optionSymbol in self.__fundamentalUniverses[universe]:
                optionChainUniverse = self.CreateOptionChain(algorithm, optionSymbol, self.__optionFilter, universe.UniverseSettings)
                allUniverses.append(optionChainUniverse)
        
        return allUniverses

    def CreateOptionChainSecurity(self, algorithm: QCAlgorithm, symbol: Symbol, settings: UniverseSettings) -> Security:
        '''Creates the canonical option chain security for a given symbol
        Args:
            algorithm: The algorithm instance to create universes for
            symbol: Symbol of the option
            settings: Universe settings define attributes of created subscriptions, such as their resolution and the minimum time in universe before they can be removed
        Returns
            Option for the given symbol'''
        config = algorithm.SubscriptionManager.SubscriptionDataConfigService.Add(symbol,
                                                                                settings.Resolution,
                                                                                settings.FillForward,
                                                                                settings.ExtendedMarketHours,
                                                                                False)
        security = algorithm.Securities.CreateSecurity(symbol, config, settings.Leverage, False)

        underlying = symbol.Underlying
        if not algorithm.Securities.ContainsKey(underlying):
            underlying_security = algorithm.AddSecurity(underlying.SecurityType,
                underlying.Value,
                settings.Resolution,
                underlying.ID.Market,
                False,
                0,
                settings.ExtendedMarketHours)
        else:
            underlying_security = algorithm.Securities[underlying]

        self.ConfigureUnderlyingSecurity(algorithm, settings.Resolution, underlying_security)

        security.Underlying = underlying_security

        return security
    
    def CreateOptionChain(self, algorithm: QCAlgorithm, symbol: Symbol, optionFilter: Optional[Callable[[OptionFilterUniverse], OptionFilterUniverse]], universeSettings: UniverseSettings) -> OptionChainUniverse:
        '''Creates a OptionChainUniverse for a given symbol
        Args:
            algorithm: The algorithm instance to create universes for
            symbol: Symbol of the option
        Returns:
            OptionChainUniverse for the given symbol'''
        if not Extensions.IsOption(symbol.SecurityType):
            raise ValueError("CreateOptionChain requires an option symbol.")

        # rewrite non-canonical symbols to be canonical
        market = symbol.ID.Market
        underlying = symbol.Underlying
        if not symbol.IsCanonical():
            alias = f"?{underlying.Value}"
            symbol = Symbol.Create(underlying.Value, SecurityType.Option, market, alias)

        # create canonical security object, but don't duplicate if it already exists
        securities = [s for s in algorithm.Securities if s.Key == symbol]
        if len(securities) == 0:
            optionChain = self.CreateOptionChainSecurity(algorithm, symbol, universeSettings)
        else:
            optionChain = securities[0]

        # set the option chain contract filter function
        optionChain.SetFilter(optionFilter)

        # force option chain security to not be directly tradable AFTER it's configured to ensure it's not overwritten
        optionChain.IsTradable = False

        return OptionChainUniverse(optionChain, universeSettings, algorithm.LiveMode)
    
    def UniverseSelectionChanged(self, sender: Object, args: EventArgs) -> None:
        self.__fundamentalUniverses[sender] = [
            Symbol.CreateCanonicalOption(symbol, symbol.ID.Market)
            for symbol in args.CurrentSelection
        ]
        self.__nextRefreshTimeUtc = datetime.min

    def ConfigureUnderlyingSecurity(self, algorithm: QCAlgorithm, resolution: Resolution, security: Security) -> None:
        configs = algorithm.SubscriptionManager.SubscriptionDataConfigService.GetSubscriptionDataConfigs(security.Symbol)
        
        for config in configs:
            if config.DataNormalizationMode != DataNormalizationMode.Raw:
                security.SetDataNormalizationMode(DataNormalizationMode.Raw)
                security.RefreshDataNormalizationModeProperty()

        if security.VolatilityModel == VolatilityModel.Null:
            if resolution == Resolution.Daily:
                updateFrequency = timedelta(days=1)
            elif resolution == Resolution.Hour:
                updateFrequency = timedelta(hours=1)
            elif resolution == Resolution.Minute:
                updateFrequency = timedelta(minutes=1)
            else:
                updateFrequency = timedelta(second=1)

            security.VolatilityModel = StandardDeviationOfReturnsVolatilityModel(resolution, updateFrequency)

# region imports
from AlgorithmImports import *

from OptionsUniverseSelection import OptionsUniverseSelectionModel
from QuantConnect.Securities.Option import OptionPriceModels
# endregion

class JumpingMagentaKitten(QCAlgorithm):

    def Initialize(self) -> None:
        self.SetStartDate(2021, 1, 19)  # Set Start Date
        self.SetEndDate(2021, 1, 25)    # Set End Date
        self.SetCash(100000)  # Set Strategy Cash
        
        self.SetSecurityInitializer(self.CustomSecurityInitializer)
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverseSelection(
            OptionsUniverseSelectionModel(
                CoarseFundamentalUniverseSelectionModel(self.SelectCoarse),
                self.OptionFilterFunction
            )
        )
        
        # set the warm-up period for the pricing model
        self.SetWarmUp(TimeSpan.FromDays(4), self.UniverseSettings.Resolution)

    def CustomSecurityInitializer(self, security: Security) -> None:
        # set the pricing model for Greeks and volatility
        # find more pricing models https://www.quantconnect.com/docs/v2/writing-algorithms/reality-modeling/options-models/pricing
        if security.Type == SecurityType.Option:
            security.PriceModel = OptionPriceModels.CrankNicolsonFD()

    def OnData(self, data: Slice) -> None:
        if self.IsWarmingUp:
            return
        
        for chain in data.OptionChains:
            volatility = self.Securities[chain.Key.Underlying].VolatilityModel.Volatility
            for contract in chain.Value:
                self.Log("{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(
                contract.Symbol.Value,
                contract.BidPrice,
                contract.AskPrice,
                contract.LastPrice,
                contract.OpenInterest,
                volatility,
                contract.TheoreticalPrice,
                contract.Greeks.Delta,
                contract.Greeks.Gamma,
                contract.Greeks.Vega,
                contract.Greeks.Rho,
                contract.Greeks.Theta / 365,
                contract.ImpliedVolatility))

    def SelectCoarse(self, coarse: List[CoarseFundamental]) -> List[Symbol]:
        selected = [c for c in coarse if c.HasFundamentalData]
        sorted_by_dollar_volume = sorted(selected, key=lambda c: c.DollarVolume, reverse=True)
        return [c.Symbol for c in sorted_by_dollar_volume[:100]] # Return most liquid assets w/ fundamentals

    def OptionFilterFunction(self, option_filter_universe: OptionFilterUniverse) -> OptionFilterUniverse:
        return option_filter_universe.Strikes(-2, +2).FrontMonth().CallsOnly()