from AlgorithmImports import *
from decimal import Decimal, ROUND_HALF_UP


class ConnorsCrash(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2024, 9, 1)
        self.set_end_date(2024, 12, 31)
        self.set_cash(100_000)
        self._rsi_period = 3
        self._streak_period = 2
        self._pct_rank_period = 100
        self._vola_period = 100
        self._crsi_entry = 90
        self._crsi_exit = 30
        self._max_shorts = 40
        self.settings.automatic_indicator_warm_up = True
        self.settings.seed_initial_prices = True
        self.universe_settings.resolution = Resolution.DAILY
        self._spy = self.add_equity("SPY")
        # Select liquid, tradeable-priced equities for the universe.
        self._universe = self.add_universe(
            lambda fundamentals: [f.symbol for f in fundamentals if f.price > 5 and f.dollar_volume > 1e6]
        )
        self.schedule.on(
            self.date_rules.every_day(self._spy),
            self.time_rules.after_market_open(self._spy, 1),
            self._rebalance
        )

    def on_securities_changed(self, changes):
        for security in changes.added_securities:
            # Attach ConnorsRSI indicator to each security.
            security.connors = self.crsi(security, self._rsi_period, self._streak_period, self._pct_rank_period)
            # Initialize and warm up the custom volatility indicator.
            security.volatility = CustomVolatility(self._vola_period)
            for bar in self.history[TradeBar](security, self._vola_period + 1):
                security.volatility.update(bar)
            self.register_indicator(security, security.volatility)
        for security in changes.removed_securities:
            self.deregister_indicator(security.connors)
            self.deregister_indicator(security.volatility)
            self.liquidate(security)

    def _rebalance(self):
        if not self._universe.selected or not Extensions.is_market_open(self._spy, False):
            return
        securities = [self.securities[symbol] for symbol in self._universe.selected]
        securities = [s for s in securities if s.connors.is_ready and s.volatility.is_ready]
        filter_vola = [s for s in securities if s.volatility.value > 100]
        # Find currently invested short positions.
        short_positions = [s for s in securities if s.holdings.is_short]
        # Liquidate short positions when ConnorsRSI falls below exit threshold.
        for security in short_positions:
            if security.connors.current.value < self._crsi_exit:
                self.liquidate(security)
        # Exclude symbols with pending open orders.
        pending_symbols = {t.symbol for t in self.transactions.get_open_order_tickets()}
        # Find short entry candidates that are not currently invested (high volatility and high CRSI).
        short_candidates = sorted([
            s for s in filter_vola
            if (s.connors.current.value > self._crsi_entry and
            not s.holdings.invested and
            s.symbol not in pending_symbols)
        ], key=lambda s: (s.connors.current.value, s.volatility.value))
        # Set union: liquidated positions are only counted once.
        short_position_symbols = {s.symbol for s in short_positions}
        occupied = short_position_symbols | pending_symbols
        available_slots = self._max_shorts - len(occupied)
        if available_slots <= 0 or not short_candidates:
            return
        n_orders = min(len(short_candidates), available_slots)
        target_weight = -1 / self._max_shorts
        for security in short_candidates[-n_orders:]:
            quantity = int(self.calculate_order_quantity(security, target_weight))
            if quantity:
                limit_price = float((Decimal('1.03') * Decimal(str(security.price))).quantize(Decimal('0.01'), rounding=ROUND_HALF_UP))
                self.limit_order(security, quantity, limit_price)


class CustomVolatility(PythonIndicator):

    def __init__(self, period):
        super().__init__()
        self.value = 0
        self._window = RollingWindow[float](period)

    def update(self, input_: BaseData):
        # Annualized log-return volatility.
        price = input_.value
        if price <= 0:
            return self.is_ready
        self._window.add(price)
        if self._window.is_ready:
            prices = np.array(list(self._window)[::-1])
            log_diffs = np.diff(np.log(prices))
            self.value = np.std(log_diffs) * math.sqrt(252) * 100.0
        return self.is_ready

    @property
    def is_ready(self) -> bool:
        return self._window.is_ready