#region imports
from AlgorithmImports import *

from Execution.ImmediateExecutionModel import ImmediateExecutionModel
from Portfolio.EqualWeightingPortfolioConstructionModel import EqualWeightingPortfolioConstructionModel
from Risk.MaximumDrawdownPercentPerSecurity import MaximumDrawdownPercentPerSecurity
#endregion


"""
COMPREHENSIVE STRATEGY EXPLANATION

This strategy uses the QuantConnect Algorithm Framework. It selects stocks from a
fixed Dow Jones Industrial Average universe and invests in the highest P/E names
within that universe.

The framework structure is preserved:

1. Universe Selection
The strategy uses coarse and fine fundamental universe selection. The starting
universe is restricted to a static list of 30 DJIA constituents. This is a
simplified pedagogical DJIA universe and is not point-in-time adjusted. Once per
month, the universe selection process filters the DJIA names for valid fundamental
data, valid price data, and valid positive P/E ratios. It then selects the five
stocks with the highest P/E ratios, subject to a reasonable P/E range.

This is a growth-oriented screen, not a value screen. High P/E stocks are usually
companies where the market is pricing in stronger expected future growth.

2. Alpha Model
The custom alpha model emits long-only Up insights for the selected universe. It
also emits Flat insights for securities removed from the universe so the framework
can exit old positions.

3. Portfolio Construction
The strategy uses QuantConnect's EqualWeightingPortfolioConstructionModel. Active
long insights are equal-weighted across the portfolio.

4. Execution and Risk
Execution is handled by ImmediateExecutionModel. Risk management is handled by
MaximumDrawdownPercentPerSecurity. If a security breaches the drawdown threshold,
the risk model liquidates it and cancels related insights.

The benchmark is DIA, the Dow Jones Industrial Average ETF. This is appropriate
because the strategy selects from the DJIA universe.
"""


class MonthlyLongOnlySelectedSymbolsAlphaModel(AlphaModel):

    def __init__(self, insight_duration_days=35):
        self.insight_duration = timedelta(days=insight_duration_days)
        self.active_symbols = []
        self.removed_symbols = []
        self.last_emit_month = None

    def Update(self, algorithm, data):

        insights = []

        current_month = (algorithm.Time.year, algorithm.Time.month)

        # Emit Flat insights for securities removed from the selected universe.
        for symbol in self.removed_symbols:
            insights.append(
                Insight.Price(
                    symbol,
                    timedelta(days=1),
                    InsightDirection.Flat
                )
            )

        self.removed_symbols = []

        # Emit long insights only once per month.
        if self.last_emit_month == current_month:
            return insights

        if len(self.active_symbols) == 0:
            return insights

        for symbol in self.active_symbols:

            if algorithm.Securities.ContainsKey(symbol):
                if algorithm.Securities[symbol].HasData:

                    insights.append(
                        Insight.Price(
                            symbol,
                            self.insight_duration,
                            InsightDirection.Up
                        )
                    )

        self.last_emit_month = current_month

        return insights

    def OnSecuritiesChanged(self, algorithm, changes):

        for security in changes.AddedSecurities:

            if security.Symbol not in self.active_symbols:
                self.active_symbols.append(security.Symbol)

        for security in changes.RemovedSecurities:

            if security.Symbol in self.active_symbols:
                self.active_symbols.remove(security.Symbol)

            if security.Symbol not in self.removed_symbols:
                self.removed_symbols.append(security.Symbol)


class DJIATopPEStrategy(QCAlgorithm):

    def Initialize(self):

        # ------------------------------------------------------------
        # 1. BACKTEST SETTINGS
        # ------------------------------------------------------------
        self.SetStartDate(2018, 1, 1)
        self.SetEndDate(2026, 5, 5)

        self.initial_cash = 100000
        self.SetCash(self.initial_cash)

        # ------------------------------------------------------------
        # 2. STATIC DJIA UNIVERSE
        # ------------------------------------------------------------
        self.djia_tickers = [
            "AAPL", "AMGN", "AXP", "BA", "CAT", "CRM", "CSCO", "CVX", "DIS", "DOW",
            "GS", "HD", "HON", "IBM", "INTC", "JNJ", "JPM", "KO", "MCD", "MMM",
            "MRK", "MSFT", "NKE", "PG", "TRV", "UNH", "V", "VZ", "WBA", "WMT"
        ]

        self.djia_symbols = [
            Symbol.Create(ticker, SecurityType.Equity, Market.USA)
            for ticker in self.djia_tickers
        ]

        self.UniverseSettings.Resolution = Resolution.Daily

        self.top_count = 5
        self.last_selection_month = None

        # ------------------------------------------------------------
        # 3. FRAMEWORK UNIVERSE SELECTION
        # ------------------------------------------------------------
        self.SetUniverseSelection(
            FineFundamentalUniverseSelectionModel(
                self.CoarseSelectionFunction,
                self.FineSelectionFunction
            )
        )

        # ------------------------------------------------------------
        # 4. FRAMEWORK MODELS
        # ------------------------------------------------------------
        self.AddAlpha(
            MonthlyLongOnlySelectedSymbolsAlphaModel(
                insight_duration_days=35
            )
        )

        self.SetPortfolioConstruction(
            EqualWeightingPortfolioConstructionModel()
        )

        self.SetExecution(
            ImmediateExecutionModel()
        )

        self.SetRiskManagement(
            MaximumDrawdownPercentPerSecurity(0.15)
        )

        # ------------------------------------------------------------
        # 5. BENCHMARK
        # ------------------------------------------------------------
        self._benchmark = self.AddEquity(
            "DIA",
            Resolution.Daily,
            Market.USA
        ).Symbol

        self.SetBenchmark(self._benchmark)

        self.initial_benchmark_price = None

    # ------------------------------------------------------------
    # 6. COARSE SELECTION
    # ------------------------------------------------------------
    def CoarseSelectionFunction(self, coarse):

        current_month = (self.Time.year, self.Time.month)

        # Re-select only once per month.
        if self.last_selection_month == current_month:
            return Universe.Unchanged

        self.last_selection_month = current_month

        filtered = [
            c for c in coarse
            if c.Symbol in self.djia_symbols
            and c.HasFundamentalData
            and c.Price is not None
            and c.Price > 5
            and c.DollarVolume is not None
            and c.DollarVolume > 0
        ]

        return [
            c.Symbol
            for c in filtered
        ]

    # ------------------------------------------------------------
    # 7. FINE SELECTION
    # ------------------------------------------------------------
    def FineSelectionFunction(self, fine):

        candidates = []

        for stock in fine:

            pe_ratio = stock.ValuationRatios.PERatio

            if pe_ratio is None:
                continue

            # Avoid invalid or extreme P/E values.
            if pe_ratio <= 5:
                continue

            if pe_ratio >= 80:
                continue

            candidates.append(stock)

        if len(candidates) == 0:
            return Universe.Unchanged

        # Highest P/E first.
        sorted_by_pe = sorted(
            candidates,
            key=lambda f: f.ValuationRatios.PERatio,
            reverse=True
        )

        selected = [
            stock.Symbol
            for stock in sorted_by_pe[:self.top_count]
        ]

        selected_text = []

        for stock in sorted_by_pe[:self.top_count]:
            selected_text.append(
                stock.Symbol.Value
                + " PE "
                + str(round(stock.ValuationRatios.PERatio, 2))
            )

        self.Debug(
            "Monthly DJIA high PE selection "
            + str(selected_text)
        )

        return selected

    # ------------------------------------------------------------
    # 8. PLOTS
    # ------------------------------------------------------------
    def OnData(self, data):

        if self._benchmark not in data or data[self._benchmark] is None:
            return

        benchmark_price = self.Securities[self._benchmark].Price

        if benchmark_price <= 0:
            return

        if self.initial_benchmark_price is None:
            self.initial_benchmark_price = benchmark_price

        self.Plot(
            "Strategy Equity",
            "Portfolio Value",
            self.Portfolio.TotalPortfolioValue
        )

        benchmark_value = (
            self.initial_cash
            * benchmark_price
            / self.initial_benchmark_price
        )

        self.Plot(
            "Strategy Equity",
            "Buy Hold DIA",
            benchmark_value
        )