# region imports
from AlgorithmImports import *

from scipy import optimize
# endregion

class MaintainHistoricalDailyUniversePriceDataAlgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2015, 3, 20)
        self.set_cash(1_000_000)

        self.settings.automatic_indicator_warm_up = True
        self._spy = self.add_equity('SPY', Resolution.DAILY).symbol
        self.set_security_initializer(
            BrokerageModelSecurityInitializer(self.brokerage_model, FuncSecuritySeeder(self.get_last_known_prices))
        )
        # Add a universe of daily data.
        self.universe_settings.resolution = Resolution.HOUR
        self._universe = self.add_universe(
            self.universe.etf(self._spy, universe_filter_func=lambda constituents: [
                    c.symbol for c in sorted(
                        [c for c in constituents if c.weight], 
                        key=lambda c: c.weight
                    )[-100:]
                ]
            )
        )

        # Define the lookback period.
        self._lookback = 252  # Trading days.

        # Create a Schedule Event to rebalance the portfolio.
        self.schedule.on(
            self.date_rules.month_start(self._spy),
            self.time_rules.after_market_open(self._spy, 31),
            self._rebalance
        )

        self.set_warm_up(timedelta(30))

    def on_securities_changed(self, changes):
        for security in changes.added_securities:
            security.factors = [
                CorrFactor(self, security.symbol, self._spy, self._lookback),
                ROCFactor(self, security.symbol, self._lookback),
                MarketCapFactor(self, security)
            ]

    def _rebalance(self):
        if self.is_warming_up or not self._universe.selected:
            return

        # Get raw factor values
        factors_df = pd.DataFrame()
        for symbol in self._universe.selected:
            for i, factors in enumerate(self.securities[symbol].factors):
               factors_df.loc[symbol, i] = factors.value
        raw_corr = factors_df.copy()
        # Calculate factor z-scores.
        factor_zscores = (factors_df - factors_df.mean()) / factors_df.std()

        # Run optimization
        trailing_return = self.history(list(self._universe.selected), self._lookback, Resolution.DAILY).close.unstack(0).pct_change(self._lookback-1).iloc[-1]
        num_factors = factors_df.shape[1]
        result = optimize.minimize(
            lambda weights: -(np.dot(factor_zscores, weights) * trailing_return).sum(), # Maximize trailing return
            x0=np.array([1/num_factors] * num_factors), # Initial guess: Equal-weighted
            method='SLSQP',
            bounds=tuple((0, 1) for _ in range(num_factors)),
            constraints=({'type': 'eq', 'fun': lambda x: np.sum(x) - 1})  # The factor weights must sum to 1
        )
        factor_weights = result.x
        for i, w in enumerate(factor_weights):
            self.plot(f"Factor Weights", str(i), w)
        # Calculate the portfolio weights.
        portfolio_weights = np.dot(factor_zscores, factor_weights)
        # Make portfolio long-only
        portfolio_weights[portfolio_weights < 0] = 0
        # Make exposure 100%
        portfolio_weights[portfolio_weights > 0] /= portfolio_weights[portfolio_weights > 0].sum()
        # Rebalance the portfolio.
        self.set_holdings([PortfolioTarget(factor_zscores.index[i], portfolio_weights[i]) for i, symbol in enumerate(portfolio_weights)], True)


class CorrFactor:

    def __init__(self, algorithm, symbol, reference, lookback):
        self._c = algorithm.c(symbol, reference, lookback, correlation_type=CorrelationType.Pearson, resolution=Resolution.DAILY)

    @property
    def value(self):
        return 1 - abs(self._c.current.value)
    
    
class ROCFactor:

    def __init__(self, algorithm, symbol, lookback):
        self._roc = algorithm.roc(symbol, lookback, resolution=Resolution.DAILY)

    @property
    def value(self):
        return self._roc.current.value


class MarketCapFactor:

    def __init__(self, algorithm, security):
        self._security = security

    @property
    def value(self):
        return self._security.fundamentals.market_cap