# region imports
from AlgorithmImports import *
# endregion


class QuiverLobbyingLongOnlyAlgorithm(QCAlgorithm):

    def initialize(self) -> None:
        self.set_start_date(2019, 2, 1)
        self.set_cash(1_000_000_000)

        # Seed prices for newly-added securities so set_holdings doesn't hit $0-price errors
        self.settings.seed_initial_prices = True

        # Leave a small cash buffer so set_holdings target weights don't exceed buying power
        self.settings.free_portfolio_value_percentage = 0.05

        # Daily resolution for the dynamic universe
        self.universe_settings.resolution = Resolution.DAILY

        # Align universe selection with the monthly rebalance schedule so the
        # selection fires on the same day as the rebalance (not every day).
        self._data_by_symbol = {}
        # Corporate Lobbying universe from Quiver Quantitative
        self._universe = self.add_universe(QuiverGovernmentContractUniverse, self._select_assets)

        # Monthly rebalance at 08:00 ET on the first trading day of the month
        self.schedule.on(
            self.date_rules.month_start("SPY"),
            self.time_rules.at(8, 0),
            self._rebalance,
        )

    def _select_assets(self, alt_coarse: List[QuiverGovernmentContractUniverse]) -> List[Symbol]:
        for datum in alt_coarse:
            symbol = datum.symbol
            if symbol not in self._data_by_symbol:
                self._data_by_symbol[symbol] = []
            self._data_by_symbol[symbol].append(datum)
        return [symbol for symbol, d in self._data_by_symbol.items() if len(d) >= 3 and sum([x.amount for x in d]) > 50000]

    def _rebalance(self) -> None:
        selected = self._universe.selected
        if not selected:
            return

        # Filter out securities that haven't received a price bar yet
        tradable = [s for s in selected if self.securities[s].price]
        if not tradable:
            return
        
        self._data_by_symbol.clear()

        weight = 1.0 / len(tradable)
        targets = [PortfolioTarget(symbol, weight) for symbol in tradable]

        self.set_holdings(targets, liquidate_existing_holdings=True)