from AlgorithmImports import *


class AggregateSalesGrowthRotationAlgorithm(QCAlgorithm):

    def initialize(self) -> None:
        self.set_start_date(2021, 7, 1)
        self.set_end_date(2026, 7, 15)
        self.set_cash(100_000)
        # Add the SPY and BIL ETFs to trade.
        self._spy = self.add_equity("SPY", Resolution.DAILY, leverage=3)
        self._bil = self.add_equity("BIL", Resolution.DAILY, leverage=3)
        # Add some members we'll need to make trading decisions.
        self._firm_data = {}
        self._asg = pd.Series()
        self._market_return = RateOfChange(1)
        self._excess_returns = pd.Series()
        self._gamma = 3
        lookback_years = 10
        self._var = Variance(lookback_years * 12)
        # Add a universe that runs selection at the start of each month.
        date_rule = self.date_rules.month_start("SPY")
        self.universe_settings.schedule.on(date_rule)
        self._universe = self.add_universe(self._select_assets)
        # Add a Scheduled Event to rebalance the portfolio each month.
        self.schedule.on(date_rule, self.time_rules.at(8, 0), self._rebalance)
        # Add a warm-up period to prime the factors and labels.
        self.set_warm_up(timedelta((lookback_years+1)*365))

    def _select_assets(self, fundamentals: List[Fundamental]) -> List[Symbol]:
        # Get revenue growth and market cap of stocks no in the Financial Services and Real Estate sectors.
        self._firm_data = {
            f.symbol: (f.operation_ratios.revenue_growth.one_year, f.market_cap)
            for f in fundamentals
            if (f.company_reference.country_id == "USA" and
                f.security_reference.is_primary_share and
                f.security_reference.security_type == "ST00000001" and  # Common stock
                f.asset_classification.morningstar_sector_code not in (MorningstarSectorCode.FINANCIAL_SERVICES, MorningstarSectorCode.REAL_ESTATE))
        }
        return []

    def _rebalance(self) -> None:
        # Get the month that just ended.
        month = pd.Period(self.time, freq="M") - 1
        # Update the excess return history.
        if self._market_return.update(self.time, self._spy.price):
            excess_return = self._market_return.current.value - self.risk_free_interest_rate_model.get_interest_rate(self.time) / 12
            self._excess_returns[month] = excess_return
            self._var.update(self.time, excess_return)
        # Calculate the market-cap-weighted ASG, winsorised at the 1st/99th percentiles.
        firms = pd.DataFrame.from_dict(self._firm_data, orient="index", columns=["growth", "market_cap"]).dropna(subset=["growth"])
        growth = firms["growth"].clip(*firms["growth"].quantile([0.01, 0.99]))
        caps = firms["market_cap"]
        usable = caps.notna() & (caps > 0)
        if usable.any():
            self._asg[month] = np.average(growth[usable], weights=caps[usable])
        # If we're still warming up, do nothing.
        if self.is_warming_up:
            return
        # Regress this month's excess return on last month's ASG.
        X, y = self._asg.shift(1, freq="M").align(self._excess_returns, join="inner")
        alpha, beta = np.polynomial.polynomial.polyfit(X, y, 1)
        # Forecast this month's excess return.
        forecast_r = alpha + beta * self._asg.get(month)
        # Calculate the target exposure to SPY using Merton's closed-form mean-variance solution.
        w_star = np.clip(forecast_r / (self._gamma * self._var.current.value), 0, 1.5)
        # Plot the ASG, beta, and target weights of each asset.
        self.plot('ASG', 'Value', self._asg.get(month))
        self.plot('Regression', 'Beta', beta)
        self.plot('Weights', 'SPY', w_star)
        self.plot('Weights', 'BIL', 1.0 - w_star)
        # Place trades to rebalance the portfolio.
        self.set_holdings([PortfolioTarget(self._spy, w_star), PortfolioTarget(self._bil, 1 - w_star)])