#region imports
from AlgorithmImports import *

from sklearn.ensemble import RandomForestRegressor
#endregion


class RandomForestAlphaModel(AlphaModel):

    _securities = []
    _scheduled_event = None
    _time = datetime.min
    _rebalance = False

    def __init__(self, algorithm, minutes_before_close, n_estimators, min_samples_split, lookback_days):
        self._algorithm = algorithm
        self._minutes_before_close = minutes_before_close
        self._n_estimators = n_estimators
        self._min_samples_split = min_samples_split
        self._lookback_days = lookback_days

    def update(self, algorithm: QCAlgorithm, data: Slice) -> List[Insight]:
        if not self._rebalance or data.quote_bars.count == 0:
            return []
        
        # Fetch history on our universe
        symbols = [s.symbol for s in self._securities]
        df = algorithm.history(symbols, 2, Resolution.DAILY, data_normalization_mode=DataNormalizationMode.SCALED_RAW)
        if df.empty: 
            return []

        self._rebalance = False
    
        # Make all of them into a single time index.
        df = df.close.unstack(level=0)
    
        # Feature engineer the data for input
        input_ = df.diff() * 0.5 + df * 0.5
        input_ = input_.iloc[-1].fillna(0).values.reshape(1, -1)
        
        # Predict the expected price
        predictions = self._regressor.predict(input_)
        
        # Get the expected return
        predictions = (predictions - df.iloc[-1].values) / df.iloc[-1].values
        predictions = predictions.flatten()
        
        insights = []

        for i in range(len(predictions)):
            insights.append( Insight.price(df.columns[i], timedelta(5), InsightDirection.UP, predictions[i]) )
        algorithm.insights.cancel(symbols)
        return insights

        # for i in range(len(predictions)):
        #     # Check if the prediction is positive (for long) or negative (for short)
        #     if predictions[i] > 0:
        #         direction = InsightDirection.UP  # Long signal
        #     else:
        #         direction = InsightDirection.DOWN  # Short signal
        #     insights.append(Insight.price(df.columns[i], timedelta(5), direction, abs(predictions[i])))
        # algorithm.insights.cancel(symbols)
        # return insights

    def _train_model(self):
        # Initialize the Random Forest Regressor
        self._regressor = RandomForestRegressor(n_estimators=self._n_estimators, min_samples_split=self._min_samples_split, random_state = 1990)
        
        # Get historical data
        history = self._algorithm.history([s.symbol for s in self._securities], self._lookback_days, Resolution.DAILY, data_normalization_mode=DataNormalizationMode.SCALED_RAW)
        
        # Select the close column and then call the unstack method.
        df = history['close'].unstack(level=0)
        
        # Feature engineer the data for input.
        input_ = df.diff() * 0.5 + df * 0.5
        input_ = input_.iloc[1:].ffill().fillna(0)
        
        # Shift the data for 1-step backward as training output result.
        output = df.shift(-1).iloc[:-1].ffill().fillna(0)
        
        # Fit the regressor
        self._regressor.fit(input_, output)


    def _before_market_close(self):
        if self._time < self._algorithm.time:
            self._train_model()
            self._time = Expiry.end_of_month(self._algorithm.time)
        self._rebalance = True

    def on_securities_changed(self, algorithm: QCAlgorithm, changes: SecurityChanges) -> None:
        for security in changes.removed_securities:
            if security in self._securities:
                self._securities.remove(security)
                
        for security in changes.added_securities:
            self._securities.append(security)

            # Add Scheduled Event
            if self._scheduled_event == None:
                symbol = security.symbol
                self._scheduled_event = algorithm.schedule.on(
                    algorithm.date_rules.every_day(symbol), 
                    algorithm.time_rules.before_market_close(symbol, self._minutes_before_close), 
                    self._before_market_close
                )

        self._train_model()

# region imports
from AlgorithmImports import *

from alpha import RandomForestAlphaModel
from portfolio import MeanVarianceOptimizationPortfolioConstructionModel
# endregion


class RandomForestAlgorithm(QCAlgorithm):

    _undesired_symbols_from_previous_deployment = []
    _checked_symbols_from_previous_deployment = False

    def initialize(self):
        # self.set_start_date(2022, 6, 1)
        # self.set_end_date(2023, 6, 1)
        self.set_start_date(2018, 1, 1)
        # self.set_end_date(2024, 8, 31)
        self.set_cash(1_000_000)
        
        self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)

        self.settings.minimum_order_margin_portfolio_percentage = 0

        self.AddEquity("XLK", Resolution.Daily)
        self.benchmarkTicker = 'XLK'
        self.SetBenchmark(self.benchmarkTicker)
        self.initBenchmarkPrice = 0
        self.benchmarkExposure = 1

        self.universe_settings.data_normalization_mode = DataNormalizationMode.RAW
        # tickers = ["AAPL", "MSFT", "NVDA", "AMZN", "GOOG", "META", "TSLA", "AVGO", "ORCL", "NFLX", 
        #             "ADBE", "CRM", "AMD", "CSCO", "IBM", "TXN", "QCOM", "NOW", "INTU", "UBER"]
        # symbols = [ Symbol.create(ticker, SecurityType.EQUITY, Market.USA) for ticker in tickers]
        # self.add_universe_selection(ManualUniverseSelectionModel(symbols))

        # Use the following method for a Classic Algorithm
        self.AddUniverse(self.Universe.ETF("SCHG", Market.USA, self.UniverseSettings, self.ETFConstituentsFilter))
        symbol = Symbol.Create("SCHG", SecurityType.Equity, Market.USA)
        # self.AddUniverse(self.Universe.ETF("VONG", Market.USA, self.UniverseSettings, self.ETFConstituentsFilter))
        # symbol = Symbol.Create("VONG", SecurityType.Equity, Market.USA)
        # self.AddUniverse(self.Universe.ETF("IWF", Market.USA, self.UniverseSettings, self.ETFConstituentsFilter))
        # symbol = Symbol.Create("IWF", SecurityType.Equity, Market.USA)
        # self.AddUniverse(self.Universe.ETF("SPYG", Market.USA, self.UniverseSettings, self.ETFConstituentsFilter))
        # symbol = Symbol.Create("SPYG", SecurityType.Equity, Market.USA)
        # self.AddUniverse(self.Universe.ETF("IWY", Market.USA, self.UniverseSettings, self.ETFConstituentsFilter))
        # symbol = Symbol.Create("IWY", SecurityType.Equity, Market.USA)
        # self.AddUniverse(self.Universe.ETF("SPGP", Market.USA, self.UniverseSettings, self.ETFConstituentsFilter))
        # symbol = Symbol.Create("SPGP", SecurityType.Equity, Market.USA)
        # Use the following method for a Framework Algorithm
        self.AddUniverseSelection(ETFConstituentsUniverseSelectionModel(symbol, self.UniverseSettings, self.ETFConstituentsFilter))

        # Add dynamic universe selection with fundamental filtering
        # self.AddUniverse(self.FundamentalUniverseSelection)

        self.add_alpha(RandomForestAlphaModel(
            self,
            self.get_parameter("minutes_before_close", 5),
            self.get_parameter("n_estimators", 100),
            self.get_parameter("min_samples_split", 5),
            self.get_parameter("lookback_days", 360)
        ))

        self.set_portfolio_construction(MeanVarianceOptimizationPortfolioConstructionModel(self, lambda time: None, PortfolioBias.LONG, period=self.get_parameter("pcm_periods", 5)))

        # self.set_portfolio_construction(MeanVarianceOptimizationPortfolioConstructionModel(
        #     self, 
        #     lambda time: None, 
        #     PortfolioBias.LONG_SHORT,  # Allows both long and short positions
        #     period=self.get_parameter("pcm_periods", 5)
        # ))
        
        self.add_risk_management(NullRiskManagementModel())

        self.set_execution(ImmediateExecutionModel())

        self.set_warm_up(timedelta(5))

    def ETFConstituentsFilter(self, constituents):
        # Get the 10 securities with the largest weight in the index
        selected = sorted([c for c in constituents if c.Weight],
            key=lambda c: c.Weight, reverse=True)[:20]
        self.weightBySymbol = {c.Symbol: c.Weight for c in selected}
        
        return list(self.weightBySymbol.keys())

    def FundamentalUniverseSelection(self, fundamental):
        """
        Selects stocks based on fundamental data.
        
        - Filters stocks in the energy sector with a positive market cap.
        - Sorts filtered stocks by market capitalization in descending order.
        - Selects the top 20 stocks by market cap.
        
        :param fundamental: List of fundamental data objects
        :return: List of selected stock symbols
        """
        energy_sector_code = MorningstarSectorCode.ENERGY  # Define sector code for energy

        # Filter stocks based on sector and market cap
        filtered = [
            x for x in fundamental
            if x.AssetClassification.MorningstarSectorCode == energy_sector_code and x.MarketCap > 0
        ]

        # Sort filtered stocks by market capitalization
        sorted_by_market_cap = sorted(filtered, key=lambda x: x.MarketCap, reverse=True)
        
        # Return top 20 stocks by market capitalization
        return [x.Symbol for x in sorted_by_market_cap][:20]

    def on_data(self, data):

        self.UpdateBenchmarkValue()
        self.Plot('Strategy Equity', self.benchmarkTicker, self.benchmarkValue)

        # Exit positions that aren't backed by existing insights.
        # If you don't want this behavior, delete this method definition.
        if not self.is_warming_up and not self._checked_symbols_from_previous_deployment:
            for security_holding in self.portfolio.values():
                if not security_holding.invested:
                    continue
                symbol = security_holding.symbol
                if not self.insights.has_active_insights(symbol, self.utc_time):
                    self._undesired_symbols_from_previous_deployment.append(symbol)
            self._checked_symbols_from_previous_deployment = True
        
        for symbol in self._undesired_symbols_from_previous_deployment[:]:
            if self.is_market_open(symbol):
                self.liquidate(symbol, tag="Holding from previous deployment that's no longer desired")
                self._undesired_symbols_from_previous_deployment.remove(symbol)

    def UpdateBenchmarkValue(self):
        ''' Simulate buy and hold the Benchmark '''
        # if self.initBenchmarkPrice is None:
        if self.initBenchmarkPrice == 0: # Use if Plotting Short Position of Benchmark
            self.initBenchmarkCash = self.Portfolio.Cash
            self.initBenchmarkPrice = self.Benchmark.Evaluate(self.Time)
            self.benchmarkValue = self.initBenchmarkCash
        else:
            currentBenchmarkPrice = self.Benchmark.Evaluate(self.Time)
            # self.benchmarkValue = (currentBenchmarkPrice / self.initBenchmarkPrice) * self.initBenchmarkCash
            
            # Use if Plotting Short Position of Benchmark
            lastReturn = ((currentBenchmarkPrice / self.initBenchmarkPrice) - 1) * self.benchmarkExposure
            self.benchmarkValue = (1 + lastReturn) * self.initBenchmarkCash

# region imports
from AlgorithmImports import *

from alpha import RandomForestAlphaModel
from portfolio import MeanVarianceOptimizationPortfolioConstructionModel
# endregion


class RandomForestAlgorithm(QCAlgorithm):

    _undesired_symbols_from_previous_deployment = []
    _checked_symbols_from_previous_deployment = False

    def initialize(self):
        self.set_start_date(2022, 6, 1)
        self.set_end_date(2023, 6, 1)
        self.set_cash(1_000_000)
        
        self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)

        self.settings.minimum_order_margin_portfolio_percentage = 0

        self.universe_settings.data_normalization_mode = DataNormalizationMode.RAW
        tickers = ["SHY", "TLT", "IEI", "SHV", "TLH", "EDV", "BIL", "SPTL", "TBT", "TMF", 
                    "TMV", "TBF", "VGSH", "VGIT", "VGLT", "SCHO", "SCHR", "SPTS", "GOVT"]
        symbols = [ Symbol.create(ticker, SecurityType.EQUITY, Market.USA) for ticker in tickers]
        self.add_universe_selection(ManualUniverseSelectionModel(symbols))

        self.add_alpha(RandomForestAlphaModel(
            self,
            self.get_parameter("minutes_before_close", 5),
            self.get_parameter("n_estimators", 100),
            self.get_parameter("min_samples_split", 5),
            self.get_parameter("lookback_days", 360)
        ))

        self.set_portfolio_construction(MeanVarianceOptimizationPortfolioConstructionModel(self, lambda time: None, PortfolioBias.LONG, period=self.get_parameter("pcm_periods", 5)))
        
        self.add_risk_management(NullRiskManagementModel())

        self.set_execution(ImmediateExecutionModel())

        self.set_warm_up(timedelta(5))

    def on_data(self, data):
        # Exit positions that aren't backed by existing insights.
        # If you don't want this behavior, delete this method definition.
        if not self.is_warming_up and not self._checked_symbols_from_previous_deployment:
            for security_holding in self.portfolio.values():
                if not security_holding.invested:
                    continue
                symbol = security_holding.symbol
                if not self.insights.has_active_insights(symbol, self.utc_time):
                    self._undesired_symbols_from_previous_deployment.append(symbol)
            self._checked_symbols_from_previous_deployment = True
        
        for symbol in self._undesired_symbols_from_previous_deployment[:]:
            if self.is_market_open(symbol):
                self.liquidate(symbol, tag="Holding from previous deployment that's no longer desired")
                self._undesired_symbols_from_previous_deployment.remove(symbol)

# We re-define the MeanVarianceOptimizationPortfolioConstructionModel because
# - The model doesn't warm-up with ScaledRaw data (https://github.com/QuantConnect/Lean/issues/7239)
# - The original definition doesn't reset the `roc` and `window` in the `MeanVarianceSymbolData` objects when corporate actions occur

from AlgorithmImports import *

from Portfolio.MinimumVariancePortfolioOptimizer import MinimumVariancePortfolioOptimizer


### <summary>
### Provides an implementation of Mean-Variance portfolio optimization based on modern portfolio theory.
### The default model uses the MinimumVariancePortfolioOptimizer that accepts a 63-row matrix of 1-day returns.
### </summary>
class MeanVarianceOptimizationPortfolioConstructionModel(PortfolioConstructionModel):
    def __init__(self,
                 algorithm,
                 rebalance = Resolution.DAILY,
                 portfolio_bias = PortfolioBias.LONG_SHORT,
                 lookback = 1,
                 period = 63,
                 resolution = Resolution.DAILY,
                 target_return = 0.10,
                 optimizer = None):
        """Initialize the model
        Args:
            rebalance: Rebalancing parameter. If it is a timedelta, date rules or Resolution, it will be converted into a function.
                              If None will be ignored.
                              The function returns the next expected rebalance time for a given algorithm UTC DateTime.
                              The function returns null if unknown, in which case the function will be called again in the
                              next loop. Returning current time will trigger rebalance.
            portfolio_bias: Specifies the bias of the portfolio (Short, Long/Short, Long)
            lookback(int): Historical return lookback period
            period(int): The time interval of history price to calculate the weight
            resolution: The resolution of the history price
            optimizer(class): Method used to compute the portfolio weights"""
        super().__init__()
        self._algorithm = algorithm
        self._lookback = lookback
        self._period = period
        self._resolution = resolution
        self._portfolio_bias = portfolio_bias
        self._sign = lambda x: -1 if x < 0 else (1 if x > 0 else 0)

        lower = algorithm.settings.min_absolute_portfolio_target_percentage*1.1 if portfolio_bias == PortfolioBias.LONG else -1
        upper = 0 if portfolio_bias == PortfolioBias.SHORT else 1
        self._optimizer = MinimumVariancePortfolioOptimizer(lower, upper, target_return) if optimizer is None else optimizer

        self._symbol_data_by_symbol = {}
        self._new_insights = False

    def is_rebalance_due(self, insights, algorithmUtc):
        if not self._new_insights:
            self._new_insights = len(insights) > 0
        is_rebalance_due = self._new_insights and not self._algorithm.is_warming_up and self._algorithm.current_slice.quote_bars.count > 0
        if is_rebalance_due:
            self._new_insights = False
        return is_rebalance_due

    def create_targets(self, algorithm, insights):
        # Reset and warm-up indicators when corporate actions occur
        data = algorithm.current_slice
        reset_symbols = []
        for symbol in set(data.dividends.keys()) | set(data.splits.keys()):
            symbol_data = self._symbol_data_by_symbol[symbol]
            if symbol_data.should_reset():
                symbol_data.clear_history()
                reset_symbols.append(symbol)
        if reset_symbols:
            self._warm_up(algorithm, reset_symbols)

        return super().create_targets(algorithm, insights)

    def should_create_target_for_insight(self, insight):
        if len(PortfolioConstructionModel.filter_invalid_insight_magnitude(self._algorithm, [insight])) == 0:
            return False

        symbol_data = self._symbol_data_by_symbol.get(insight.symbol)
        if insight.magnitude is None:
            self._algorithm.set_run_time_error(ArgumentNullException('MeanVarianceOptimizationPortfolioConstructionModel does not accept \'None\' as Insight.magnitude. Please checkout the selected Alpha Model specifications.'))
            return False
        symbol_data.add(self._algorithm.time, insight.magnitude)

        return True

    def determine_target_percent(self, activeInsights):
        """
         Will determine the target percent for each insight
        Args:
        Returns:
        """
        targets = {}

        # If we have no insights just return an empty target list
        if len(activeInsights) == 0:
            return targets

        symbols = [insight.symbol for insight in activeInsights]

        # 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.return_ for symbol, data in self._symbol_data_by_symbol.items() if symbol in symbols }
        returns = pd.DataFrame(returns)

        # 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
        for insight in activeInsights:
            weight = weights[str(insight.symbol.id)]

            # don't trust the optimizer
            if self._portfolio_bias != PortfolioBias.LONG_SHORT and self._sign(weight) != self._portfolio_bias:
                weight = 0
            targets[insight] = weight

        return targets

    def on_securities_changed(self, algorithm, changes):
        # clean up data for removed securities
        super().on_securities_changed(algorithm, changes)
        for removed in changes.removed_securities:
            symbol_data = self._symbol_data_by_symbol.pop(removed.symbol, None)
            symbol_data.reset()

        # initialize data for added securities
        symbols = [x.symbol for x in changes.added_securities]
        for symbol in [x for x in symbols if x not in self._symbol_data_by_symbol]:
            self._symbol_data_by_symbol[symbol] = self.MeanVarianceSymbolData(symbol, self._lookback, self._period)
        self._warm_up(algorithm, symbols)
    
    def _warm_up(self, algorithm, symbols):
        history = algorithm.history[TradeBar](symbols, self._lookback * self._period + 1, self._resolution, data_normalization_mode=DataNormalizationMode.SCALED_RAW)
        for bars in history:
            for symbol, bar in bars.items():
                self._symbol_data_by_symbol.get(symbol).update(bar.end_time, bar.value)


    class MeanVarianceSymbolData:
        def __init__(self, symbol, lookback, period):
            self._symbol = symbol
            self._roc = RateOfChange(f'{symbol}.ROC({lookback})', lookback)
            self._roc.updated += self._on_rate_of_change_updated
            self._window = RollingWindow[IndicatorDataPoint](period)

        def should_reset(self):
            # Don't need to reset when the `window` only contain data from the insight.magnitude
            return self._window.samples < self._window.size * 2
        
        def clear_history(self):
            self._roc.reset()
            self._window.reset()

        def reset(self):
            self._roc.updated -= self._on_rate_of_change_updated
            self.clear_history()

        def update(self, time, value):
            return self._roc.update(time, value)

        def _on_rate_of_change_updated(self, roc, value):
            if roc.is_ready:
                self._window.add(value)

        def add(self, time, value):
            item = IndicatorDataPoint(self._symbol, time, value)
            self._window.add(item)

        # Get symbols' returns, we use simple return according to
        # Meucci, Attilio, Quant Nugget 2: Linear vs. Compounded Returns – Common Pitfalls in Portfolio Management (May 1, 2010). 
        # GARP Risk Professional, pp. 49-51, April 2010 , Available at SSRN: https://ssrn.com/abstract=1586656
        @property
        def return_(self):
            return pd.Series(
                data = [x.value for x in self._window],
                index = [x.end_time for x in self._window])

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

#region imports
from AlgorithmImports import *
#endregion
# 05/25/2023 -Set the universe data normalization mode to raw
#            -Added warm-up
#            -Made the following updates to the portfolio construction model:
#                - Added IsRebalanceDue to only rebalance after warm-up finishes and there is quote data
#                - Reset the MeanVarianceSymbolData indicator and window when corporate actions occur
#                - Changed the minimum portfolio weight to be algorithm.Settings.MinAbsolutePortfolioTargetPercentage*1.1 to avoid errors
#            -Adjusted the history requests to use scaled raw data normalization
#            https://www.quantconnect.com/terminal/processCache?request=embedded_backtest_587cc09bd82676a2ede5c88b100ef70b.html
#
# 07/13/2023: -Fixed warm-up logic to liquidate undesired portfolio holdings on re-deployment
#             -Set the MinimumOrderMarginPortfolioPercentage to 0
#             https://www.quantconnect.com/terminal/processCache?request=embedded_backtest_fa3146d7b1b299f4fc23ef0465540be0.html