#region imports
from AlgorithmImports import *
#endregion

LOOKBACK = 20
RES = Resolution.HOUR

class YouTubeSeries(QCAlgorithm):
    def initialize(self):
        self.set_start_date(2024, 1, 1)
        self.set_cash(100000)

        self.universe_settings.asynchronous = True
        self.universe_settings.data_normalization_mode = DataNormalizationMode.RAW
        self.universe_settings.resolution = RES
        
        universe = self.add_universe(self._select_coarse)
        self.add_universe_options(universe, self._option_filter_function)

        self.Indics = {}
        self.dte = 15
        self.uni_size = 10
        
        self.AddRiskManagement(MaximumUnrealizedProfitPercentPerSecurity(1.00))
        self.AddRiskManagement(MaximumDrawdownPercentPerSecurity(0.50))

    def _select_coarse(self, coarse: List[CoarseFundamental]) -> List[Symbol]:
        selected = [c for c in coarse if c.has_fundamental_data]
        sorted_by_dollar_volume = sorted(selected, key=lambda c: c.dollar_volume, reverse=True)
        symbols = [c.symbol for c in sorted_by_dollar_volume[:self.uni_size]]
        return symbols

    def _option_filter_function(self, universe: OptionFilterUniverse) -> OptionFilterUniverse:
        return universe.IncludeWeeklys().Strikes(5, 5).Expiration(timedelta(self.dte), timedelta(self.dte))

    def on_data(self, data: Slice):

        option_invested = [x.Key for x in self.Portfolio if x.Value.Invested and x.Value.Type==SecurityType.OPTION]
        if option_invested:
            if self.Time + timedelta(5) > option_invested[0].ID.Date:
                self.Liquidate(option_invested[0], "Too close to expiration")
            
        equity_invested = [x.Key for x in self.Portfolio if x.Value.Invested and x.Value.Type==SecurityType.EQUITY]
        if equity_invested:
                self.Liquidate(equity_invested[0], tag="Equity in here")

        for _, chain in data.option_chains.items():

            symbol = chain.underlying.symbol
            spot = chain.underlying.price

            exp = sorted(chain, key = lambda x: x.Expiry, reverse=False)
            if len(exp) == 0: continue
            expiry = exp[0].Expiry

            calls = [i for i in chain if i.Expiry == expiry and i.Right == OptionRight.CALL]
            call_contracts = sorted(calls, key = lambda x: abs(x.Strike - spot))
            if len(call_contracts) == 0: continue
            call_con = call_contracts[0]

            if call_con.ask_price > 10.00: continue

            puts = [i for i in chain if i.Expiry == expiry and i.Right == OptionRight.PUT]
            put_contracts = sorted(puts, key = lambda x: abs(x.Strike - spot))
            if len(put_contracts) == 0: continue
            put_con = put_contracts[0]

            if put_con.ask_price > 10.00: continue

            if not (data.contains_key(symbol) and data[symbol] is not None):
                continue

            if not symbol in self.Indics: continue

            bar = data.bars.get(symbol)
            if bar: 
                self.Indics[symbol].st.update(bar)
                self.Indics[symbol].stch.update(bar)

            if not self.Indics[symbol].st.is_ready or not self.Indics[symbol].stch.is_ready: continue


            if not self.portfolio[call_con.symbol].Invested or not self.portfolio[put_con.symbol].Invested:

                c = self.securities[symbol].close
                st = self.Indics[symbol].st.current.value
                stch = self.Indics[symbol].stch.current.value

                if stch > 70 and st < c:
                    self.market_order(call_con.symbol, 1)

                elif stch < 30 and st > c:
                    self.market_order(put_con.symbol, 1)

        
    def on_securities_changed(self, changes: SecurityChanges) -> None:

        for security in changes.added_securities:
            if security.type == SecurityType.EQUITY:
                self.Indics[security.symbol] = Indics()

                trade_bars = self.history[TradeBar](security.symbol, LOOKBACK+1, RES)
                for trade_bar in trade_bars:
                    self.Indics[security.symbol].st.update(trade_bar)
                    self.Indics[security.symbol].stch.update(trade_bar)

        for security in changes.removed_securities:
            if security.type == SecurityType.EQUITY:
                indic = self.Indics.pop(security.symbol, None)

class Indics:
    def __init__(self):
        self.st = SuperTrend(LOOKBACK, 2, MovingAverageType.Wilders)
        self.stch = Stochastic(LOOKBACK, 10, 20)