# region imports
from AlgorithmImports import *
# endregion

class SymbolData():
    def __init__(self, period: int) -> None:
        self._prices: RollingWindow = RollingWindow[float](period)

    def update_price(self, price: float) -> None:
        self._prices.add(price)

    def get_returns(self, period) -> np.ndarray:
        return pd.Series(list(self._prices)[::-1][:period]).pct_change().dropna()

    def get_MAX(self, period: int) -> float:
        returns: np.ndarray = self.get_returns(period)
        return max(returns)

    def get_momentum(self, period: List[int]) -> float:
        return self._prices[period[1]] / self._prices[period[0] - 1] - 1

    def is_ready(self) -> bool:
        return self._prices.is_ready

# custom fee model
class CustomFeeModel(FeeModel):
    def get_order_fee(self, parameters) -> OrderFee:
        fee: float = parameters.security.price * parameters.order.absolute_quantity * 0.00005
        return OrderFee(CashAmount(fee, "USD"))

# https://quantpedia.com/strategies/switching-between-momentum-and-reversal-strategies-in-equities
# 
# The investment universe for this strategy includes all common stocks listed on the NYSE, AMEX, and NASDAQ, as indicated by share codes 10 and 11 in the CRSP dataset.
# (The study employs an extensive dataset from the Center for Research in Security Prices (CRSP).)
# Fundamental Recapitulation: To introduce market-state dependency, define the market condition using the sign of the past 24‑month cumulative return of the CRSP 
# value‑weighted index (B. Strategy Construction, pg. 5). Under this approach, a positive cumulative return indicates a bullish market state that triggers the 
# traditional momentum strategy (long winners, short losers). In contrast, a negative cumulative return signals a bearish market state, prompting a reversal strategy 
# (long losers, short winners).
# Sorting: Individual instruments are selected based on their past returns. Stocks are ranked into deciles according to their prior 12-month returns. This ranking is 
# used to identify the top decile (winners) and the bottom decile (losers) for constructing the portfolio.
# Decisional Process: The strategy uses the cumulative return of the CRSP value-weighted index over the past 24 months to classify market states. A positive cumulative 
# return indicates a positive market state, while a negative cumulative return signals a negative market state. In a positive market state, the strategy follows a 
# momentum approach by taking long positions in the top decile of stocks and short positions in the bottom decile. Conversely, in a negative market state, the strategy 
# adopts a reversal approach by taking long positions in the bottom decile and short positions in the top decile.
# Strategy Execution: A zero-cost momentum strategy involves taking long positions in the top decile (winners) and short positions in the bottom decile (losers). 
# Reversal requires opposite directional trades.
# Rebalancing & Weighting: The strategy involves rebalancing the portfolio monthly, with positions held for three months before liquidation. This creates an overlapping 
# portfolio framework. Maintain value-weighted portfolios at all times.
# 
# QC Implementation changes:
#   - Universe consists of 3000 largest from NYSE, AMEX and NASDAQ.
#   - Momentum period is set to 6-1 months.

# region imports
from AlgorithmImports import *
from pandas.core.frame import DataFrame
from dateutil.relativedelta import relativedelta
from dataclasses import dataclass
import data_tools
# endregion

class SwitchingBetweenMomentumAndReversalStrategiesInEquities(QCAlgorithm):

    def initialize(self) -> None:
        self.set_start_date(2000, 1, 1)
        self.set_cash(1_000_000)
        
        self._tickers_to_ignore: List[str] = ['DCIX', 'PIXY']
        self._exchange_codes: List[str] = ['NYS', 'NAS', 'ASE']    

        self._mom_period: List[int] = [6, 1]
        self._period: int = 12
        self._market_performance_period: int = 24
        self._quantile: int = 10
        self._holding_period: int = 3 # Months
        self._min_share_price: int = 5

        self._data: Dict[Symbol, data_tools.SymbolData] = {}
        self._portions: Dict[Symbol, float] = {}
        self._stock_holdings: List[HoldingItem] = []

        self._market: Symbol = self.add_equity('SPY', Resolution.DAILY).symbol

        self.fundamental_count: int = 3_000
        self.fundamental_sorting_key = lambda x: x.market_cap

        self._selection_flag: bool = False
        self.settings.minimum_order_margin_portfolio_percentage = 0
        self.settings.daily_precise_end_time = False
        self.universe_settings.resolution = Resolution.DAILY
        self.universe_settings.leverage = 10
        self.add_universe(self._fundamental_selection_function)
        self.set_security_initializer(lambda security: security.set_fee_model(data_tools.CustomFeeModel()))

        self.schedule.on(
            self.date_rules.month_start(self._market), 
            self.time_rules.after_market_open(self._market), 
            self._selection)

    def _fundamental_selection_function(self, fundamental: Iterable[Fundamental]) -> Iterable[Symbol]:
        # Monthly selection.
        if not self._selection_flag:
            return Universe.UNCHANGED

        # Store monthly stock prices.
        for stock in fundamental:
            symbol: Symbol = stock.symbol

            if symbol in self._data:
                self._data[symbol].update_price(stock.adjusted_price)
        
        selected: List[Fundamental] = [
            x for x in fundamental 
            if x.has_fundamental_data
            and x.market_cap != 0
            and x.market == 'usa'
            and x.price > self._min_share_price
            and x.security_reference.exchange_id in self._exchange_codes
            and x.symbol.value not in self._tickers_to_ignore
        ]

        if len(selected) > self.fundamental_count:
            selected = [x for x in sorted(selected, key=self.fundamental_sorting_key, reverse=True)[:self.fundamental_count]]

        # Warmup price rolling windows.
        for stock in selected:
            symbol = stock.symbol
            if symbol in self._data:
                continue
            
            self._data[symbol] = data_tools.SymbolData(self._period) 

            history: DataFrame = self.history(symbol, self._period, Resolution.DAILY)
            if history.empty:
                self.log(f"Not enough data for {symbol} yet.")
                continue
            closes: pd.Series = history.loc[symbol].close
            for time, close in closes.items():
                self._data[symbol].update_price(close)

        # Evaluate signal based on market performance.
        market_history: DataFrame | Quantconnect.Series = self.history(
            self._market,
            start=self.time - relativedelta(months=self._market_performance_period),
            end=self.time
        ).unstack(level=0)
        if not market_history.empty:
            monthly_closes: DataFrame = market_history.groupby(pd.Grouper(freq='MS')).first()
            if len(monthly_closes) < self._market_performance_period:
                self.log('Not enough data for signal evaluation.')

        market_performance: float = monthly_closes.iloc[-1].values[0] / monthly_closes.iloc[0].values[0] - 1

        momentum_signal: bool = True if market_performance > 0 else False

        # Sort and divide based on signal.
        sorted_by_momentum: List[tuple[Fundamental, Any]] = sorted(
            {
                stock: self._data[stock.symbol].get_momentum(self._mom_period)
                for stock in selected
                if self._data[stock.symbol].is_ready()
            }.items(),
            key=lambda x: x[1],
            reverse=momentum_signal
        )

        if len(sorted_by_momentum) < self._quantile:
            return Universe.UNCHANGED

        quantile: int = len(sorted_by_momentum) // self._quantile
        
        long: List[Fundamental] = [x[0] for x in sorted_by_momentum][:quantile]
        short: List[Fundamental] = [x[0] for x in sorted_by_momentum][-quantile:]

        # Value weighting.
        for i, portfolio in enumerate([long, short]):
            mc_sum: float = sum(list(map(lambda stock: stock.market_cap, portfolio)))
            for stock in portfolio:
                if stock.symbol not in self._portions:
                    self._portions[stock.symbol] = (
                        ((-1) ** i) * stock.market_cap / mc_sum
                    ) * (self.portfolio.total_portfolio_value / self._holding_period)

        return list(self._portions.keys())

    def on_data(self, slice: Slice) -> None:
        if not self._selection_flag:
            return
        self._selection_flag = False

        # Hold for 3 months, then liquidate.
        items_to_remove:List[HoldingItem] = []
        for item in self._stock_holdings:
            item.holding_period += 1
            if item.holding_period >= self._holding_period:
                self.market_order(item.symbol, -item.quantity)
                items_to_remove.append(item)   

        # Remove from collection.
        for item in items_to_remove:
            self._stock_holdings.remove(item)  

        # Trade execution.
        for symbol, portion in self._portions.items():
            if slice.contains_key(symbol) and slice[symbol]:
                quantity: int = portion // slice[symbol].price
                if quantity != 0:
                    self.market_order(symbol, quantity)
                    self._stock_holdings.append(HoldingItem(symbol, quantity))

        self._portions.clear()
    
    def _selection(self) -> None:
        self._selection_flag = True

@dataclass
class HoldingItem():
    symbol: Symbol
    quantity: int
    holding_period: int = 0