from AlgorithmImports import *
from itertools import combinations

class TechUniverseSelection(QCAlgorithm):
    
    def Initialize(self):
        self.SetStartDate(2023, 12, 11)  
        self.SetEndDate(2024, 12, 1) 
        self.SetCash(100000)
        
        # Add universe with coarse and fine selection filters
        self.AddUniverse(self.CoarseSelectionFilter, self.FineSelectionFunction)
        
        self.UniverseSettings.Resolution = Resolution.Daily
        self.SetSecurityInitializer(lambda x: x.SetDataNormalizationMode(DataNormalizationMode.Raw))
        
        # Set leverage to 2 as in example
        self.UniverseSettings.Leverage = 2
        
        # Store the pairs
        self.stock_pairs = []
        
        self.Liquidate()
                
    def CoarseSelectionFilter(self, coarse):
        # Sort by dollar volume and filter by price > $10
        sortedByDollarVolume = sorted(coarse, key=lambda c: c.DollarVolume, reverse=True)
        filteredByPrice = [c.Symbol for c in sortedByDollarVolume if c.Price > 10]
        
        # Keep top 100 for fine filtering
        return filteredByPrice[:100]
        
    def FineSelectionFunction(self, fine):
        # Filter for technology sector stocks
        tech_stocks = [x for x in fine if x.AssetClassification.MorningstarSectorCode == MorningstarSectorCode.Technology]
        
        # Sort by market cap
        sortedByMarketCap = sorted(tech_stocks, key=lambda c: c.MarketCap, reverse=True)
        
        # Get company names to check for duplicates
        top_companies = []
        filtered_symbols = []
        count = 0
        
        for stock in sortedByMarketCap:
            company_name = stock.CompanyReference.StandardName
            if company_name not in top_companies:
                top_companies.append(company_name)
                filtered_symbols.append(stock.Symbol)
                count += 1
                if count == 3:
                    break
        
        # Generate all possible pairs using combinations
        self.stock_pairs = list(combinations(filtered_symbols, 2))
        
        # Log the pairs for verification
        self.Debug(f"Generated pairs: {self.stock_pairs}")
        self.Debug(f"Number of pairs: {len(self.stock_pairs)}")
        
        return filtered_symbols
        
    def OnSecuritiesChanged(self, changes):
        # Log changes in selected securities
        for security in changes.AddedSecurities:
            self.Debug(f"Added: {security.Symbol}")
            
        for security in changes.RemovedSecurities:
            self.Debug(f"Removed: {security.Symbol}")
    
    def OnData(self, data):
        pass