from AlgorithmImports import *


class BBandsAlgorithm(QCAlgorithm):

    def initialize(self) -> None:
        self.set_start_date(self.end_date - timedelta(5*365))
        self.set_cash(100_000)
        self.settings.automatic_indicator_warm_up = True
        self.settings.free_portfolio_value_percentage = 0.05
        # Define some parameters.
        self._universe_size = self.get_parameter("selected", 20)
        self._bb_std = self.get_parameter("stdBB", 4)
        self._bb_period = self.get_parameter("BB", 15)
        self._roc_period = self.get_parameter("ROC", 79)
        # Add a universe of US Equities.
        self.universe_settings.resolution = Resolution.DAILY
        self._universe = self.add_universe(self.universe.etf('SPY', universe_filter_func=self._select_assets))
        # Add some alternative investments.
        self._uup, self._tlt, self._gld = [
            self.add_equity(ticker, Resolution.DAILY) 
            for ticker in ['UUP', 'TLT', 'GLD']
        ]
    
    def _select_assets(self, constituents: List[ETFConstituentData]) -> List[Symbol]:
        # Select the largets stocks in the ETF.
        filtered = sorted([c for c in constituents if c.weight], key=lambda c: c.weight)[-self._universe_size:]
        return [f.symbol for f in filtered]

    def on_securities_changed(self, changes: SecurityChanges) -> None:
        # As assets enter the universe, add some indicators for them.
        for security in changes.added_securities:
            security.bb = self.bb(security, self._bb_period, self._bb_std)
            security.roc = self.roc(security, self._roc_period)
            security.last_trade_time = None
        for security in changes.removed_securities:
            self.liquidate(security)
            self.deregister_indicator(security.bb)
            self.deregister_indicator(security.roc)     

    def on_data(self, data: Slice) -> None:
        if not data.bars or not self._universe.selected:
            return
        # Get the ROC value for all assets in the universe.
        roc_by_security = {}
        for symbol in self._universe.selected:
            security = self.securities[symbol]
            if security.roc.is_ready and security.roc.current.value > 0:
                roc_by_security[security] = security.roc.current.value
        # If no assets in the universe have a positive ROC, resort to 
        # one of the alternative investements.
        if not roc_by_security:
            # Decide which alternative to buy based on their ROC values.
            roc_uup = self._uup.roc.current.value
            roc_tlt = self._tlt.roc.current.value
            if roc_uup > 0 or roc_tlt > 0:
                if roc_uup > roc_tlt:
                    self.set_holdings(self._uup, 1, True, tag="UUP ROC > 0")
                else:
                    self.set_holdings(self._tlt, 1, True, tag="TLT ROC > 0")
            else:
                # If both UUP and TLT have negative ROC, invest in GLD.
                self.set_holdings(self._gld, 1, True, tag="GLD ROC > 0")
            return
        # Scan for an entry: If the assets with the greatest ROC is in the buy zone of the BB, buy it.
        security = max(roc_by_security, key=lambda security: roc_by_security[security])
        if (not security.invested and 
            security.bb.middle_band.Current.Value < security.price < security.bb.upper_band.Current.Value):
            self.set_holdings(security, 1, True, tag='BUY CONDITIONS MET')
            security.last_trade_time = self.time
        # Scan for exits.
        for symbol in self._universe.selected:
            security = self.securities[symbol]
            if not security.invested:
                continue
            # Profit taking
            tag = ''
            if security.holdings.unrealized_profit_percent >= 0.05:
                tag = "5%+ PROFIT"
            # Stop Loss
            elif security.holdings.unrealized_profit_percent <= -0.05:
                tag = "STOP LOSS"
            # Time-based exit
            elif (self.time - security.last_trade_time).days >= 10:
                tag = "HELD FOR 10 DAYS"
            else:
                continue
            self.liquidate(security, tag)
            security.last_trade_time = None