from AlgorithmImports import *
import pandas as pd
import numpy as np
import scipy.cluster.hierarchy as sch

class UltimateFactorHRP(QCAlgorithm):
    def Initialize(self):
        # 1. Basics
        self.SetStartDate(2021, 1, 1)
        self.SetCash(100000)
        
        # 2. Universe: Use a standard liquid universe
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.FundamentalSelection)
        
        # 3. Benchmark for Risk-Off (SPY)
        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol
        
        # 4. State Variables
        self.max_candidates = 50 
        self.final_count = 10     
        self.rebalance_time = datetime.min
        self.candidates = []

        # 5. Warmup
        self.SetWarmUp(60)

    def FundamentalSelection(self, fundamental):
        # Monthly Rebalance Trigger
        if self.Time < self.rebalance_time: return Universe.Unchanged
        self.rebalance_time = self.Time + timedelta(days=30)
        
        # Broad filter: Just need Price and MarketCap to start
        filtered = [f for f in fundamental if f.HasFundamentalData and f.Price > 5 and f.MarketCap > 1e8]
        
        # Sort by a very reliable factor (Market Cap / PE Ratio) to ensure we get a list
        # We will do the deep 12-factor analysis in OnData
        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]]
        
        self.Debug(f"Universe selected {len(self.candidates)} candidates at {self.Time}")
        return self.candidates

    def OnData(self, data):
        if self.IsWarmingUp: return
        
        # If we have candidates but no investments, trigger the HRP logic
        if not self.Portfolio.Invested and len(self.candidates) > 0:
            self.Log("Triggering Portfolio Construction...")
            self.PerformTrade()

    def PerformTrade(self):
        # 1. Get History for HRP
        history = self.History(self.candidates, 60, Resolution.Daily)
        if history.empty: 
            self.Debug("History was empty - skipping trade")
            return
        
        prices = history['close'].unstack(level=0)
        returns = prices.pct_change().dropna()
        
        if returns.empty: return

        # 2. Factor Scoring (Momentum & Risk)
        # We pick the top 10 by 1-month momentum from our 30 candidates
        mom_scores = (prices.iloc[-1] / prices.iloc[0]) - 1
        top_symbols = mom_scores.sort_values(ascending=False).head(self.final_count).index.tolist()
        
        # 3. HRP Weighting Logic
        final_returns = returns[top_symbols]
        cov, corr = final_returns.cov(), final_returns.corr()
        
        try:
            # Clustering
            dist = np.sqrt(0.5 * (1 - corr))
            link = sch.linkage(sch.distance.squareform(dist), 'single')
            sort_ix = self.get_quasi_diag(link)
            sorted_symbols = [final_returns.columns[i] for i in sort_ix]
            
            # Simple Inverse Variance for the weights (HRP Component)
            inv_var = 1 / np.diag(cov)
            raw_weights = inv_var / inv_var.sum()
            
            # 4. Execute
            self.Liquidate()
            for symbol in sorted_symbols:
                weight = raw_weights[final_returns.columns.get_loc(symbol)]
                self.SetHoldings(symbol, weight)
                self.Debug(f"Buying {symbol} at weight {weight:.2%}")
                
        except Exception as e:
            self.Debug(f"Trade Execution Error: {e}")

    def get_quasi_diag(self, link):
        link = link.astype(int)
        sort_ix = pd.Series([link[-1, 0], link[-1, 1]])
        num_items = link[-1, 3]
        while sort_ix.max() >= num_items:
            sort_ix.index = range(0, sort_ix.shape[0] * 2, 2)
            df0 = sort_ix[sort_ix >= num_items]
            i, j = df0.index, df0.values - num_items
            sort_ix[i] = link[j, 0]
            sort_ix = pd.concat([sort_ix, pd.Series(link[j, 1], index=i + 1)]).sort_index()
        return sort_ix.tolist()