# region imports
from AlgorithmImports import *
import numpy as np
# endregion

class AlphaTrendComposite(QCAlgorithm):

    def initialize(self):
        # 1. Setup: 5 Year backtest to capture 2020 (Crash) and 2022 (Inflation)
        self.set_start_date(2019, 1, 1)
        self.set_end_date(2024, 6, 1)
        self.set_cash(100000)
        self.settings.rebalance_portfolio_on_insight_changes = False
        self.settings.rebalance_portfolio_on_security_changes = False
        
        # 2. Universe: The Dow 30 (Blue Chips) + Big Tech
        # We manually list them to ensure no data errors with delisted stocks
        self.tickers = [
            'AAPL', 'AXP', 'BA', 'CAT', 'CSCO', 'CVX', 'DIS', 'GS', 'HD', 'IBM', 
            'INTC', 'JPM', 'KO', 'MCD', 'MMM', 'MRK', 'MSFT', 'NKE', 'PG', 'TRV', 
            'UNH', 'V', 'VZ', 'WMT', 'XOM', 'AMZN', 'GOOG', 'META', 'NVDA', 'AMD'
        ]
        
        self.symbols = [self.add_equity(t, Resolution.DAILY).symbol for t in self.tickers]

        # 3. The Benchmark (SPY) for Alpha calculation and Trend Filtering
        self.spy = self.add_equity('SPY', Resolution.DAILY).symbol
        
        # 4. The Trend Filter (200 Day Moving Average)
        self.spy_ma = self.sma(self.spy, 200, Resolution.DAILY)
        
        # 5. Parameters
        self.lookback = 62 # ~3 months for regression stability
        self.top_n = 3     # Concentrate on the top 3 winners
        
        # 6. Schedule: Rebalance at the start of every month
        self.schedule.on(self.date_rules.month_start(self.spy),
                         self.time_rules.after_market_open(self.spy, 30),
                         self.rebalance)
                         
        self.set_warm_up(200)

    def rebalance(self):
        if self.is_warming_up: return

        # --- SAFETY CHECK ---
        # If SPY is below its 200-day average, the market is in a downtrend.
        # We go to CASH. This saves you from -50% drawdowns.
        if self.securities[self.spy].price < self.spy_ma.current.value:
            self.liquidate()
            self.plot("State", "Cash", 1)
            return
        
        # If we are here, the market is safe. Let's find the best stocks.
        self.plot("State", "Invested", 1)
        
        # --- ALPHA CALCULATION ---
        history = self.history(self.symbols + [self.spy], self.lookback, Resolution.DAILY).close.unstack(level=0)
        
        if history.empty: return
        
        alphas = {}
        
        # Get Benchmark (SPY) returns
        if self.spy not in history: return
        spy_returns = history[self.spy].pct_change().dropna()
        
        for symbol in self.symbols:
            if symbol not in history: continue
            
            # Get Stock returns
            stock_returns = history[symbol].pct_change().dropna()
            
            # Align data lengths
            min_len = min(len(spy_returns), len(stock_returns))
            if min_len < 30: continue # Not enough data
            
            y = stock_returns.tail(min_len).values
            x = spy_returns.tail(min_len).values
            
            # Linear Regression: Stock = Alpha + Beta * SPY
            # We stack [SPY, 1] to get the intercept (Alpha)
            A = np.vstack([x, np.ones(len(x))]).T
            
            try:
                # result[0] is [slope, intercept]. Intercept is Alpha.
                result = np.linalg.lstsq(A, y, rcond=None)
                alpha = result[0][1] # The intercept
                alphas[symbol] = alpha
            except:
                continue
                
        # --- SELECTION & EXECUTION ---
        # Sort by Alpha (Highest first)
        sorted_by_alpha = sorted(alphas.items(), key=lambda x: x[1], reverse=True)
        
        # Select Top N
        selected_symbols = [x[0] for x in sorted_by_alpha[:self.top_n]]
        
        if not selected_symbols: return
        
        # Allocate 100% equity divided equally among the winners
        # We liquidate others first to free up cash
        self.set_holdings([PortfolioTarget(s, 1.0/len(selected_symbols)) for s in selected_symbols], True)