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 = 30     # Increased to 30 stocks
        self.candidates = []
        
        # 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 (Weekly Rebalance)
        self.Schedule.On(self.DateRules.MonthStart(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 weekly 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 25 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]
            
            # Set Target Weights Directly (Buffer Removed)
            self.pending_weights = hrp_weights_series.to_dict()
            
            # 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]

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]