# 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)

        self.day = 1
        self.universe1 = []
        self.universe2 = []

    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)
        self.universe1 = [f.Symbol for f in sorted_by_pe_ratio[:10]]
        return self.universe1

    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)
        self.universe2 = [f.Symbol for f in sorted_by_pe_ratio[:10]]
        return self.universe2

    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() == 4:

            if self.Time.hour == 10 and self.Time.minute == 30:

                self.day = -1

        if self.day < 0:
        
            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 not calls or not puts:
                    continue
                
                if chain.Underlying.Symbol in self.universe1:
                    symbol1 = calls[0].Symbol
                    self.MarketOrder(symbol1, 1)
                    symbol2 = puts[0].Symbol
                    self.MarketOrder(symbol2, 1)
                if chain.Underlying.Symbol in self.universe2:
                    symbol1 = calls[0].Symbol
                    self.MarketOrder(symbol1, -1)
                    symbol2 = puts[0].Symbol
                    self.MarketOrder(symbol2, -1)

            self.day = 1