from AlgorithmImports import *
import numpy as np
import json
from datetime import datetime, timedelta

# WARNING: This is an example to illustrate a "flaw" I'm seeing in a lot of QuantConnect published algorithms.
# That is, the use of Resolution.DAILY when algorithm leverages trade within last N minutes of market closing.
# Why is this a problem? When the equity symbol is subscribed on a daily basis,
# backtesting will fill the price at yesterday's closing price. When you backtest this algorithm, you'll see 
# that 2025/4/18  3:55pm (ET) will fill the price on the close price of 2025/4/17.
class TrendBarbellMetaAlgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2015, 1, 1)
        self.set_end_date(2025, 1, 1)  # ignored in live trading
        self.set_cash(10000)
        self.settings.rebalance_portfolio_on_insight_changes = False
        self.set_brokerage_model(BrokerageName.CHARLES_SCHWAB, AccountType.MARGIN)

        # ── 1. THE REGIME ARBITER ─────────────────────────────────────────────
        self.spy = self.add_equity("SPY", Resolution.MINUTE).symbol
        self.inverse_spy = self.add_equity("SH", Resolution.MINUTE).symbol

        self.spy_sma_200 = self.sma(self.spy, 200, Resolution.DAILY)
        self.spy_ema = self.ema(self.spy, 2, Resolution.DAILY)
        self.spy_roc_3 = self.roc(self.spy, 3, Resolution.DAILY)

        # ── 2. MICRO SETUP ────────────────────────────────────────────────────
        self.vwap_ind = self.vwap(self.spy)
        self.imbalance_sma = SimpleMovingAverage(5)
        # Note: This is the standard deviation of the closing price, not the VWAP itself.
        self.price_std = self.std(self.spy, 390, Resolution.MINUTE)

        self.trades_today = 0
        self.micro_entry_price = 0.0
        self.micro_highest_price = 0.0
        self.micro_lowest_price = 99999.0
        self._last_state_save = datetime.min

        # ── OPENING RANGE BREAKOUT (ORB) ─────────────────────────────────────
        self._or_high = 0.0
        self._or_low = 99999.0
        self._or_complete = False
        self._orb_trade = False
        
        # Unique ObjectStore key per deployment
        self._deploy = self.get_parameter("deployment_id") or "paper"
        self._state_key = "barbell_state" if self._deploy == "live" else f"barbell_state_{self._deploy}"

        # ── 3. MACRO BASKET SETUP ─────────────────────────────────────────────
        self.tqqq = self.add_equity("TQQQ", Resolution.DAILY).symbol
        self.soxl = self.add_equity("SOXL", Resolution.DAILY).symbol
        self.tna = self.add_equity("TNA", Resolution.DAILY).symbol

        self.bear_tickers = ["GLD", "TLT", "TMF", "USO", "UUP", "FXY"]
        self.bear_symbols = {t: self.add_equity(t, Resolution.DAILY).symbol for t in self.bear_tickers}
        self.bear_momentum = {
            t: self.momp(self.bear_symbols[t], 21, Resolution.DAILY)
            for t in self.bear_tickers
        }
        self.last_bear_rebalance = datetime(2000, 1, 1)

        self.schedule.on(self.date_rules.every_day(self.spy), self.time_rules.midnight, self.reset_daily_limits)
        self.schedule.on(self.date_rules.every_day(self.spy), self.time_rules.after_market_open(self.spy, 5), self.macro_rebalance)
        # self.schedule.on(self.date_rules.every_day(self.spy), self.time_rules.before_market_close(self.spy, 5), self.macro_rebalance)
        self.set_warm_up(210, Resolution.DAILY)

    def reset_daily_limits(self):
        self.trades_today = 0
        if not self.portfolio[self.spy].is_long and not self.portfolio[self.inverse_spy].invested:
            self.micro_entry_price = 0.0
            self.micro_highest_price = 0.0
            self.micro_lowest_price = 99999.0
        self._or_high = 0.0
        self._or_low = 99999.0
        self._or_complete = False
        self._orb_trade = False
        self._save_state()

    def on_warmup_finished(self):
        today = self.time.date()
        self.trades_today = sum(
            1 for o in self.transactions.get_orders()
            if o.time.date() == today and o.status == OrderStatus.FILLED
        )
        spy_price = self.securities[self.spy].price
        if self.portfolio[self.spy].is_long:
            self.micro_entry_price = self.portfolio[self.spy].average_price
            self.micro_highest_price = spy_price
        elif self.portfolio[self.inverse_spy].invested:
            self.micro_entry_price = self.portfolio[self.inverse_spy].average_price
            self.micro_lowest_price = spy_price
        self._load_state()
        self.debug(f"Indicator Status - SMA: {self.spy_sma_200.is_ready}, EMA: {self.spy_ema.is_ready}")

    def _save_state(self):
        state = {
            "last_bear_rebalance": self.last_bear_rebalance.isoformat(),
            "state_date": self.time.date().isoformat(),
            "trades_today": self.trades_today,
            "micro_entry_price": self.micro_entry_price,
            "micro_highest_price": self.micro_highest_price,
            "micro_lowest_price": self.micro_lowest_price,
        }
        self.object_store.save(self._state_key, json.dumps(state))
        self._last_state_save = self.time

    def _load_state(self):
        if not self.object_store.contains_key(self._state_key):
            self.debug("No state found, starting fresh.")
            return
        try:
            state = json.loads(self.object_store.read(self._state_key))
            self.debug(f"Loading state: {state}")
            self.last_bear_rebalance = datetime.fromisoformat(state["last_bear_rebalance"])
            saved_date = datetime.strptime(state["state_date"], "%Y-%m-%d").date()
            if saved_date == self.time.date():
                self.trades_today = max(self.trades_today, state.get("trades_today", 0))
            if self.portfolio[self.spy].is_long:
                self.micro_entry_price = self.portfolio[self.spy].average_price
                self.micro_highest_price = max(self.micro_highest_price, state.get("micro_highest_price", 0))
            elif self.portfolio[self.inverse_spy].invested:
                self.micro_entry_price = self.portfolio[self.inverse_spy].average_price
                self.micro_lowest_price = min(self.micro_lowest_price, state.get("micro_lowest_price", 99999))
        except Exception as e:
            self.log(f"State restore failed, using safe defaults: {e}")

    def get_top_momentum_bear_assets(self, n=3):
        scored = {}
        for t in self.bear_tickers:
            mom = self.bear_momentum[t]
            if mom.is_ready:
                scored[t] = mom.current.value
        if not scored:
            return list(self.bear_symbols.keys())[:n]
        ranked = sorted(scored.keys(), key=lambda t: scored[t], reverse=True)
        return ranked[:n]

    def macro_rebalance(self):
        self.debug(f"Macro rebalance triggered at {self.time}")
        if self.is_warming_up or not self.spy_sma_200.is_ready or not self.spy_ema.is_ready:
            return
        
        spy_price = self.securities[self.spy].close
        sma200 = self.spy_sma_200.current.value
        ema_val = self.spy_ema.current.value

        if spy_price > ema_val and spy_price > sma200:
            self.debug("Regime: Bullish")
            
            # FIX: Clear micro positions to free up margin for the leveraged basket
            if self.portfolio[self.spy].invested:
                self.liquidate(self.spy, "Bull: Liquidating Micro Long to free margin")
            if self.portfolio[self.inverse_spy].invested:
                self.liquidate(self.inverse_spy, "Bull: Liquidating Micro Short to free margin")

            for t in self.bear_tickers:
                s = self.bear_symbols[t]
                if self.portfolio[s].invested:
                    self.liquidate(s, "Bull: Liquidating Protective Basket")

            tna_available = self.securities[self.tna].price > 0
            total_val = max(self.portfolio.total_portfolio_value, 1)
            
            if tna_available:
                tqqq_w = abs(self.portfolio[self.tqqq].holdings_value) / total_val
                soxl_w = abs(self.portfolio[self.soxl].holdings_value) / total_val
                tna_w = abs(self.portfolio[self.tna].holdings_value) / total_val
                if abs(tqqq_w - 0.30) > 0.05 or abs(soxl_w - 0.60) > 0.05 or abs(tna_w - 0.10) > 0.05:
                    self.set_holdings(self.tqqq, 0.30)
                    self.set_holdings(self.soxl, 0.60)
                    self.set_holdings(self.tna, 0.10)
            else:
                tqqq_w = abs(self.portfolio[self.tqqq].holdings_value) / total_val
                soxl_w = abs(self.portfolio[self.soxl].holdings_value) / total_val
                if abs(tqqq_w - 0.33) > 0.05 or abs(soxl_w - 0.67) > 0.05:
                    self.set_holdings(self.tqqq, 0.33)
                    self.set_holdings(self.soxl, 0.67)

        elif spy_price > sma200:
            self.debug("Regime: Weakening Bull")
            for t in self.bear_tickers:
                s = self.bear_symbols[t]
                if self.portfolio[s].invested:
                    mom = self.bear_momentum[t]
                    if not mom.is_ready or mom.current.value <= 0:
                        self.liquidate(s, "Weakening Bull: Clearing Negative Momentum Asset")
            if self.portfolio[self.tqqq].invested: self.liquidate(self.tqqq, "Weakening Bull: Exiting TQQQ")
            if self.portfolio[self.soxl].invested: self.liquidate(self.soxl, "Weakening Bull: Exiting SOXL")
            if self.portfolio[self.tna].invested: self.liquidate(self.tna, "Weakening Bull: Exiting TNA")

        else:
            self.debug("Regime: Bear market")
            if self.portfolio[self.tqqq].invested: self.liquidate(self.tqqq, "Bear Market: Liquidating TQQQ")
            if self.portfolio[self.soxl].invested: self.liquidate(self.soxl, "Bear Market: Liquidating SOXL")
            if self.portfolio[self.tna].invested: self.liquidate(self.tna, "Bear Market: Liquidating TNA")

            days_since_rebalance = (self.time - self.last_bear_rebalance).days
            if days_since_rebalance >= 30:
                top_assets = self.get_top_momentum_bear_assets(n=3)
                self.last_bear_rebalance = self.time
                self._save_state()

                total = self.portfolio.total_portfolio_value
                micro_weight = (abs(self.portfolio[self.spy].holdings_value) +
                                abs(self.portfolio[self.inverse_spy].holdings_value)) / total if total > 0 else 0.0
                per_asset = max(0.0, 1.0 - micro_weight) / len(top_assets)

                for t in self.bear_tickers:
                    s = self.bear_symbols[t]
                    if t not in top_assets and self.portfolio[s].invested:
                        self.liquidate(s, "Bear: Rotating Out Low Momentum Asset")

                for t in top_assets:
                    s = self.bear_symbols[t]
                    current_w = abs(self.portfolio[s].holdings_value) / max(total, 1)
                    if abs(current_w - per_asset) > 0.05:
                        self.set_holdings(s, per_asset)

    def check_micro_stops(self, price):
        # FIX: Evaluated unconditionally so stops trigger regardless of regime shifts
        if self.portfolio[self.spy].is_long:
            self.micro_highest_price = max(self.micro_highest_price, price)
            if price < self.micro_highest_price * 0.99 or price < self.micro_entry_price * 0.98:
                self.liquidate(self.spy, "Sniper Stop Loss Triggered (Long)")
                self._save_state()
        elif self.portfolio[self.inverse_spy].invested:
            self.micro_lowest_price = min(self.micro_lowest_price, price)
            if price > self.micro_lowest_price * 1.01 or price > self.micro_entry_price * 1.02:
                self.liquidate(self.inverse_spy, "Sniper Stop Loss Triggered (SH)")
                self._save_state()

    def on_data(self, data: Slice):
        if self.is_warming_up or not self.spy_sma_200.is_ready or not self.spy_ema.is_ready: 
            return

        # Safely extract quotes
        quote = data.quote_bars.get(self.spy) if data.quote_bars.contains_key(self.spy) else None
        if quote:
            tot = quote.last_bid_size + quote.last_ask_size
            if tot > 0:
                self.imbalance_sma.update(self.time, (quote.last_bid_size - quote.last_ask_size) / tot)

        if not data.bars.contains_key(self.spy): return
        spy_price = data.bars[self.spy].close

        # Check stops immediately, independently of regime
        self.check_micro_stops(spy_price)

        sma200_val = self.spy_sma_200.current.value
        ema_val = self.spy_ema.current.value
        in_non_bull = spy_price < sma200_val or (spy_price < ema_val and spy_price >= sma200_val)

        # ── OPENING RANGE BREAKOUT (ORB) ─────────────────────────────────────
        if self.time.hour == 9 and self.time.minute >= 30:
            self._or_high = max(self._or_high, spy_price)
            self._or_low = min(self._or_low, spy_price)
        elif self.time.hour == 10 and not self._or_complete:
            self._or_complete = True

        if (in_non_bull and self._or_complete and
                self.time.hour == 10 and self.time.minute <= 30 and
                not self._orb_trade and self.trades_today < 2 and
                not self.portfolio[self.spy].is_long and
                not self.portfolio[self.inverse_spy].invested and
                self.imbalance_sma.is_ready):
            
            or_range = self._or_high - self._or_low
            imbalance = self.imbalance_sma.current.value
            
            if or_range > 0 and or_range / max(self._or_low, 1) > 0.0005:
                if spy_price < self._or_low * 0.999 and imbalance < -0.08:
                    self.set_holdings(self.inverse_spy, 0.20)
                    self.micro_entry_price = self.micro_lowest_price = spy_price
                    self._orb_trade = True
                    self.trades_today += 1
                    self._save_state()

        if in_non_bull:
            self.run_micro_execution(spy_price)

        # FIX: Only save state frequently during Live deployments to prevent backtest I/O throttling
        if self._deploy == "live" and (self.time - self._last_state_save).total_seconds() >= 300:
            self._save_state()

    def run_micro_execution(self, price):
        if not self.vwap_ind.is_ready or not self.price_std.is_ready: return

        # FIX: Removed the 15:00 restriction so regular trades trigger all day
        if self.trades_today >= 3:
            return

        vwap_val = self.vwap_ind.current.value
        std_dev = self.price_std.current.value if self.price_std.current.value > 0 else 1.0
        imbalance = self.imbalance_sma.current.value

        strong_bull = price > (vwap_val + std_dev) and imbalance > 0.15
        roc3_bearish = self.spy_roc_3.is_ready and self.spy_roc_3.current.value < 0
        strong_bear = price < (vwap_val - std_dev) and imbalance < -0.15 and roc3_bearish
        target_size = 0.50 if abs(imbalance) < 0.30 else 0.95

        spy_long = self.portfolio[self.spy].is_long
        sh_invested = self.portfolio[self.inverse_spy].invested
        
        if not spy_long and not sh_invested:
            if strong_bull:
                self.set_holdings(self.spy, target_size)
                self.micro_entry_price = self.micro_highest_price = price
                self.trades_today += 1
                self._save_state()
            elif strong_bear:
                self.set_holdings(self.inverse_spy, target_size)
                self.micro_entry_price = self.micro_lowest_price = price
                self.trades_today += 1
                self._save_state()
        else:
            if spy_long and strong_bear:
                self.liquidate(self.spy, "Flipping Bear: Liquidating SPY Long")
                self.set_holdings(self.inverse_spy, target_size)
                self.micro_entry_price = self.micro_lowest_price = price
                self.trades_today += 1
                self._save_state()
            elif sh_invested and strong_bull:
                self.liquidate(self.inverse_spy, "Flipping Bull: Liquidating SH")
                self.set_holdings(self.spy, target_size)
                self.micro_entry_price = self.micro_highest_price = price
                self.trades_today += 1
                self._save_state()

    def on_order_event(self, orderEvent):
        # We only care when an order actually fills
        if orderEvent.status == OrderStatus.FILLED:
            symbol = orderEvent.symbol
            
            # Check if it's one of your daily macro assets
            if symbol == self.tqqq or symbol == self.soxl:
                fill_price = orderEvent.fill_price
                fill_time = self.time
                engine_price = self.securities[symbol].price
                
                self.log(f"--- FAKE FILL CHECK ---")
                self.log(f"Fill Time: {fill_time}")
                self.log(f"Asset: {symbol}")
                self.log(f"Fill Price: ${fill_price}")
                self.log(f"Engine's Current Price: ${engine_price}")
                self.log(f"-----------------------")