from AlgorithmImports import *
import pandas as pd
import numpy as np
from sklearn.covariance import OAS

# ==============================================================================
# MLFINLAB: MEAN VARIANCE OPTIMISATION (MVO) CLASSES
# ==============================================================================

class MeanVarianceOptimisation:
    def __init__(self):
        self.weights = list()

    def allocate(self, asset_prices, solution='inverse_variance', resample_by='B'):
        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)
        assets = asset_prices.columns

        if solution == 'inverse_variance':
            cov = asset_returns.cov()
            self.weights = self._inverse_variance(covariance=cov)
        else:
            raise ValueError("Unknown solution string specified. Supported solutions - inverse_variance.")
        
        self.weights = pd.DataFrame(self.weights)
        self.weights.index = assets
        self.weights = self.weights.T

    @staticmethod
    def _calculate_returns(asset_prices, resample_by):
        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 _inverse_variance(covariance):
        ivp = 1. / np.diag(covariance)
        ivp /= ivp.sum()
        return ivp


class MeanVarianceOptimisationModified(MeanVarianceOptimisation):
    def allocate(self, asset_prices, covariance, solution='inverse_variance', resample_by='B'):
        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)
        assets = asset_prices.columns

        if solution == 'inverse_variance':
            # Using the modified covariance matrix (e.g., Shrinkage)
            cov = pd.DataFrame(covariance, columns=assets, index=assets)
            self.weights = self._inverse_variance(covariance=cov)
        else:
            raise ValueError("Unknown solution string specified. Supported solutions - inverse_variance.")
        
        self.weights = pd.DataFrame(self.weights)
        self.weights.index = assets
        self.weights = self.weights.T

# ==============================================================================
# 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):
        super().Initialize(security)
        # Apply 0.1% custom slippage penalty
        security.SetSlippageModel(ConstantSlippageModel(0.001))

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

class UltimateFactorMVO(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
        
        self.max_candidates = 50 
        self.final_count = 25     
        self.candidates = []
        
        self.next_universe_time = self.Time
        self.pending_weights = {}
        self.pending_liquidations = []

        self.SetWarmUp(60)
        
        self.SetBrokerageModel(BrokerageName.ALPACA)
        self.SetSecurityInitializer(CustomSecurityInitializer(self.BrokerageModel, SecuritySeeder.Null))

        # Weekly Rebalance
        self.Schedule.On(self.DateRules.WeekStart(self.spy), 
                         self.TimeRules.AfterMarketOpen(self.spy, 30), 
                         self.QueueTrades)

    def FundamentalSelection(self, fundamental):
        if self.Time < self.next_universe_time:
            return Universe.Unchanged
            
        self.next_universe_time = self.Time + timedelta(days=7)

        # 1. Base Filter: Price/Cap, valid P/B ratio, and EXCLUDE 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 Market Cap to define boundaries
        sorted_by_size = sorted(filtered, key=lambda x: x.MarketCap, reverse=True)
        
        # 3. Stratify into Cap Buckets
        large_caps = sorted_by_size[:200]        # Top 200 Largest Companies
        mid_caps = sorted_by_size[200:500]       # Next 300 Companies
        small_caps = sorted_by_size[500:1000]    # Next 500 Companies

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

        # 4. Extract Value Stocks from each strata (Total = 50 Candidates)
        large_value_symbols = get_value_stocks(large_caps, 20)  # 40% Large-Value
        mid_value_symbols = get_value_stocks(mid_caps, 15)      # 30% Mid-Value
        small_value_symbols = get_value_stocks(small_caps, 15)  # 30% Small-Value
        
        self.candidates = large_value_symbols + mid_value_symbols + small_value_symbols
        return self.candidates

    def QueueTrades(self):
        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 from our Fama-French Stratified Universe
        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:
            # Generate Sklearn Shrinkage Covariance Matrix
            # (Note: Inverse Variance only uses the diagonal, but we retain OAS for clean volatility estimates)
            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)
            
            # Execute Mean-Variance Optimization (MVO) -> Inverse Variance
            mvo = MeanVarianceOptimisationModified()
            mvo.allocate(asset_prices=target_prices, covariance=target_cov, resample_by='B', solution='inverse_variance')
            
            # Extract optimal target weights
            mvo_weights_series = mvo.weights.iloc[0]
            
            # Clean floating point dust
            clean_weights = {}
            for sym, w in mvo_weights_series.items():
                if w > 0.0001:  
                    clean_weights[sym] = round(w, 4) 
            
            # 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 clean_weights]
            
            # Set Target Weights Directly
            self.pending_weights = clean_weights
            
            # 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]