# region imports
from AlgorithmImports import *
# endregion
import math, numpy as np
from datetime import timedelta
from math import floor
from decimal import Decimal

from datetime import timedelta
from QuantConnect.DataSource import *


class ChainedUniverseAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2023, 8, 2)
        self.SetCash(100000)
        self.UniverseSettings.Asynchronous = True
        self.UniverseSettings.DataNormalizationMode = DataNormalizationMode.Raw

        universe1 = self.AddUniverse(self.FundamentalFunction)
        self.AddUniverseOptions(universe1, self.OptionFilterFunction)
        universe2 = self.AddUniverse(self.FundamentalFunction2)
        self.AddUniverseOptions(universe2, self.OptionFilterFunction)

    def FundamentalFunction(self, fundamental: List[Fundamental]) -> List[Symbol]:
        filtered = (f for f in fundamental if not np.isnan(f.ValuationRatios.PERatio))
        sorted_by_pe_ratio = sorted(filtered, key=lambda f: f.ValuationRatios.PERatio)
        return [f.Symbol for f in sorted_by_pe_ratio[:10]]

    def FundamentalFunction2(self, fundamental: List[Fundamental]) -> List[Symbol]:
        filtered = (f for f in fundamental if not np.isnan(f.ValuationRatios.PERatio))
        sorted_by_pe_ratio = sorted(filtered, key=lambda f: f.ValuationRatios.PERatio, reverse=True)
        return [f.Symbol for f in sorted_by_pe_ratio[:10]]

    def OptionFilterFunction(self, option_filter_universe: OptionFilterUniverse) -> OptionFilterUniverse:
        return option_filter_universe.Strikes(-2, +2).Expiration(30, 60)

    def OnData(self, data: Slice) -> None:
       
        if self.Time.weekday() == 5:

            if self.Time.hour == 10 and self.Time.minute == 30:
        
                for symbol, chain in data.OptionChains.items():
                    if self.Portfolio[chain.Underlying.Symbol].Invested:
                        self.Liquidate(chain.Underlying.Symbol)

                    spot = chain.Underlying.Price
                    atm_strike = sorted(chain, key=lambda x: abs(x.Strike - chain.Underlying.Price))[0].Strike
                    calls = [i for i in chain if i.Strike == atm_strike and i.Right == OptionRight.Call]
                    puts = [i for i in chain if i.Strike == atm_strike and i.Right == OptionRight.Put]
                   
                    if chain.Underlying.Symbol in universe1:
                        symbol1 = calls[0].Symbol
                        self.MarketOrder(symbol1, 1)
                        symbol2 = puts[0].Symbol
                        self.MarketOrder(symbol2, 1)
                    if chain.Underlying.Symbol in universe2:
                        symbol1 = calls[0].Symbol
                        self.MarketOrder(symbol1, -1)
                        symbol2 = puts[0].Symbol
                        self.MarketOrder(symbol2, -1)