from AlgorithmImports import *
from QuantConnect.DataSource import *

class ExtractAlphaCrossAssetModelAlgorithm(QCAlgorithm):

    def initialize(self) -> None:
        self.set_start_date(2019, 1, 1)
        self.set_end_date(2019, 12, 31)
        self.set_cash(100000)
        
        self.last_time = datetime.min
        
        self.add_universe(self.my_coarse_filter_function)

        self.dataset_symbol_by_symbol = {}
        self.points = {}
        
    def my_coarse_filter_function(self, coarse: List[CoarseFundamental]) -> List[Symbol]:
        sorted_by_dollar_volume = sorted([x for x in coarse if x.has_fundamental_data and x.price > 4], 
                                key=lambda x: x.dollar_volume, reverse=True)
        selected = [x.symbol for x in sorted_by_dollar_volume[:100]]
        return selected

    def on_data(self, slice: Slice) -> None:
        if self.last_time > self.time: return
    
        # Accessing Data
        points = slice.Get(ExtractAlphaCrossAssetModel)
        if points:
            self.points = points

        if slice.time.time() < time(10): return

        sorted_by_score = sorted([x for x in self.points.items() if x[1].score != None], 
            key=lambda x: x[1].score, reverse=True)
        long_symbols = [x[0].underlying for x in sorted_by_score[:10]]
        short_symbols = [x[0].underlying for x in sorted_by_score[-10:]]

        portfolio_targets = []
        for symbol, security_holding in self.portfolio.items():
            weight = 0
            if symbol in long_symbols:
                weight = 0.05
            elif symbol in short_symbols:
                weight = -0.05
            elif not security_holding.invested:
                continue
            portfolio_targets.append(PortfolioTarget(symbol, weight))
        self.set_holdings(portfolio_targets)
        
        self.last_time = Expiry.END_OF_DAY(self.time)
        
    def on_securities_changed(self, changes: SecurityChanges) -> None:
        for security in changes.added_securities:
            # Requesting data
            self.dataset_symbol_by_symbol[security.symbol] = self.add_data(ExtractAlphaCrossAssetModel, security.symbol).symbol

        for security in changes.removed_securities:
            dataset_symbol = self.dataset_symbol_by_symbol.pop(security.symbol, None)
            if dataset_symbol:
                self.remove_security(dataset_symbol)