# region imports
from AlgorithmImports import *
# endregion

class ContinuousFuturesConstantAlphaModel(AlphaModel):
    ''' Provides an implementation of IAlphaModel that always returns the same insight for each security'''

    def __init__(self, type, direction, period, magnitude = None, confidence = None):
        '''Initializes a new instance of the ConstantAlphaModel class
        Args:
            type: The type of insight
            direction: The direction of the insight
            period: The period over which the insight with come to fruition
            magnitude: The predicted change in magnitude as a +- percentage
            confidence: The confidence in the insight'''
        self.type = type
        self.direction = direction
        self.period = period
        self.magnitude = magnitude
        self.confidence = confidence
        self.securities = []
        self.insights_time_by_symbol = {}

    def update(self, algorithm, data):
        ''' Creates a constant insight for each security as specified via the constructor
        Args:
            algorithm: The algorithm instance
            data: The new data available
        Returns:
            The new insights generated'''
        insights = []

        for security in self.securities:
            # security price could be zero until we get the first data point. e.g. this could happen
            # when adding both forex and equities, we will first get a forex data point
            if security.price != 0 and self.should_emit_insight(algorithm.utc_time, security):
                symbol = security.mapped if security.type == SecurityType.FUTURE else security.symbol
                insights.append(Insight(symbol, self.period, self.type, self.direction, self.magnitude, self.confidence))

        return insights


    def on_securities_changed(self, algorithm, changes):
        ''' Event fired each time the we add/remove securities from the data feed
        Args:
            algorithm: The algorithm instance that experienced the change in securities
            changes: The security additions and removals from the algorithm'''
        for added in changes.added_securities:
            if ((added.type == SecurityType.FUTURE and added.symbol.is_canonical()) or \
                added.type == SecurityType.EQUITY):
                self.securities.append(added)

        # this will allow the insight to be re-sent when the security re-joins the universe
        for removed in changes.removed_securities:
            if removed in self.securities:
                self.securities.remove(removed)
            if removed.symbol in self.insights_time_by_symbol:
                self.insights_time_by_symbol.pop(removed.symbol)


    def should_emit_insight(self, utc_time, security):
        if not security.is_tradable or security.is_delisted:
            return False
            
        generated_time_utc = self.insights_time_by_symbol.get(security.symbol)

        if generated_time_utc is not None:
            # we previously emitted a insight for this symbol, check it's period to see
            # if we should emit another insight
            if utc_time - generated_time_utc < self.period:
                return False

        # we either haven't emitted a insight for this symbol or the previous
        # insight's period has expired, so emit a new insight now for this symbol
        self.insights_time_by_symbol[security.symbol] = utc_time
        return True

# region imports
from AlgorithmImports import *
from universe import FrontMonthFutureUniverseSelectionModel
from portfolio import ModifiedRiskParityPortfolioConstructionModel
# endregion

class AllWeatherPortfolio(QCAlgorithm):

    def initialize(self):
        self.set_end_date(datetime.now())
        self.set_start_date(self.end_date - timedelta(5*365))
        self.set_cash(100000)  # Set Strategy Cash

        self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)
        
        self.universe_settings.resolution = Resolution.MINUTE
        self.universe_settings.leverage = 1
        # We try to include various intuitively unrelated assets to reduce individual asset's risk
        symbols = [Symbol.create(ticker, SecurityType.EQUITY, Market.USA)
            for ticker in [
                "SPY",  # US Equity
                "EEM",  # Emerging market Equity
                "UUP"   # cash
        ]]
        # self.add_universe_selection(ManualUniverseSelectionModel(symbols))
        self.add_universe_selection(FrontMonthFutureUniverseSelectionModel(self.select_future_chain_symbols))
        
        self.add_alpha(ConstantAlphaModel(InsightType.PRICE, InsightDirection.UP, timedelta(30)))

        # RiskParityPortfolioConstructionModel seeks to build a portfolio with the 
        # equal contribution of risk to the total portfolio risk from all assets
        # We build a modified version on top of that to handle corporate events
        # https://www.quantconnect.com/docs/v2/writing-algorithms/algorithm-framework/portfolio-construction/supported-models#11-Risk-Parity-Model
        self.set_portfolio_construction(ModifiedRiskParityPortfolioConstructionModel(
            self, lambda x: Expiry.END_OF_WEEK(x)))

        self.set_warm_up(timedelta(252))

        self.set_benchmark("SPY")
    
    def select_future_chain_symbols(self, utc_time):
        return [ 
            # Symbol.create(Futures.Financials.MICRO_Y_2_TREASURY_BOND, SecurityType.FUTURE, Market.CBOT),
            # Symbol.create(Futures.Financials.MICRO_Y_10_TREASURY_NOTE, SecurityType.FUTURE, Market.CBOT),
            # Symbol.create(Futures.Energies.MICRO_CRUDE_OIL_WTI, SecurityType.FUTURE, Market.NYMEX),
            # Symbol.create(Futures.Metals.MICRO_GOLD, SecurityType.FUTURE, Market.COMEX)
            Symbol.create(Futures.Softs.SUGAR_11, SecurityType.FUTURE, Market.ICE),
            Symbol.create(Futures.Softs.COTTON_2, SecurityType.FUTURE, Market.ICE),
            Symbol.create(Futures.Softs.COFFEE, SecurityType.FUTURE, Market.ICE)
        ]

# region imports
from AlgorithmImports import *
from scipy.optimize import *
# endregion

### <summary>
### Provides an implementation of a risk parity portfolio optimizer that calculate the optimal weights 
### with the weight range from 0 to 1 and equalize the risk carried by each asset
### </summary>
class RiskParityPortfolioOptimizer:
    
    def __init__(self, 
                 algorithm,
                 minimum_weight = 1e-05, 
                 maximum_weight = sys.float_info.max):
        self.algorithm = algorithm
        self.minimum_weight = minimum_weight if minimum_weight >= 1e-05 else 1e-05
        self.maximum_weight = maximum_weight if maximum_weight >= minimum_weight else minimum_weight

    def optimize(self, historical_returns, budget = None, covariance = None):
        if covariance is None:
            covariance = np.cov(historical_returns.T)

        size = historical_returns.columns.size   # K x 1
        if size == 1: 
            return np.array([1])
        
        # Optimization Problem
        # minimize_{x >= 0} f(x) = 1/2 * x^T.S.x - b^T.log(x)
        # b = 1 / num_of_assets (equal budget of risk)
        # df(x)/dx = S.x - b / x
        # H(x) = S + Diag(b / x^2)
        # lw <= x <= up
        x0 = np.array(size * [1. / size])
        budget = budget if budget is not None else x0
        objective = lambda weights: 0.5 * weights.T @ covariance @ weights - budget.T @ np.log(weights)
        gradient = lambda weights: covariance @ weights - budget / weights
        hessian = lambda weights: covariance + np.diag((budget / weights**2).flatten())
        bounds = tuple((self.minimum_weight, self.maximum_weight) for _ in range(size))
        solver = minimize(objective, jac=gradient, hess=hessian, x0=x0, bounds=bounds, method="trust-constr")

        if not solver["success"]:
            if self.algorithm.live_mode:
                self.algorithm.log("RiskParityPortfolioOptimizer.optimize: Did not converge. Return equal weighted")
            return x0
        # Normalize weights: w = x / x^T.1
        return solver["x"]/np.sum(solver["x"])

# region imports
from AlgorithmImports import *
from optimizer import RiskParityPortfolioOptimizer
# endregion

class ModifiedRiskParityPortfolioConstructionModel(PortfolioConstructionModel):
    '''Modified version of LEAN's Risk Parity Portfolio Construction Model to handle corporate events also'''
    def __init__(self,
                 algorithm,
                 rebalance = Resolution.DAILY,
                 portfolio_bias = PortfolioBias.LONG_SHORT,
                 lookback = 1,
                 period = 252,
                 optimizer = None) -> None:
        super().__init__()
        if portfolio_bias == PortfolioBias.SHORT:
            raise ArgumentException("Long position must be allowed in RiskParityPortfolioConstructionModel.")
        
        self.algorithm = algorithm
        self.lookback = lookback
        self.period = period
        self.sign = lambda x: -1 if x < 0 else (1 if x > 0 else 0)

        self.optimizer = RiskParityPortfolioOptimizer(algorithm) if optimizer is None else optimizer

        self.symbol_data_by_symbol = {}

        # If the argument is an instance of Resolution or Timedelta
        # Redefine rebalancing_func
        rebalancing_func = rebalance
        if isinstance(rebalance, int):
            rebalance = Extensions.to_time_span(rebalance)
        if isinstance(rebalance, timedelta):
            rebalancing_func = lambda dt: dt + rebalance
        if rebalancing_func:
            self.set_rebalancing_func(rebalancing_func)

    def determine_target_percent(self, active_insights: List[Insight]) -> Dict[Insight, float]:
        targets = {}

        # Reset indicators when corporate actions occur
        data = self.algorithm.current_slice
        for symbol in set(list(data.splits.keys()) + list(data.dividends.keys())):
            if symbol in self.symbol_data_by_symbol:
                self.symbol_data_by_symbol[symbol].roc.reset()
                self.symbol_data_by_symbol[symbol].window.reset()
                self.symbol_data_by_symbol[symbol].warm_up_indicators()

        # If we have no insights just return an empty target dict
        if not active_insights:
            self.algorithm.log('PortfolioContructionModel: No active insights. Create zero-quantity targets.')
            return targets

        symbols = [insight.symbol for insight in active_insights]

        # Create a dictionary keyed by the symbols in the insights with an pandas.series as value to create a data frame
        returns = { str(symbol.id) : data.returns for symbol, data in self.symbol_data_by_symbol.items() if symbol in symbols }
        returns = pd.DataFrame(returns).dropna(axis=1, how='all')

        # The portfolio optimizer finds the optional weights for the given data
        weights = self.optimizer.optimize(returns)
        weights = pd.Series(weights, index = returns.columns)

        # Create portfolio targets from the specified insights with a 10% buffer
        for insight in active_insights:
            targets[insight] = .9 * weights.get(str(insight.symbol.id), 0) / self.symbol_data_by_symbol[insight.symbol].multiplier

        return targets

    def on_securities_changed(self, algorithm: QCAlgorithm, changes: SecurityChanges) -> None:
        algorithm.log(f'PortfolioContructionModel.on_securities_changed: Changes: {changes}')
        
        # clean up data for removed securities
        super().on_securities_changed(algorithm, changes)
        for removed in changes.removed_securities:
            symbolData = self.symbol_data_by_symbol.pop(removed.symbol, None)
            symbolData.dispose()

        # initialize data for added securities
        for added in changes.added_securities:
            symbol = added.symbol
            if symbol not in self.symbol_data_by_symbol:
                symbolData = self.RiskParitySymbolData(
                    self.algorithm, added, self.lookback, self.period)
                self.symbol_data_by_symbol[symbol] = symbolData

    class RiskParitySymbolData:
        def __init__(self, algorithm, security, lookback, period):
            self.algorithm = algorithm
            self.symbol = security.symbol
            self.lookback = lookback
            self.period = period
            self.multiplier = security.symbol_properties.contract_multiplier \
                if security.type == SecurityType.FUTURE and security.symbol_properties.contract_multiplier > 1 \
                else 1

            self.roc = RateOfChange(f'{self.symbol}.ROC({lookback})', lookback)
            self.roc.updated += self.on_rate_of_change_updated
            self.window = RollingWindow[IndicatorDataPoint](period)

            self.consolidator = TradeBarConsolidator(1)
            algorithm.subscription_manager.add_consolidator(self.symbol, self.consolidator)
            algorithm.register_indicator(self.symbol, self.roc, self.consolidator)

            self.warm_up_indicators()

        def dispose(self):
            self.algorithm.subscription_manager.remove_consolidator(self.symbol, self.consolidator)
            self.roc.updated -= self.on_rate_of_change_updated
            self.roc.reset()
            self.window.reset()

        def warm_up_indicators(self):
            self.algorithm.warm_up_indicator(self.symbol, self.roc, Resolution.DAILY)

        def on_rate_of_change_updated(self, roc, updated):
            if roc.is_ready:
                self.window.add(updated)

        @property
        def returns(self):
            series = pd.Series(
                data = [x.value for x in self.window],
                index = [x.end_time.date() for x in self.window]
            )
            # if same date, only keep latest datum
            series = series[~series.index.duplicated(keep='last')]
            return series

        @property
        def is_ready(self):
            return self.window.is_ready

#region imports
from AlgorithmImports import *
from Selection.FutureUniverseSelectionModel import FutureUniverseSelectionModel
#endregion

class FrontMonthFutureUniverseSelectionModel(FutureUniverseSelectionModel):
    '''Creates futures chain universes that select the front month contract and runs a user
    defined futureChainSymbolSelector every day to enable choosing different futures chains'''
    def __init__(self, select_future_chain_symbols):
        super().__init__(timedelta(7), select_future_chain_symbols)

    def filter(self, filter):
        '''Defines the futures chain universe filter'''
        return (filter.front_month()
                      .only_apply_filter_at_market_open())