#region imports
from AlgorithmImports import *
from market_surge_stochastic import MarketSurgeStochastic
from dateutil.relativedelta import relativedelta
#endregion

# consolicator example: https://github.com/QuantConnect/Lean/blob/master/Algorithm.Python/BasicTemplateFuturesConsolidationAlgorithm.py
class ValueAlphaModel(AlphaModel):    
    def __init__(self, stoch_k = 0, stoch_d = 0):
        '''Initializes a new instance of the ValueAlphaModel class
        Args:
            _securities: list of securities received from Universe Selection
            _stoch_k: the last _stoch_k from the stochastics indicator
            _stoch_d: the last _stoch_d from the stochastics indicator'''
        self._stoch_k = 0
        self._stoch_d = 0
        self._securities = []

        # Parameters for rules
        self._ticker = "TQQQ"
        self._over_sold = 12
        self._over_bought = 65

        # PD stop (per day?) percent.
        self._one_day_drop = 12.9

        # Local status
        self.move_above_delay = False
        self.cross_under_delay  = False
        self.can_rebuy = False

    # when resolution is daily, update() enters every day at 16:00, no matter whether it uses a weekly consolidator.
    # when resolution is hour, update() enters every day at [10, 16:00], no matter whether it uses a weekly consolidator.
    def update(self, algorithm: QCAlgorithm, data: Slice) -> List[Insight]:
        # algorithm.debug(" (david) update() enters: " + str(algorithm.time))

        # https://www.quantconnect.com/docs/v2/writing-algorithms/historical-data/warm-up-periods#03-Warm-Up-Vs-Reality
        # https://www.quantconnect.com/docs/v2/writing-algorithms/indicators/automatic-indicators#08-Warm-Up-Indicators
        if algorithm.is_warming_up:
            raise Exception("algorithm should have warmed up, but is_warming_up is actually false.")
        
        if not self._ms_sto.is_ready:
            raise Exception(f"{self._ms_sto.name} should have been ready, but it wasn't. The indicator received {self._ms_sto.samples} samples.")

        # It shows daily chart, because update() enters daily.
        # algorithm.plot("MarketSurgeStochastic", "fast_stoch", self._ms_sto.fast_stoch)
        algorithm.plot("MarketSurgeStochastic", "stoch_k", self._ms_sto.stoch_k)
        algorithm.plot("MarketSurgeStochastic", "stoch_d", self._ms_sto.stoch_d)

        insights = []

        # For corporate actions, returns directly.
        # https://www.quantconnect.com/docs/v2/writing-algorithms/securities/asset-classes/us-equity/corporate-actions
        if data[self._ticker] == None:
            return insights

        # Reset cross_under_delay and move_above_delay
        if self._ms_sto.stoch_k >= self._ms_sto.stoch_d:
            # because self._ms_sto.stoch_k >= self._ms_sto.stoch_d, so it's not cross under.
            self.cross_under_delay = False
        else:
            # because self._ms_sto.stoch_k < self._ms_sto.stoch_d, so it's not move above.
            self.move_above_delay = False

        # Buy Rule #1: Buy if the weekly fast line crosses above the slow line EOD, but only if the fast line is not below the Over Sold Line* (a variable). If it is below, wait to purchase.
        # The purchase will be on a Friday. If Friday is a holiday, purchase EOD Thursday. -- this cannot happen on Friday, because _ms_sto only updates on Monday. 
        if (not algorithm.portfolio[self._ticker].invested) and self._ms_sto.stoch_k >= self._ms_sto.stoch_d and (self._stoch_k < self._stoch_d or self.move_above_delay):
            if self._ms_sto.stoch_k >= self._over_sold:
                algorithm.debug(f" (david) update() {algorithm.time}: Buy Rule #1 triggers, create UP insights, self._ms_sto.stoch_k = {self._ms_sto.stoch_k}, self._ms_sto.stoch_d = {self._ms_sto.stoch_d}")
                # https://www.quantconnect.com/forum/discussion/16788/how-can-construct-portfolio-properly-without-re-balancing-if-there-wasn-039-t-signal-to-change-direction/p1/comment-48027
                # insights.append(Insight.price(self._ticker, Expiry.ONE_YEAR, InsightDirection.UP))
                insights.append(Insight.price("TQQQ", datetime.now() + relativedelta(years=100), InsightDirection.UP))
                self.move_above_delay = False
            else:
                # Wait until self._ms_sto.stoch_k moves above oversold
                self.move_above_delay = True

        # Re-buy rule: If you sell because of Sell Rules 2, 3, or 4, and you verify both Buy conditions are still true
        #   1. Fast Line is higher than the Slow Line
        #   2. Fast Line is above the Oversold Line
        # Re-Buy the next day at or near EOD for both the Daily and Weekly Strategies.
        # It's exactly on the next day.
        if self.can_rebuy and algorithm.time.date() == (self.sell_time + relativedelta(days=1)).date() and self._ms_sto.stoch_k >= self._ms_sto.stoch_d and self._ms_sto.stoch_k >= self._over_sold:
            algorithm.debug(f" (david) update() {algorithm.time}: Re-buy Rule triggers, create UP insights, self._ms_sto.stoch_k = {self._ms_sto.stoch_k}, self._ms_sto.stoch_d = {self._ms_sto.stoch_d}")
            # https://www.quantconnect.com/forum/discussion/16788/how-can-construct-portfolio-properly-without-re-balancing-if-there-wasn-039-t-signal-to-change-direction/p1/comment-48027
            # insights.append(Insight.price(self._ticker, Expiry.ONE_YEAR, InsightDirection.UP))
            insights.append(Insight.price("TQQQ", datetime.now() + relativedelta(years=100), InsightDirection.UP))
            self.can_rebuy = False
            self.move_above_delay = False

        # Sell Rule #1: Sell if the Weekly Fast Line drops below the Slow Line, but only if the Fast Line is lower than the Over Bought Line* (a variable), sell at EOD price.
        # If the Fast Line is above the Over Bought Line*, wait to sell. This will be on a Friday. If Friday is a holiday, sell EOD Thursday.
        if algorithm.insights.contains_key("TQQQ") and self._ms_sto.stoch_k < self._ms_sto.stoch_d and (self._stoch_k >= self._stoch_d or self.cross_under_delay):
            if self._ms_sto.stoch_k < self._over_bought:
                algorithm.debug(f" (david) update() {algorithm.time}: Sell Rule #1 triggers, cancel UP insights")
                # https://www.quantconnect.com/docs/v2/writing-algorithms/algorithm-framework/insight-manager#10-Cancel-Insights
                algorithm.insights.cancel(["TQQQ"])
                self.sell_time = algorithm.time
                self.cross_under_delay = False
                self.can_rebuy = False
            else:
                self.cross_under_delay = True

        # Sell Rule #2a: Sell at the Open if the Open Price is Below the 1-Day Drop Percentage* (a variable) Price, based on previous EOD.
        elif algorithm.insights.contains_key("TQQQ") and data["TQQQ"].open / self._daily_close_identity.current.price < (1 - self._one_day_drop / 100):
            algorithm.debug(f" (david) update() {algorithm.time}: Sell Rule #2a triggers, cancel UP insights")
            # https://www.quantconnect.com/docs/v2/writing-algorithms/algorithm-framework/insight-manager#10-Cancel-Insights
            algorithm.insights.cancel(["TQQQ"])
            self.sell_time = algorithm.time
            self.can_rebuy = True

        # The current resolution is hourly, so it cannot sell immediately. It sells when the hourly low is lower the threshold.
        # Sell Rule #2b: Sell intraday if the intraday price is Below the 1-Day Drop Percentage* (a variable) Price, based on previous EOD.
        elif algorithm.insights.contains_key("TQQQ") and (data["TQQQ"].low / self._daily_close_identity.current.price) < (1 - self._one_day_drop / 100):
            algorithm.debug(f" (david) update() {algorithm.time}: Sell Rule #2b triggers, cancel UP insights")
            # https://www.quantconnect.com/docs/v2/writing-algorithms/algorithm-framework/insight-manager#10-Cancel-Insights
            algorithm.insights.cancel(["TQQQ"])
            self.sell_time = algorithm.time
            self.can_rebuy = True

        self._stoch_k = self._ms_sto.stoch_k
        self._stoch_d = self._ms_sto.stoch_d
        return insights


    # on_securities_changed() enters only once, and it's before update().
    def on_securities_changed(self, algorithm: QCAlgorithm, changes: SecurityChanges) -> None:
        algorithm.debug(" (david) " + str(algorithm.time) + " on_securities_changed: changes.removed_securities=" + ', '.join(map(str, changes.removed_securities)) + "; changes.added_securities=" + ', '.join(map(str, changes.added_securities)))

        for security in changes.removed_securities:
            if security in self._securities:
                self._securities.remove(security)
        self._securities.extend(changes.added_securities)
        
        # https://www.quantconnect.com/forum/discussion/12428/how-to-compare-current-price-to-previous-day-039-s-close/p1/comment-36582
        # https://www.quantconnect.com/docs/v2/writing-algorithms/indicators/supported-indicators/identity
        self._daily_close_identity = algorithm.Identity(self._ticker, Resolution.Daily, Field.Close)
        algorithm.warm_up_indicator(self._ticker, self._daily_close_identity, Resolution.DAILY)

        # https://www.quantconnect.com/docs/v2/writing-algorithms/consolidating-data/consolidator-types/calendar-consolidators
        # self._consolidator = algorithm.consolidate(self._ticker, Calendar.WEEKLY, self._consolidation_handler)
        # self._consolidator = algorithm.create_consolidator(Calendar.WEEKLY, TradeBar) # API error
        self._consolidator = TradeBarConsolidator(Calendar.WEEKLY)

        # https://www.quantconnect.com/docs/v2/writing-algorithms/indicators/custom-indicators
        self._ms_sto = MarketSurgeStochastic("MarketSurgeStochastic", 5, 4, 3)
        # algorithm.register_indicator(self._ticker, self._ms_sto, Resolution.DAILY, self._consolidator)
        algorithm.register_indicator(self._ticker, self._ms_sto, self._consolidator)

        # https://www.quantconnect.com/docs/v2/writing-algorithms/indicators/automatic-indicators#08-Warm-Up-Indicators
        # This function will invoke MarketSurgeStochastic.update 5 times right away, before running assert() below.
        # Make sure that there are enough history before the start date, otherwise it won't work for warmup.
        # algorithm.warm_up_indicator(self._ticker, self._ms_sto, Resolution.Daily)
        # By using `timedelta(weeks=1)`, it will use weekly data to warm up it {warm_up_period} times.
        algorithm.warm_up_indicator(self._ticker, self._ms_sto, timedelta(weeks=1))

        # samples is a field defined in the parent class PythonIndicator.
        assert(self._ms_sto.samples >= 5), f"_ms_sto indicator was expected to have processed 5 datapoints already, but self._ms_sto.samples=" + str(self._ms_sto.samples)
        assert(self._ms_sto.is_ready), "_ms_sto indicator was expected to be ready"
    
    # It seems to be useless, but I have to define _consolidation_handler() for algorithm.consolidate()
    # Define the consolidation handler.
    def _consolidation_handler(self, consolidated_bar: TradeBar) -> None:
        # it does not have algorithm, so it cannot call algorithm.log().
        # self.log("_consolidation_handler enters:" + str(self.time))
        # self.log(str(consolidated_bar))
        pass

# region imports
from AlgorithmImports import *
from market_surge_stochastic import MarketSurgeStochastic

# from universe import LowPBRatioUniverseSelectionModel
from alpha import ValueAlphaModel
# from portfolio import EqualWeightingRebalanceOnInsightsPortfolioConstructionModel
# endregion

# 4/11/2022 Buy: i don't have this buy, becuase I'm selling on that day. why they buy on that day?
# 7/1/2022 Buy: i don't have this buy, why they buy on that day?
# 11/3/2022 Buy: i don't have this buy, why they buy on that day?
# 
class BuffetBargainHunterAlgorithm(QCAlgorithm):

    def initialize(self):
        # self.set_start_date(2022, 12, 26) # consolidator warm up starts from 2022/11/4
        # self.set_start_date(2022, 12, 27) # consolidator warm up starts from 2022/11/4, 
        self.set_start_date(2022, 12, 28) # consolidator warm up starts from 2022/11/7, Monday
        # self.set_start_date(2022, 12, 29) # consolidator warm up starts from 2022/11/8, Tuesday
        self.set_end_date(2024, 7, 23)
        self.set_cash(100000)
        self._ticker = "TQQQ"
        self._trailing_stop = 35.5
        self._profit_target = 33.6

        # `add_equity` may be duplicated with universe selection.
        # Default is ADJUSTED, it's preferred because it's easier to handle splits.
        # https://www.quantconnect.com/docs/v2/research-environment/datasets/us-equity#07-Data-Normalization
        # self._tqqq = self.add_equity(self._ticker, Resolution.DAILY, data_normalization_mode=DataNormalizationMode.RAW).symbol

        # seems to be useless
        # https://www.quantconnect.com/docs/v2/writing-algorithms/algorithm-framework/universe-selection/universe-settings#03-Leverage
        # self.universe_settings.leverage = 0

        # DAILY is not enough, becuase we need to sell on the open / intraday.
        # https://www.quantconnect.com/docs/v2/writing-algorithms/algorithm-framework/universe-selection/universe-settings#02-Resolution
        # self.universe_settings.resolution = Resolution.DAILY
        self.universe_settings.resolution = Resolution.HOUR
        # self.universe_settings.resolution = Resolution.MINUTE

        # The default is ADJUSTED, which seems to be much different from Yahoo.
        # SPLIT_ADJUSTED and TOTAL_RETURN are close to Yahoo.
        # https://www.quantconnect.com/docs/v2/research-environment/datasets/us-equity#07-Data-Normalization
        # https://www.quantconnect.com/docs/v2/writing-algorithms/universes/settings#07-Data-Normalization-Mode
        self.universe_settings.data_normalization_mode = DataNormalizationMode.RAW
        # self.universe_settings.data_normalization_mode = DataNormalizationMode.SPLIT_ADJUSTED
        # self.universe_settings.data_normalization_mode = DataNormalizationMode.TOTAL_RETURN
        
        tickers = [self._ticker]
        symbols = [Symbol.create(ticker, SecurityType.EQUITY, Market.USA) for ticker in tickers]
        self.add_universe_selection(ManualUniverseSelectionModel(symbols))

        # https://www.quantconnect.com/docs/v2/writing-algorithms/initialization#11-Set-Benchmark
        self.set_benchmark(self._ticker)

        ## Set trading fees to $0
        # https://www.quantconnect.com/docs/v2/writing-algorithms/reality-modeling/transaction-fees/key-concepts#02-Set-Models
        self.set_security_initializer(lambda security: security.set_fee_model(ConstantFeeModel(0)))

        self.add_alpha(ValueAlphaModel())

        # It does not create any orders, which is not what I want.
        # self.set_portfolio_construction(NullPortfolioConstructionModel())

        # EqualWeightingPortfolioConstructionModel sells when the insights is expired.
        # Set `rebalance = lambda time:None`, so that there is no rebalance.
        # https://www.quantconnect.com/docs/v2/writing-algorithms/algorithm-framework/portfolio-construction/supported-models#03-Equal-Weighting-Model
        self.set_portfolio_construction(EqualWeightingPortfolioConstructionModel(rebalance = lambda time:None))

        # When using `rebalance=lambda time:None`, these settings cannot be set as False, otherwise it will not create any orders.
        # https://www.quantconnect.com/docs/v2/writing-algorithms/algorithm-framework/portfolio-construction/key-concepts#08-Rebalance-Settings
        # Disable automatic portfolio rebalancing upon insight change, allowing for manual control over when portfolio adjustments are made based on insights.
        # self.settings.rebalance_portfolio_on_insight_changes = False
        # Disable automatic portfolio rebalancing upon security change, allowing for manual control over when portfolio adjustments are made based on security additions or removals.
        # self.settings.rebalance_portfolio_on_security_changes = False

        # https://www.quantconnect.com/docs/v2/writing-algorithms/algorithm-framework/risk-management/key-concepts#03-Multi-Model-Algorithms
        # https://www.quantconnect.com/docs/v2/writing-algorithms/algorithm-framework/risk-management/supported-models#07-Trailing-Stop-Model
        # https://github.com/QuantConnect/Lean/blob/master/Algorithm.Framework/Risk/TrailingStopRiskManagementModel.py
        # https://github.com/QuantConnect/Lean/blob/master/Algorithm.Framework/Risk/MaximumUnrealizedProfitPercentPerSecurity.py
        # self.add_risk_management(NullRiskManagementModel())
        self.add_risk_management(TrailingStopRiskManagementModel(maximum_drawdown_percent = self._trailing_stop / 100))
        self.add_risk_management(MaximumUnrealizedProfitPercentPerSecurity(maximum_unrealized_profit_percent = self._profit_target / 100))
        self.set_execution(ImmediateExecutionModel())

    # on_data() enters every day at 16:00, including warm up period.
    # def on_data(self, slice: Slice) -> None:
    #     if self._sto.is_ready:
    #         # The current value of self._sto is represented by self._sto.current.value
    #         self.plot("Stochastic", "sto", self._sto.current.value)
    #         # Plot all attributes of self._sto
    #         self.plot("Stochastic", "fast_stoch", self._sto.fast_stoch.current.value)
    #         self.plot("Stochastic", "stoch_k", self._sto.stoch_k.current.value)
    #         self.plot("Stochastic", "stoch_d", self._sto.stoch_d.current.value)

    # def on_data(self, data):
    #     if not self.is_warming_up:
    #         pass
        
    #     if not self._sto.is_ready:
    #         pass

    #     if not self.insights.has_active_insights(self._tqqq, self.time):
    #         pass

    #     # The current value of self._sto is represented by self._sto.current.value
    #     self.plot("Stochastic", "sto", self._sto.current.value)
    #     # Plot all attributes of self._sto
    #     self.plot("Stochastic", "fast_stoch", self._sto.fast_stoch.current.value)
    #     self.plot("Stochastic", "stoch_k", self._sto.stoch_k.current.value)
    #     self.plot("Stochastic", "stoch_d", self._sto.stoch_d.current.value)

# region imports
from AlgorithmImports import *
from collections import deque
# endregion

# Your New Python File

# https://www.quantconnect.com/docs/v2/writing-algorithms/indicators/custom-indicators
# https://github.com/QuantConnect/Lean/blob/master/Algorithm.Python/CustomIndicatorAlgorithm.py
# https://github.com/QuantConnect/Lean/blob/master/Algorithm.Python/CustomWarmUpPeriodIndicatorAlgorithm.py

######################################################
# Stochastic that matches MarketSmith
# declare lower;

# input period = 14;
# input k_period = 5;
# input d_period = 5;
# #Get the aggregation data from chart instead
# #input tsagg = AggregationPeriod.DAY;
# input upper = 80;
# input lower = 10;

# #def price = close(period = tsagg);
# def price = close(period = GetAggregationPeriod());

# def range = 100*(price - Lowest(price, period)) / (Highest(price, period) - Lowest(price, period));

# def k1 = 2 / (k_period + 1);
# def slow_k = RoundDown(range * k1, 0) + RoundDown(slow_k[1] * (1 - k1), 0); 
# slow_k[1] means the former vaule of slow_k

# def k2 = 2 / (d_period + 1);
# def slow_d = RoundDown(slow_k * k2, 0) + RoundDown(slow_d[1] * (1 - k2), 0);
# slow_d[1] means the former value of slow_d

# plot Up = upper;
# plot Lo = lower;

# plot Fast = K;
# plot Slow = D;
######################################################


class MarketSurgeStochastic(PythonIndicator):
    def __init__(self, name, period, k_period, d_period):
        # name, time, value are required for each indicator.
        self.name = name
        self.time = datetime.min
        self.value = 0
        
        # When warm_up_period=464, it start from 2010/2/11, which is a Thursday.
        self.warm_up_period = 464
        self.queue = deque(maxlen=period)
        self.period = period
        self.k_period = k_period
        self.d_period = d_period
        self.fast_stoch = 0
        self.stoch_k = 0
        self.stoch_d = 0

    def update(self, input: BaseData) -> bool:
        if not isinstance(input, TradeBar):
            raise TypeError('MarketSurgeStochastic.update: input must be a TradeBar')

        # `input.time` seems to be at 9:30 am on Monday, but the input.high/low/close are the prices of the week when I use RAW.
        self.time = input.time
        self.queue.appendleft(input.close)
        count = len(self.queue)

        if max(self.queue) == min(self.queue):
            # if there's no range, just return constant zero
            self.fast_stoch = 0
        else:
            self.fast_stoch = 100 * (input.close - min(self.queue)) / (max(self.queue) - min(self.queue))

        k1 = 2 / (self.k_period + 1)
        self.stoch_k = int(self.fast_stoch * k1) + int(self.stoch_k * (1 - k1))

        k2 = 2 / (self.d_period + 1)
        self.stoch_d = int(self.stoch_k * k2) + int(self.stoch_d * (1 - k2))

        # It does not have algorithm, so it cannot call algorithm.plot().
        # algorithm.plot("MarketSurgeStochastic", "stoch_k", self.stoch_k)
        # algorithm.plot("MarketSurgeStochastic", "stoch_d", self.stoch_d)

        # return if it's ready to use.
        return count == self.queue.maxlen

#region imports
from AlgorithmImports import *
#endregion


class EqualWeightingRebalanceOnInsightsPortfolioConstructionModel(EqualWeightingPortfolioConstructionModel):
    def __init__(self, algorithm):
        super().__init__()
        self._algorithm = algorithm
        self._new_insights = False

    def is_rebalance_due(self, insights: List[Insight], algorithm_utc: datetime) -> bool:
        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

#region imports
from AlgorithmImports import *
#endregion
# 05/24/2023: -Updated universe selection timing to run at the start of each month.
#             -Added warm-up.
#             -Removed the risk management model so the algorithm could warm-up properly.
#             -Added OnWarmupFinished to liquidate existing holdings that aren't backed by active insights.
#             https://www.quantconnect.com/terminal/processCache?request=embedded_backtest_210cae3651bd5aa87f79cd88b5f8109b.html
#
# 05/24/2023: -Updated universe selection timing so that the first trading day of each month always has the latest universe already selected.
#              Before this change, if the first trading day of the month was a Monday, the universe selection would run on Tuesday morning, 
#              cancelling some of the month's insights during warm-up. 
#             https://www.quantconnect.com/terminal/processCache?request=embedded_backtest_3806b7739d53f7b3687a948d0613cb25.html
#
# 05/26/2023: -Updated IsRebalanceDue function in the portfolio construction model to avoid MOO orders when deploying outside of regular trading hours
#             https://www.quantconnect.com/terminal/processCache?request=embedded_backtest_d15053c45241cf1803495ca83f5452c5.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_ac226bd3f017d0d0812d3b41fc25772a.html
#
# 10/27/2023: - Implement new Fundamental Universe Selection Model, merging coarse and fine selections
#
# 04/15/2023: -Updated to PEP8 style
#             -Added nan filter to universe selection
#             https://www.quantconnect.com/terminal/processCache?request=embedded_backtest_69e07892dfd0829a67f228ab5c7d7bdd.html

#region imports
from AlgorithmImports import *
#endregion


class LowPBRatioUniverseSelectionModel(FundamentalUniverseSelectionModel):
    def __init__(self, universe_settings: UniverseSettings = None, coarse_size: int = 1000, fine_size: int = 250) -> None:
        def select(fundamentals):
            fundamentals = [f for f in fundamentals if not np.isnan(f.valuation_ratios.pb_ratio)]
            sorted_by_dollar_volume = sorted(fundamentals, key=lambda c: c.dollar_volume, reverse=True)[:coarse_size]
            return [c.symbol for c in sorted(sorted_by_dollar_volume, key=lambda x: x.valuation_ratios.pb_ratio)[:fine_size]]
        super().__init__(None, select, universe_settings)