# region imports
from AlgorithmImports import *
from datetime import timedelta
import math
import numpy as np
# endregion


class ConnorsCrash(QCAlgorithm):

    def initialize(self):
        self.set_start_date(self.end_date - timedelta(5*365))
        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 = 20
        self.settings.automatic_indicator_warm_up = True
        self.universe_settings.resolution = Resolution.DAILY
        self._universe = self.add_universe(
            lambda fundamentals: [f.symbol for f in fundamentals if f.price > 5 and f.volume > 1e6] # switch to dollar volume 
        )

    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 custom volatility indicator and.
            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.volatility)
            self.liquidate(security)

    def on_data(self, data):
        # Filter securities with volatility > 100%.
        securities = [self.securities[symbol] for symbol in self._universe.selected]
        filter_vola = [s for s in securities if s.volatility.is_ready and s.volatility.value > 100]
        # Find invested short positions positions.
        short_positions = [s for s in securities if self.portfolio[s.symbol].is_short]
        # Liquidate short positions when ConnorsRSI falls below exit threshold.
        for security in short_positions:
            if security.connors.is_ready and security.connors.current.value < self._crsi_exit:
                self.liquidate(security)
        # Track symbols with open pending orders to avoid duplicate entries.
        pending_symbols = {t.symbol for t in self.transactions.get_open_order_tickets()}
        # Find short entry candidates (high volatility and high CRSI)
        short_candidates = [
            s for s in filter_vola
            if (s.connors.is_ready and
                s.connors.current.value > self._crsi_entry and
                not self.portfolio[s.symbol].invested and
                s.symbol not in pending_symbols)
        ]
        # Calculate available position slots
        available_slots = self._max_shorts - len(short_positions) - len(pending_symbols) 
        if available_slots <= 0 or not short_candidates:
            return
        # Place short orders for candidates using available margin for shorts.
        n_orders = min(len(short_candidates), available_slots)
        for i, security in enumerate(short_candidates[:n_orders]):
            if security.price <= 0:
                continue
            remaining = n_orders - i 
            # Equal weight the available margin per slot. ### shorts are 50% margin so we can use 2x the margin for position sizing
            max_position_value = self.portfolio.margin_remaining * 0.70 / remaining
            quantity = int(max_position_value / security.price)
            if quantity > 0:
                self.limit_order(security, -quantity, round(1.03 * security.price, 2))



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._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