from AlgorithmImports import *
import numpy as np
import pandas as pd
from datetime import timedelta

class LargeCapDriftDispersionAlgorithm(QCAlgorithm):

    def Initialize(self):
        # 1. Backtest Settings
        self.SetStartDate(2024, 1, 1)
        self.SetCash(25000)
        self.SetTimeZone(TimeZones.Chicago)

        self.resolution = Resolution.Minute
        self.spy = self.AddEquity("SPY", self.resolution).Symbol
        self.SetBenchmark("SPY")

        # 2. Brokerage Model (Margin is required for shorting)
        self.SetBrokerageModel(
            BrokerageName.InteractiveBrokersBrokerage, 
            AccountType.Margin
        )

        # 3. Universe Selection (Large Cap Liquid Stocks)
        self.UniverseSettings.Resolution = self.resolution
        tickers = [
            "AAPL","MSFT","AMZN","NVDA","META","GOOGL","BRK.B","JPM","JNJ","V","PG","HD","UNH","MA",
            "XOM","CVX","MRK","KO","PEP","ABBV","AVGO","COST","WMT","MCD","CSCO","IBM","CAT","BA","DIS","GS",
            "AMD","TSLA","GOOG","SMCI","PLTR","NFLX","CRM","ADBE","INTU",
            "NOW","ORCL","QCOM","TXN","MU","PANW", "ABT", "T", "TMUS", "GEV", "ANET", "ACN", "PGR", "CEG", "ASML",
            "SCHW", "UL", "COF", "TD", "PFE", "SYK", "BBVA", "BN", "ARM", "PLD", "MDT", "CME", "SBUX", "DASH"
        ]
        symbols = [Symbol.Create(t, SecurityType.Equity, Market.USA) for t in tickers]
        self.SetUniverseSelection(ManualUniverseSelectionModel(symbols))

        # 4. Alpha Model: Systematic Beta Selection
        self.SetAlpha(TridentBetaNeutralAlpha(
            reference_ticker="SPY",
            lookback=252,      # 1 Year of trading days
            sma_period=200,    # Market Regime Filter
            resolution=self.resolution,
            n_each_side=2,     # Top 2 Long, Top 2 Short
            emit_after_open_min=245
        ))

        # 5. Portfolio Construction: Equal Weighting (10% Gross Exposure)
        self.SetPortfolioConstruction(GrossCappedInsightWeightingPCM(max_gross=0.10))
        self.set_warm_up(252*3)
        # 6. Execution Model: 45-minute delay on short entries
        self.SetExecution(StaggeredLongShortExecutionModel(delay_minutes=45))

        # 7. Charting: Real-time P/L visualization
        pnl_chart = Chart('Trade Profits')
        pnl_chart.AddSeries(Series('Realized P/L ($)', SeriesType.Bar, '$'))
        self.AddChart(pnl_chart)

    def OnOrderEvent(self, orderEvent: OrderEvent):
            """Logs Realized P/L and correctly identifies Long vs Short sleeves."""
            if orderEvent.Status != OrderStatus.Filled:
                return

            symbol = orderEvent.Symbol
            
            # We only log when the position is fully closed (not invested)
            if not self.Portfolio[symbol].Invested:
                closed_trades = list(self.TradeBuilder.ClosedTrades)
                if not closed_trades:
                    return
                    
                last_t = closed_trades[-1]
                
                # Ensure we are looking at the trade for this specific symbol
                if last_t.Symbol == symbol:
                    # 1. Get the P/L
                    pnl = self.Portfolio[symbol].LastTradeProfit
                    if pnl == 0: pnl = getattr(last_t, 'Profit', 0)

                    # 2. THE FIX: Correctly identify the sleeve using TradeDirection
                    # TradeDirection.Long means you bought first.
                    # TradeDirection.Short means you sold first.
                    if last_t.Direction == TradeDirection.Long:
                        sleeve = "LONG (Low Beta)"
                    else:
                        sleeve = "SHORT (High Beta)"
                    
                    # 3. Calculate Math for the log
                    entry_p = last_t.EntryPrice
                    exit_p = last_t.ExitPrice
                    total_cost = entry_p * abs(last_t.Quantity)
                    pnl_pct = (pnl / total_cost) * 100 if total_cost > 0 else 0

                    # 4. Final Accurate Log
                    if pnl != 0:
                        log_msg = (f"[{sleeve}] Closed {symbol} | "
                                f"P/L: ${pnl:.2f} ({pnl_pct:.2f}%) | "
                                f"Entry: {entry_p:.2f} | Exit: {exit_p:.2f}")
                        
                        self.Log(log_msg)
                        self.Plot('Trade Profits', 'Realized P/L ($)', pnl)

class StaggeredLongShortExecutionModel(ExecutionModel):
    def __init__(self, delay_minutes: int = 20):
        self.delay = timedelta(minutes=int(delay_minutes))
        self.pending = {}

    def Execute(self, algorithm: QCAlgorithm, targets):
        now = algorithm.Time
        
        # Handle delayed short entries
        due = [sym for sym, rec in self.pending.items() if rec["time"] <= now]
        for sym in due:
            rec = self.pending.pop(sym, None)
            if rec: algorithm.SetHoldings(sym, float(rec["percent"]), tag=rec["tag"])

        for t in targets:
            sym = t.Symbol
            if sym not in algorithm.Securities: continue
            
            price = algorithm.Securities[sym].Price
            total_value = algorithm.Portfolio.TotalPortfolioValue
            if price <= 0 or total_value <= 0: continue
            
            target_percent = (float(t.Quantity) * float(price)) / float(total_value)

            # Liquidation
            if abs(target_percent) < 1e-6:
                if algorithm.Portfolio[sym].Invested:
                    algorithm.Liquidate(sym, tag="Exit: Selection Changed")
                if sym in self.pending: self.pending.pop(sym, None)
                continue

            insight_tag = "Alpha Trade"
            active_insights = algorithm.Insights.GetActiveInsights(now)
            for insight in active_insights:
                if insight.Symbol == sym:
                    insight_tag = insight.Tag
                    break

            if target_percent < 0: # Delayed Shorts
                self.pending[sym] = {"time": now + self.delay, "percent": float(target_percent), "tag": insight_tag}
            else: # Immediate Longs
                algorithm.SetHoldings(sym, float(target_percent), tag=insight_tag)

class TridentBetaNeutralAlpha(AlphaModel):
    def __init__(self, reference_ticker, lookback, sma_period, resolution, n_each_side=2, emit_after_open_min=245):
        super().__init__()
        self.reference_ticker = reference_ticker
        self.lookback = lookback
        self.sma_period = sma_period
        self.resolution = resolution
        self.n_each_side = n_each_side
        self.emit_after_open_min = emit_after_open_min
        self.cache = {}
        self.reference = None
        self.last_emit_date = None
        self.ref_sma200 = None

    def OnSecuritiesChanged(self, algorithm, changes):
        if self.reference is None:
            ref_sec = algorithm.AddEquity(self.reference_ticker, self.resolution)
            self.reference = ref_sec.Symbol
            self.ref_sma200 = algorithm.SMA(self.reference, self.sma_period, Resolution.Daily)
        self._init_symboldata_from_history(algorithm, [x.Symbol for x in changes.AddedSecurities])

    def Update(self, algorithm, data):
        if algorithm.IsWarmingUp or not self._after_open(algorithm): return []
        today = algorithm.Time.date()
        if self.last_emit_date == today: return []

        regime = "BULL" if algorithm.Securities[self.reference].Price >= self.ref_sma200.Current.Value else "BEAR"
        ref_sd = self.cache.get(self.reference)
        if not ref_sd: return []
        ref_series = ref_sd.GetSeries()
        
        betas = {}
        for sym, sd in self.cache.items():
            if sym == self.reference: continue
            beta = sd.compute_beta_vs(ref_series)
            if beta is not None: betas[sym] = beta

        sorted_betas = sorted(betas.items(), key=lambda x: x[1])
        long_candidates = sorted_betas[:self.n_each_side]
        short_candidates = sorted_betas[-self.n_each_side:]

        insights = []
        period = timedelta(days=1)
        
        for sym, beta in long_candidates:
            tag = f"LONG | Beta: {beta:.2f} | {regime}"
            insights.append(Insight.Price(sym, period, InsightDirection.Up, None, None, None, 1.0, tag))
        
        for sym, beta in short_candidates:
            tag = f"SHORT | Beta: {beta:.2f} | {regime}"
            insights.append(Insight.Price(sym, period, InsightDirection.Down, None, None, None, 1.0, tag))

        self.last_emit_date = today
        return insights

    def _after_open(self, algorithm):
        sec = algorithm.Securities[self.reference]
        open_time = sec.Exchange.Hours.GetNextMarketOpen(algorithm.Time - timedelta(days=1), False)
        return algorithm.Time >= open_time + timedelta(minutes=self.emit_after_open_min)

    def _init_symboldata_from_history(self, algorithm, symbols):
        hist = algorithm.History(symbols, self.lookback + self.sma_period, self.resolution)
        if hist.empty: return
        close_df = hist["close"].unstack(level=0)
        for sym in symbols:
            if sym in close_df.columns:
                self.cache[sym] = SymbolData(algorithm, sym, self.lookback, self.sma_period, self.resolution, close_df[sym])

class GrossCappedInsightWeightingPCM(PortfolioConstructionModel):
    def __init__(self, max_gross=0.10):
        super().__init__()
        self.max_gross = max_gross

    def CreateTargets(self, algorithm, insights):
        if not insights: return []
        weight = self.max_gross / len(insights)
        return [PortfolioTarget.Percent(algorithm, i.Symbol, (1 if i.Direction == InsightDirection.Up else -1) * weight) for i in insights]

class SymbolData:
    def __init__(self, algorithm, symbol, lookback, sma_period, resolution, history_series):
        self.window = RollingWindow[IndicatorDataPoint](lookback)
        self.sma = algorithm.SMA(symbol, sma_period, resolution)
        self.sma.Updated += lambda s, u: self.window.Add(u)
        for t, p in history_series.ffill().items(): self.sma.Update(t, p)

    def GetSeries(self):
        if not self.window.IsReady: return None
        return pd.Series({x.EndTime: x.Value for x in self.window}).iloc[::-1]

    def compute_beta_vs(self, reference_series):
        sym_series = self.GetSeries()
        if sym_series is None or reference_series is None: return None
        df = pd.concat([sym_series.rename("s"), reference_series.rename("r")], axis=1).dropna()
        if len(df) < 50: return None
        s_ret, r_ret = np.diff(df["s"]) / df["s"][:-1], np.diff(df["r"]) / df["r"][:-1]
        n = min(len(s_ret), len(r_ret))
        A = np.vstack([s_ret[:n], np.ones(n)]).T
        beta, _ = np.linalg.lstsq(A, r_ret[:n], rcond=None)[0]
        return float(beta)

# region imports
from AlgorithmImports import *
# endregion
OrderTypeKeys = [
    'Market', 'Limit', 'StopMarket', 'StopLimit', 'MarketOnOpen',
    'MarketOnClose', 'OptionExercise',
]

OrderTypeCodes = dict(zip(range(len(OrderTypeKeys)), OrderTypeKeys))