# region imports
from AlgorithmImports import *
# endregion
import numpy as np
import pandas as pd
from scipy.optimize import minimize
import statsmodels.api as sm

class DynamicBlackLittermanWithWarmUp(QCAlgorithm):
    def Initialize(self):
        # 1. Basics & WarmUp
        self.SetStartDate(2020, 1, 1)
        self.SetCash(100000)
        # Warm up the algorithm for 252 trading days to fill history
        self.SetWarmUp(252, Resolution.Daily)
        
        # 2. Universe Settings
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelection, self.FineSelection)
        
        # 3. Reference Symbols
        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol
        self.shy = self.AddEquity("SHY", Resolution.Daily).Symbol # Defensive Bond ETF
        
        self.lookback = 252 
        self.active_symbols = []
        self.Settings.MinAbsolutePortfolioTargetPercentage = 0.001
        
        # 4. Schedule Rebalance
        self.Schedule.On(self.DateRules.MonthStart("SPY"), 
                         self.TimeRules.AfterMarketOpen("SPY", 30), 
                         self.Rebalance)

        self.set_brokerage_model(BrokerageName.ALPACA) 

    def CoarseSelection(self, coarse):
        # Liquidity filter
        sorted_by_volume = sorted([x for x in coarse if x.HasFundamentalData and x.Price > 10], 
                                  key=lambda x: x.DollarVolume, reverse=True)
        return [x.Symbol for x in sorted_by_volume[:100]]

    def FineSelection(self, fine):
        # Select top 25 by Market Cap
        sorted_by_mcap = sorted(fine, key=lambda x: x.MarketCap, reverse=True)
        return [x.Symbol for x in sorted_by_mcap[:25]]

    def OnSecuritiesChanged(self, changes):
        for security in changes.AddedSecurities:
            if security.Symbol not in self.active_symbols:
                self.active_symbols.append(security.Symbol)
        for security in changes.RemovedSecurities:
            if security.Symbol in self.active_symbols:
                self.active_symbols.remove(security.Symbol)

    def Rebalance(self):
        # Skip if still warming up
        if self.IsWarmingUp: return
        if not self.active_symbols: return

        # 1. Market Regime Filter (200-day SMA on SPY)
        spy_hist = self.History(self.spy, 200, Resolution.Daily)
        if spy_hist.empty: return
        
        current_spy_price = self.Securities[self.spy].Price
        spy_sma = spy_hist['close'].mean()

        # If SPY is below 200 SMA, move to SHY and exit others
        if current_spy_price < spy_sma:
            self.Log("Bear Market Detected: Moving to SHY")
            self.SetHoldings(self.shy, 1.0)
            for symbol in self.active_symbols:
                if symbol != self.spy: self.Liquidate(symbol)
            return

        # 2. Fetch History for BL Model
        history = self.History(self.active_symbols + [self.spy], self.lookback, Resolution.Daily)
        if history.empty: return
        history_unstacked = history['close'].unstack(level=0)
        returns_df = history_unstacked.pct_change().dropna()
        
        valid_symbols = [s for s in self.active_symbols if str(s) in returns_df.columns]
        if len(valid_symbols) < 10: return
        
        cov_matrix = returns_df[list(map(str, valid_symbols))].cov().values * 252
        market_caps = [self.Securities[s].Fundamentals.MarketCap for s in valid_symbols]
        mcap_weights = np.array(market_caps) / sum(market_caps)

        # 3. Hybrid Views (Regression Alpha + RSI Adjustment)
        spy_ret = returns_df[str(self.spy)]
        views = []
        for symbol in valid_symbols:
            s_ret = returns_df[str(symbol)]
            model = sm.OLS(s_ret, sm.add_constant(spy_ret)).fit()
            alpha = model.params[0] * 252
            
            # 14-day RSI
            prices = history_unstacked[str(symbol)].tail(14)
            diff = prices.diff()
            gain = (diff.where(diff > 0, 0)).mean()
            loss = (-diff.where(diff < 0, 0)).mean()
            rsi = 100 - (100 / (1 + (gain / (loss + 1e-9))))
            
            # Penalize overbought, reward oversold
            view_adj = alpha * 0.5 if rsi > 70 else (alpha * 1.5 if rsi < 30 else alpha)
            views.append(view_adj)

        # 4. Black-Litterman Optimization
        opt = BlackLittermanOptimizer(valid_symbols, mcap_weights, cov_matrix, views)
        target_dict = opt.calculate_weights()

        # 5. Build Portfolio Targets
        targets = []
        # Liquidate SHY if we were in defensive mode
        if self.Portfolio[self.shy].Invested:
            targets.append(PortfolioTarget(self.shy, 0))

        for symbol in valid_symbols:
            weight = target_dict.get(symbol, 0)
            targets.append(PortfolioTarget(symbol, weight if weight > 0.005 else 0))

        # Liquidate securities no longer in universe
        for holding in self.Portfolio.Values:
            if holding.Symbol not in valid_symbols and holding.Symbol != self.shy and holding.Invested:
                targets.append(PortfolioTarget(holding.Symbol, 0))

        if targets:
            self.SetHoldings(targets)

class BlackLittermanOptimizer:
    def __init__(self, symbols, mcap_weights, cov, views, tau=0.025, delta=2.5):
        self.symbols, self.w_mkt, self.sigma, self.Q = symbols, mcap_weights, cov, np.array(views)
        self.tau, self.delta = tau, delta

    def calculate_weights(self):
        n = len(self.symbols)
        pi = self.delta * np.dot(self.sigma, self.w_mkt)
        P = np.eye(n)
        omega = np.diag(np.diag(np.dot(np.dot(P, self.tau * self.sigma), P.T)))
        
        # BL Posterior
        term1 = np.linalg.inv(np.linalg.inv(self.tau * self.sigma) + np.dot(np.dot(P.T, np.linalg.inv(omega)), P))
        term2 = np.dot(np.linalg.inv(self.tau * self.sigma), pi) + np.dot(np.dot(P.T, np.linalg.inv(omega)), self.Q)
        er = np.dot(term1, term2)

        def obj(w):
            port_return = np.dot(w, er)
            port_vol = np.sqrt(np.dot(w.T, np.dot(self.sigma, w)))
            return -port_return / (port_vol + 1e-9)
        
        res = minimize(obj, self.w_mkt, bounds=[(0, 0.15) for _ in range(n)], 
                       constraints={'type': 'eq', 'fun': lambda x: np.sum(x) - 1})
        return dict(zip(self.symbols, res.x))