from AlgorithmImports import *
from QuantConnect.DataSource import *

class ExtractAlphaTrueBeatsDataAlgorithm(QCAlgorithm):

    def initialize(self) -> None:
        self.set_start_date(2019, 1, 1)
        self.set_end_date(2020, 1, 1)
        self.set_cash(100000)
        # A variable to control the time of rebalancing
        self.last_time = datetime.min
        
        self.add_universe(self.my_coarse_filter_function)
        self.universe_settings.resolution = Resolution.MINUTE
        
    def my_coarse_filter_function(self, coarse: List[CoarseFundamental]) -> List[Symbol]:
        # Select the stocks with highest dollar volume due to better informed information from more market activities
        # Only the ones with fundamental data are supported by true beats data
        sorted_by_dollar_volume = sorted([x for x in coarse if x.has_fundamental_data], 
                                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.time > self.time: return
        
        # Trade only based on the updated true beats data
        points = slice.Get(ExtractAlphaTrueBeats)
        if not points: return
        
        # Extract the true beats data (earning and revenue estimates) as trading signals
        true_beats = {point.Key: trueBeat for point in points for trueBeat in point.Value}

        # Long the ones with the highest earning and revenue estimates due to fundamental factor may bring stock price up
        # Short the lowest that predicted to bring stock price down
        sorted_by_true_beat = sorted(true_beats.items(), key=lambda x: x[1].true_beat)
        long_symbols = [x[0].underlying for x in sorted_by_true_beat[-10:]]
        short_symbols = [x[0].underlying for x in sorted_by_true_beat[:10]]
        
        # Liquidate the ones without a strong earning and revenue that support stock price direction
        for symbol in self.portfolio.Keys:
            if self.portfolio[symbol].invested \
                and symbol not in long_symbols \
                and symbol not in short_symbols:
                self.liquidate(symbol)
        
        # Invest equally and dollar-neutral to evenly dissipate capital risk and hedge systematic risk
        long_targets = [PortfolioTarget(symbol, 0.05) for symbol in long_symbols]
        short_targets = [PortfolioTarget(symbol, -0.05) for symbol in short_symbols]
        self.set_holdings(long_targets + short_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 true beats data for trading signal generation
            extract_alpha_true_beats_symbol = self.add_data(ExtractAlphaTrueBeats, security.symbol).symbol
            
            # Historical Data
            history = self.history(extract_alpha_true_beats_symbol, 10, Resolution.DAILY)
            self.log(f"We got {len(history)} items from our history request for {security.symbol} ExtractAlpha True Beats data")