# region imports
from AlgorithmImports import *
# endregion


class MarketCapFactor:

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

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


class SortinoFactor:

    def __init__(self, algorithm, symbol, lookback):
        self._sortino = algorithm.sortino(symbol, lookback, resolution=Resolution.DAILY)

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


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

# region imports
from AlgorithmImports import *

from scipy import optimize
from scipy.optimize import Bounds
from factors import *
# endregion


class FactorWeightOptimizationAlgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2014, 12, 31)
        self.set_cash(10_000_000)
        self.settings.automatic_indicator_warm_up = True
        spy = Symbol.create('SPY', SecurityType.EQUITY, Market.USA)
        # Add a universe of hourly data.
        self.universe_settings.resolution = Resolution.HOUR
        universe_size = self.get_parameter('universe_size', 20)
        self._universe = self.add_universe(self.universe.etf(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)[-universe_size:]]))
        self._lookback = self.get_parameter('lookback', 21) # Set a 21-day trading lookback.
        # Create a Schedule Event to rebalance the portfolio.
        self.schedule.on(self.date_rules.month_start(spy), self.time_rules.after_market_open(spy, 31), self._rebalance)

    def on_securities_changed(self, changes):
        for security in changes.added_securities: # Create factors for assets that enter the universe.
            security.factors = [MarketCapFactor(security), SortinoFactor(self, security.symbol, self._lookback)]

    def _rebalance(self):
        # Get raw factor values of the universe constituents.
        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
        # Calculate the factor z-scores.
        factor_zscores = (factors_df - factors_df.mean()) / factors_df.std()
        # Run an optimization to find optimal factor weights. Objective: Maximize trailing return. Initial guess: Equal-weighted.
        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]
        factor_weights = optimize.minimize(lambda weights: -(np.dot(factor_zscores, weights) * trailing_return).sum(), x0=np.array([1.0/num_factors] * num_factors), method='Nelder-Mead', bounds=Bounds([0] * num_factors, [1] * num_factors), options={'maxiter': 10}).x
        # Calculate the portfolio weights. Ensure the portfolio is long-only with 100% exposure, then rebalance the portfolio.
        portfolio_weights = (factor_zscores * factor_weights).sum(axis=1)
        portfolio_weights = portfolio_weights[portfolio_weights > 0]
        self.set_holdings([PortfolioTarget(symbol, weight/portfolio_weights.sum()) for symbol, weight in portfolio_weights.items()], True)