from AlgorithmImports import *
import pandas as pd
import numpy as np
from scipy.cluster.hierarchy import dendrogram, linkage
from scipy.spatial.distance import squareform
from sklearn.covariance import OAS

# ==============================================================================
# 1. HIERARCHICAL RISK PARITY (HRP) CLASSES
# ==============================================================================

class HierarchicalRiskParity:
    def __init__(self):
        self.weights = list()
        self.seriated_correlations = None
        self.seriated_distances = None
        self.ordered_indices = None
        self.clusters = None

    @staticmethod
    def _tree_clustering(correlation, method='single'):
        distances = np.sqrt((1 - correlation).round(5) / 2)
        clusters = linkage(squareform(distances.values), method=method)
        return distances, clusters

    def _quasi_diagnalization(self, num_assets, curr_index):
        if curr_index < num_assets:
            return [curr_index]
        left = int(self.clusters[curr_index - num_assets, 0])
        right = int(self.clusters[curr_index - num_assets, 1])
        return (self._quasi_diagnalization(num_assets, left) + self._quasi_diagnalization(num_assets, right))

    def _get_seriated_matrix(self, assets, distances, correlations):
        ordering = assets[self.ordered_indices]
        seriated_distances = distances.loc[ordering, ordering]
        seriated_correlations = correlations.loc[ordering, ordering]
        return seriated_distances, seriated_correlations

    def _recursive_bisection(self, covariances, assets):
        self.weights = pd.Series(1.0, index=self.ordered_indices)
        clustered_alphas = [self.ordered_indices]

        while clustered_alphas:
            clustered_alphas = [cluster[start:end]
                                for cluster in clustered_alphas
                                for start, end in ((0, len(cluster) // 2), (len(cluster) // 2, len(cluster)))
                                if len(cluster) > 1]

            for subcluster in range(0, len(clustered_alphas), 2):
                left_cluster = clustered_alphas[subcluster]
                right_cluster = clustered_alphas[subcluster + 1]

                left_subcovar = covariances.iloc[left_cluster, left_cluster]
                inv_diag = 1 / np.diag(left_subcovar.values)
                parity_w = inv_diag * (1 / np.sum(inv_diag))
                left_cluster_var = np.dot(parity_w, np.dot(left_subcovar, parity_w))

                right_subcovar = covariances.iloc[right_cluster, right_cluster]
                inv_diag = 1 / np.diag(right_subcovar.values)
                parity_w = inv_diag * (1 / np.sum(inv_diag))
                right_cluster_var = np.dot(parity_w, np.dot(right_subcovar, parity_w))

                alloc_factor = 1 - left_cluster_var / (left_cluster_var + right_cluster_var)
                self.weights[left_cluster] *= alloc_factor
                self.weights[right_cluster] *= 1 - alloc_factor

        self.weights.index = assets[self.ordered_indices]
        self.weights = pd.DataFrame(self.weights).T

    @staticmethod
    def _calculate_returns(asset_prices, resample_by):
        if resample_by is not None:
            asset_prices = asset_prices.resample(resample_by).last()
        asset_returns = asset_prices.pct_change()
        asset_returns = asset_returns.dropna(how='all')
        return asset_returns

    @staticmethod
    def _shrink_covariance(covariance):
        oas = OAS()
        oas.fit(covariance)
        shrinked_covariance = oas.covariance_
        return pd.DataFrame(shrinked_covariance, index=covariance.columns, columns=covariance.columns)

    @staticmethod
    def _cov2corr(covariance):
        d_matrix = np.zeros_like(covariance)
        diagnoal_sqrt = np.sqrt(np.diag(covariance))
        np.fill_diagonal(d_matrix, diagnoal_sqrt)
        d_inv = np.linalg.inv(d_matrix)
        corr = np.dot(np.dot(d_inv, covariance), d_inv)
        corr = pd.DataFrame(corr, index=covariance.columns, columns=covariance.columns)
        return corr

class HierarchicalRiskParityModified(HierarchicalRiskParity):
    def allocate(self, asset_prices, covariance, resample_by='B', use_shrinkage=False):
        if not isinstance(asset_prices, pd.DataFrame):
            raise ValueError("Asset prices matrix must be a dataframe")
        if not isinstance(asset_prices.index, pd.DatetimeIndex):
            raise ValueError("Asset prices dataframe must be indexed by date.")

        asset_returns = self._calculate_returns(asset_prices, resample_by=resample_by)
        num_assets = asset_returns.shape[1]
        assets = asset_returns.columns
        
        cov = pd.DataFrame(covariance, columns=assets, index=assets)
        
        if use_shrinkage:
            cov = self._shrink_covariance(covariance=cov)
        corr = self._cov2corr(covariance=cov)

        distances, self.clusters = self._tree_clustering(correlation=corr)
        self.ordered_indices = self._quasi_diagnalization(num_assets, 2 * num_assets - 2)
        self.seriated_distances, self.seriated_correlations = self._get_seriated_matrix(assets=assets, distances=distances, correlations=corr)
        self._recursive_bisection(covariances=cov, assets=assets)


# ==============================================================================
# 2. QUANTCONNECT ALGORITHM: 60-STOCK HRP (QUARTERLY REBALANCE + SMART STOP)
# ==============================================================================

class UltimateFactorHRP(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2021, 1, 1)
        self.SetCash(100000)
        
        self.UniverseSettings.Resolution = Resolution.Minute
        self.AddUniverse(self.FundamentalSelection)
        self.spy = self.AddEquity("SPY", Resolution.Minute).Symbol
        
        # Target exact 60 stock portfolio
        self.final_count = 60     
        self.candidates = []
        
        # Execution & Protection Variables
        self.current_quarter = 0 # Tracks when to update the universe
        self.pending_weights = {}
        self.pending_liquidations = []
        self.weight_buffer = 0.02  # 2% buffer
        
        # Smart Stop Loss Variables
        self.trailing_stop_pct = 0.15 # 15% trailing stop loss
        self.high_water_marks = {}
        self.stop_loss_blacklist = set()

        self.SetWarmUp(60)
        self.SetBrokerageModel(BrokerageName.ALPACA)

        # QUARTERLY Rebalance Schedule (Triggered monthly, filtered internally to quarters)
        self.Schedule.On(self.DateRules.MonthStart(self.spy), 
                         self.TimeRules.AfterMarketOpen(self.spy, 30), 
                         self.QueueTrades)

    def FundamentalSelection(self, fundamental):
        """Builds a strict 60-stock Value universe: 36 Large, 18 Mid, 6 Small every Quarter"""
        
        # Calculate the current quarter (1, 2, 3, or 4)
        current_q = (self.Time.month - 1) // 3 + 1
        
        # Only update the universe if we have entered a new quarter
        if current_q == self.current_quarter:
            return Universe.Unchanged
            
        # Lock in the new quarter and reset the stop-loss blacklist
        self.current_quarter = current_q
        self.stop_loss_blacklist.clear()

        # 1. Base Filter (Ex-Financials/Real Estate)
        filtered = [f for f in fundamental if f.HasFundamentalData 
                                          and f.Price > 5 
                                          and f.MarketCap > 1e8 
                                          and f.ValuationRatios.PBRatio > 0
                                          and f.AssetClassification.MorningstarSectorCode != MorningstarSectorCode.FinancialServices
                                          and f.AssetClassification.MorningstarSectorCode != MorningstarSectorCode.RealEstate]
        
        if len(filtered) < 1000: return Universe.Unchanged

        # 2. Sort by size to define strata
        sorted_by_size = sorted(filtered, key=lambda x: x.MarketCap, reverse=True)
        
        large_caps = sorted_by_size[:200]       
        mid_caps = sorted_by_size[200:500]       
        small_caps = sorted_by_size[500:1000]    

        def get_value_stocks(bucket, count):
            # Sort by highest Book-to-Market
            sorted_bucket = sorted(bucket, key=lambda x: 1 / x.ValuationRatios.PBRatio, reverse=True)
            return [x.Symbol for x in sorted_bucket[:count]]

        # 3. Extract exact 60/30/10 ratio for a 60-stock portfolio
        large_value_symbols = get_value_stocks(large_caps, 36)  # 60% Large Cap
        mid_value_symbols = get_value_stocks(mid_caps, 18)      # 30% Mid Cap
        small_value_symbols = get_value_stocks(small_caps, 6)   # 10% Small Cap
        
        self.candidates = large_value_symbols + mid_value_symbols + small_value_symbols
        self.Debug(f"Q{self.current_quarter} Universe Generated: 60 Stratified Value Targets")
        return self.candidates

    def QueueTrades(self):
        if self.IsWarmingUp or not self.candidates: return

        # Enforce Quarterly Execution (Jan, Apr, Jul, Oct)
        if self.Time.month not in [1, 4, 7, 10]:
            return

        # Risk-Off Check
        spy_history = self.History(self.spy, 200, Resolution.Daily)
        if not spy_history.empty:
            spy_current = spy_history['close'].iloc[-1]
            spy_sma = spy_history['close'].mean()
            if spy_current < spy_sma:
                self.Liquidate() 
                self.pending_weights.clear() 
                self.Debug("Market Risk-Off: Liquidating to Cash")
                return

        # Ensure we don't buy stocks that were stopped out this quarter
        active_candidates = [c for c in self.candidates if c not in self.stop_loss_blacklist]

        # Fetch 60 days of history
        history = self.History(active_candidates, 60, Resolution.Daily)
        if history.empty: return
        
        prices = history['close'].unstack(level=0).ffill().dropna(axis=1)
        if prices.empty: return

        # Rank by momentum, retaining up to our 60 stock limit
        mom_scores = (prices.iloc[-1] / prices.iloc[0]) - 1
        top_symbols = mom_scores.sort_values(ascending=False).head(self.final_count).index.tolist()
        
        target_prices = prices[top_symbols]
        target_prices.index = pd.to_datetime(target_prices.index) 
        
        try:
            # 1. Base Math Data
            target_returns = target_prices.pct_change().dropna(how='all')
            oas = OAS()
            oas.fit(target_returns)
            target_cov = pd.DataFrame(oas.covariance_, index=target_returns.columns, columns=target_returns.columns)
            
            # 2. Execute Hierarchical Risk Parity (HRP)
            hrp = HierarchicalRiskParityModified()
            hrp.allocate(asset_prices=target_prices, covariance=target_cov, resample_by='B', use_shrinkage=True)
            hrp_weights = hrp.weights.iloc[0]
            
            # 3. Queue Liquidations for stocks no longer in target
            current_holdings = [x.Key for x in self.Portfolio if x.Value.Invested]
            self.pending_liquidations = [sym for sym in current_holdings if sym not in hrp_weights.index]
            
            # 4. Apply Weight Delta Buffer to Save Fees
            self.pending_weights.clear()
            for symbol, target_weight in hrp_weights.items():
                if target_weight <= 0.0001: continue # Ignore floating point dust
                
                current_weight = 0
                if self.Portfolio.ContainsKey(symbol) and self.Portfolio[symbol].Invested:
                    current_weight = self.Portfolio[symbol].HoldingsValue / self.Portfolio.TotalPortfolioValue
                
                weight_delta = abs(target_weight - current_weight)
                
                # Only queue the trade if the weight change is larger than 2%
                if weight_delta >= self.weight_buffer:
                    self.pending_weights[symbol] = round(target_weight, 4)
            
            # Fire first execution attempt
            self.ExecutePendingTrades()
                
        except Exception as e:
            self.Debug(f"HRP Execution Error: {e}")

    def OnData(self, data):
        """Intraday Smart Stop Loss & Fallback Execution Loop"""
        
        # 1. Intraday High-Water Mark Trailing Stop Loss
        for symbol in list(self.Portfolio.Keys):
            if not self.Portfolio[symbol].Invested: 
                # Cleanup high water marks for liquidated stocks
                if symbol in self.high_water_marks:
                    del self.high_water_marks[symbol]
                continue
                
            price = self.Securities[symbol].Price
            if price == 0: continue
            
            # Update high water mark
            hwm = self.high_water_marks.get(symbol, price)
            if price > hwm:
                self.high_water_marks[symbol] = price
                
            # Check trailing stop (15% drop from peak)
            elif price < hwm * (1 - self.trailing_stop_pct):
                self.SetHoldings(symbol, 0) # Instant Liquidation
                self.stop_loss_blacklist.add(symbol)
                del self.high_water_marks[symbol]
                self.Debug(f"Smart Stop Loss triggered: Liquidated {symbol.Value} at {price}. Blacklisted until next quarter.")

        # 2. Fallback loop: Retries stuck execution queues every 10 minutes
        if not self.pending_weights and not self.pending_liquidations: return
            
        if self.Time.minute % 10 == 0:
            self.ExecutePendingTrades()

    def ExecutePendingTrades(self):
        """Safely processes the queues only if valid live price data exists"""
        completed_liquidations = []
        for symbol in self.pending_liquidations:
            if self.Securities.ContainsKey(symbol) and self.Securities[symbol].Price > 0:
                self.SetHoldings(symbol, 0)
                completed_liquidations.append(symbol)
                
        for symbol in completed_liquidations:
            self.pending_liquidations.remove(symbol)

        completed_allocations = []
        for symbol, weight in self.pending_weights.items():
            if self.Securities.ContainsKey(symbol) and self.Securities[symbol].Price > 0:
                self.SetHoldings(symbol, weight)
                completed_allocations.append(symbol)

        for symbol in completed_allocations:
            del self.pending_weights[symbol]

from AlgorithmImports import *
import pandas as pd
import numpy as np
from scipy.cluster.hierarchy import dendrogram, linkage
from scipy.spatial.distance import squareform
from sklearn.covariance import OAS

# ==============================================================================
# MARCOS LÓPEZ DE PRADO HRP CLASSES (with Shrinkage & Modifications)
# ==============================================================================

class HierarchicalRiskParity:
    def __init__(self):
        self.weights = list()
        self.seriated_correlations = None
        self.seriated_distances = None
        self.ordered_indices = None
        self.clusters = None

    @staticmethod
    def _tree_clustering(correlation, method='single'):
        distances = np.sqrt((1 - correlation).round(5) / 2)
        clusters = linkage(squareform(distances.values), method=method)
        return distances, clusters

    def _quasi_diagnalization(self, num_assets, curr_index):
        if curr_index < num_assets:
            return [curr_index]
        left = int(self.clusters[curr_index - num_assets, 0])
        right = int(self.clusters[curr_index - num_assets, 1])
        return (self._quasi_diagnalization(num_assets, left) + self._quasi_diagnalization(num_assets, right))

    def _get_seriated_matrix(self, assets, distances, correlations):
        ordering = assets[self.ordered_indices]
        seriated_distances = distances.loc[ordering, ordering]
        seriated_correlations = correlations.loc[ordering, ordering]
        return seriated_distances, seriated_correlations

    def _recursive_bisection(self, covariances, assets):
        self.weights = pd.Series(1.0, index=self.ordered_indices)
        clustered_alphas = [self.ordered_indices]

        while clustered_alphas:
            clustered_alphas = [cluster[start:end]
                                for cluster in clustered_alphas
                                for start, end in ((0, len(cluster) // 2), (len(cluster) // 2, len(cluster)))
                                if len(cluster) > 1]

            for subcluster in range(0, len(clustered_alphas), 2):
                left_cluster = clustered_alphas[subcluster]
                right_cluster = clustered_alphas[subcluster + 1]

                left_subcovar = covariances.iloc[left_cluster, left_cluster]
                inv_diag = 1 / np.diag(left_subcovar.values)
                parity_w = inv_diag * (1 / np.sum(inv_diag))
                left_cluster_var = np.dot(parity_w, np.dot(left_subcovar, parity_w))

                right_subcovar = covariances.iloc[right_cluster, right_cluster]
                inv_diag = 1 / np.diag(right_subcovar.values)
                parity_w = inv_diag * (1 / np.sum(inv_diag))
                right_cluster_var = np.dot(parity_w, np.dot(right_subcovar, parity_w))

                alloc_factor = 1 - left_cluster_var / (left_cluster_var + right_cluster_var)
                self.weights[left_cluster] *= alloc_factor
                self.weights[right_cluster] *= 1 - alloc_factor

        self.weights.index = assets[self.ordered_indices]
        self.weights = pd.DataFrame(self.weights).T

    @staticmethod
    def _calculate_returns(asset_prices, resample_by):
        if resample_by is not None:
            asset_prices = asset_prices.resample(resample_by).last()
        asset_returns = asset_prices.pct_change()
        asset_returns = asset_returns.dropna(how='all')
        return asset_returns

    @staticmethod
    def _shrink_covariance(covariance):
        oas = OAS()
        oas.fit(covariance)
        shrinked_covariance = oas.covariance_
        return pd.DataFrame(shrinked_covariance, index=covariance.columns, columns=covariance.columns)

    @staticmethod
    def _cov2corr(covariance):
        d_matrix = np.zeros_like(covariance)
        diagnoal_sqrt = np.sqrt(np.diag(covariance))
        np.fill_diagonal(d_matrix, diagnoal_sqrt)
        d_inv = np.linalg.inv(d_matrix)
        corr = np.dot(np.dot(d_inv, covariance), d_inv)
        corr = pd.DataFrame(corr, index=covariance.columns, columns=covariance.columns)
        return corr


class HierarchicalRiskParityModified(HierarchicalRiskParity):
    def allocate(self, asset_prices, covariance, resample_by='B', use_shrinkage=False):
        if not isinstance(asset_prices, pd.DataFrame):
            raise ValueError("Asset prices matrix must be a dataframe")
        if not isinstance(asset_prices.index, pd.DatetimeIndex):
            raise ValueError("Asset prices dataframe must be indexed by date.")

        asset_returns = self._calculate_returns(asset_prices, resample_by=resample_by)
        num_assets = asset_returns.shape[1]
        assets = asset_returns.columns
        
        cov = pd.DataFrame(covariance, columns=assets, index=assets)
        
        if use_shrinkage:
            cov = self._shrink_covariance(covariance=cov)
        corr = self._cov2corr(covariance=cov)

        distances, self.clusters = self._tree_clustering(correlation=corr)
        self.ordered_indices = self._quasi_diagnalization(num_assets, 2 * num_assets - 2)
        self.seriated_distances, self.seriated_correlations = self._get_seriated_matrix(assets=assets, distances=distances, correlations=corr)
        self._recursive_bisection(covariances=cov, assets=assets)


# ==============================================================================
# CUSTOM SECURITY INITIALIZER (For Slippage & Brokerage Models)
# ==============================================================================

class CustomSecurityInitializer(BrokerageModelSecurityInitializer):
    def __init__(self, brokerage_model, security_seeder):
        super().__init__(brokerage_model, security_seeder)

    def Initialize(self, security):
        # 1. Apply the default Alpaca brokerage models (fees, margin, etc.)
        super().Initialize(security)
        # 2. Apply our custom 0.1% slippage penalty on top
        security.SetSlippageModel(ConstantSlippageModel(0.001))


# ==============================================================================
# QUANTCONNECT ALGORITHM
# ==============================================================================

class UltimateFactorHRP(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2021, 1, 1)
        self.SetCash(100000)
        
        # 1. Environment & Universe
        self.UniverseSettings.Resolution = Resolution.Minute
        self.AddUniverse(self.FundamentalSelection)
        self.spy = self.AddEquity("SPY", Resolution.Minute).Symbol
        
        # 2. Strategy Variables
        self.max_candidates = 50 
        self.final_count = 15     
        self.candidates = []
        self.weight_buffer = 0.02  # 2% weight delta buffer
        
        # 3. State Variables
        self.next_universe_time = self.Time
        self.pending_weights = {}
        self.pending_liquidations = []

        self.SetWarmUp(60)
        
        # 4. Brokerage & Custom Security Initializer
        self.SetBrokerageModel(BrokerageName.ALPACA)
        self.SetSecurityInitializer(CustomSecurityInitializer(self.BrokerageModel, SecuritySeeder.Null))

        # 5. Scheduling (Daily Rebalance)
        self.Schedule.On(self.DateRules.EveryDay(self.spy), 
                         self.TimeRules.AfterMarketOpen(self.spy, 30), 
                         self.QueueTrades)

    def FundamentalSelection(self, fundamental):
        """Locks universe updates to a weekly cycle"""
        if self.Time < self.next_universe_time:
            return Universe.Unchanged
            
        self.next_universe_time = self.Time + timedelta(days=7)

        filtered = [f for f in fundamental if f.HasFundamentalData and f.Price > 5 and f.MarketCap > 1e8]
        sorted_by_cap = sorted(filtered, key=lambda x: x.MarketCap, reverse=True)
        self.candidates = [x.Symbol for x in sorted_by_cap[:self.max_candidates]]
        
        return self.candidates

    def QueueTrades(self):
        """Calculates daily targets and pushes them into the execution queues"""
        if self.IsWarmingUp or not self.candidates: return

        # Risk-Off Check
        spy_history = self.History(self.spy, 200, Resolution.Daily)
        if not spy_history.empty:
            spy_current = spy_history['close'].iloc[-1]
            spy_sma = spy_history['close'].mean()
            if spy_current < spy_sma:
                self.Liquidate() 
                self.pending_weights.clear() 
                self.Debug("Market Risk-Off: Liquidating to Cash")
                return

        # Fetch candidate history
        history = self.History(self.candidates, 60, Resolution.Daily)
        if history.empty: return
        
        prices = history['close'].unstack(level=0).ffill().dropna(axis=1)
        if prices.empty or len(prices.columns) < self.final_count: return

        # Target Top 15 Momentum Stocks
        mom_scores = (prices.iloc[-1] / prices.iloc[0]) - 1
        top_symbols = mom_scores.sort_values(ascending=False).head(self.final_count).index.tolist()
        
        target_prices = prices[top_symbols]
        target_prices.index = pd.to_datetime(target_prices.index) 
        
        try:
            # HRP Math & Allocation
            target_returns = target_prices.pct_change().dropna(how='all')
            target_cov = target_returns.cov()
            
            hrp = HierarchicalRiskParityModified()
            hrp.allocate(asset_prices=target_prices, covariance=target_cov, resample_by='B', use_shrinkage=True)
            
            hrp_weights_series = hrp.weights.iloc[0]
            
            # Queue Liquidations
            current_holdings = [x.Key for x in self.Portfolio if x.Value.Invested]
            self.pending_liquidations = [sym for sym in current_holdings if sym not in hrp_weights_series.index]
            
            # Apply Weight Delta Filter
            self.pending_weights.clear()
            for symbol, target_weight in hrp_weights_series.items():
                current_weight = 0
                if self.Portfolio.ContainsKey(symbol) and self.Portfolio[symbol].Invested:
                    current_weight = self.Portfolio[symbol].HoldingsValue / self.Portfolio.TotalPortfolioValue
                
                weight_delta = abs(target_weight - current_weight)
                
                if weight_delta >= self.weight_buffer:
                    self.pending_weights[symbol] = target_weight
            
            # Trigger first attempt
            self.ExecutePendingTrades()
                
        except Exception as e:
            self.Debug(f"Trade Execution Error: {e}")

    def OnData(self, data):
        """Fallback loop: Retries stuck trades every 10 minutes"""
        if not self.pending_weights and not self.pending_liquidations: return
            
        if self.Time.minute % 10 == 0:
            self.ExecutePendingTrades()

    def ExecutePendingTrades(self):
        """Safely processes the queues only if valid live price data exists"""
        completed_liquidations = []
        for symbol in self.pending_liquidations:
            if self.Securities.ContainsKey(symbol) and self.Securities[symbol].Price > 0:
                self.SetHoldings(symbol, 0)
                completed_liquidations.append(symbol)
                
        for symbol in completed_liquidations:
            self.pending_liquidations.remove(symbol)

        completed_allocations = []
        for symbol, weight in self.pending_weights.items():
            if self.Securities.ContainsKey(symbol) and self.Securities[symbol].Price > 0:
                self.SetHoldings(symbol, weight)
                completed_allocations.append(symbol)

        for symbol in completed_allocations:
            del self.pending_weights[symbol]