# region imports
from AlgorithmImports import *
# endregion


class PredictionUniverseAlgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2026, 2, 10) # Set this 1 day before the dataset begins.
        self.set_cash(100_000)
        self.settings.seed_initial_prices = True
        # Define some parameters.
        self._return_prediction_threshold = 0
        # Add the universe of custom data.
        self.universe_settings.resolution = Resolution.DAILY
        self._universe = self.add_universe(PredictionUniverse, self._select_assets)
        # Add a Scheduled Event to rebalance the portfolio every day.
        self.schedule.on(self.date_rules.every_day('SPY'), self.time_rules.at(8, 0), self._rebalance)

    def _select_assets(self, data):
        # Select the assets with a return prediction above some threshold.
        return [stock.symbol for stock in data if stock.value > self._return_prediction_threshold]
    
    def _rebalance(self):
        # Form an equal-weighted portfolio with the assets in the universe.
        symbols = self._universe.selected
        if not symbols:
            return
        targets = [PortfolioTarget(symbol, 1/len(symbols)) for symbol in symbols]
        self.set_holdings(targets, True)


# Define a custom universe class that reads from the JSON file
# we created in the research notebook.
class PredictionUniverse(PythonData):

    def get_source(self, config, date, is_live_mode):
        # Read from the JSON file we saved in the Object Store.
        return SubscriptionDataSource(
            'research-to-backtest-factors.json', 
            SubscriptionTransportMedium.OBJECT_STORE, 
            FileFormat.UNFOLDING_COLLECTION
        )

    def reader(self, config, line, date, is_live):
        objects = []
        # Iterate through each day.
        for obj in json.loads(line):
            end_time = datetime.strptime(obj["date"], "%Y-%m-%d")
            # Iterate through each stock on this day.
            for symbol, prediction in obj['prediction_by_symbol'].items():
                stock = PredictionUniverse()
                stock.symbol = Symbol(SecurityIdentifier.parse(symbol), "")
                stock.end_time = end_time
                stock.value = prediction
                objects.append(stock)
        # The BaseDataCollection contains multiple PredictionUniverse objects
        # for each day. Set its `time` to match the `end_time` of one of the objects.
        # We pack the collection and timestamp it with the last timestamp (-1) to signal 
        # that the collection has data prior to (or including) that timestamp.
        return BaseDataCollection(objects[-1].end_time, config.symbol, objects)