# region imports
from AlgorithmImports import *
from pandas.tseries.offsets import BDay
from dateutil.relativedelta import relativedelta
from dataclasses import dataclass
# endregion

class FOMC():
    def __init__(self, algorithm: QCAlgorithm, offset_days: int) -> None:
        csv_string_file = algorithm.Download(
            'data.quantpedia.com/backtesting_data/economic/fed_days.csv'
        )
        dates = csv_string_file.split('\r\n')
        
        self._dates: List[datetime.date] = [
            (datetime.strptime(x, "%Y-%m-%d") - BDay(offset_days)).date() for x in dates
        ]
        
    @property
    def Dates(self) -> List[datetime.date]:
        return self._dates

# NOTE: Manager for new trades. It's represented by certain count of equally weighted brackets for long and short positions.
# If there's a place for new trade, it will be managed for time of holding period.
class TradeManager():
    def __init__(self, 
                algorithm: QCAlgorithm, 
                long_size: int, 
                short_size: int, 
                holding_period: int,) -> None:
        self.algorithm: QCAlgorithm = algorithm  # algorithm to execute orders in.
        
        self.long_size: int = long_size
        self.short_size: int = short_size
        
        self.long_len: int = 0
        self.short_len: int = 0

        # Arrays of ManagedSymbols
        self.symbols: List[ManagedSymbol] = []
        
        self.holding_period: int = holding_period    # Days of holding.
    
    # Add stock symbol object
    def add(
        self, 
        symbol: Symbol, 
        long_flag: bool, 
        price: float) -> None:
    
        # Open new long trade.
        quantity: Union[None, int] = None
        
        if long_flag:
            # If there's a place for it.
            if self.long_len < self.long_size:
                quantity: int = int(self.algorithm.portfolio.total_portfolio_value / self.long_size / price)

                self.algorithm.market_order(symbol, quantity)
                self.long_len += 1
            # else:
            #     self.algorithm.Log("There's not place for additional trade.")

        # Open new short trade.
        else:
            # If there's a place for it.
            if self.short_len < self.short_size:
                quantity: int = int(self.algorithm.portfolio.total_portfolio_value / self.short_size / price)

                self.algorithm.market_order(symbol, -quantity)
                self.short_len += 1
            # else:
                # self.algorithm.Log("There's not place for additional trade.")
   
        if quantity:
            managed_symbol: ManagedSymbol = ManagedSymbol(symbol, self.algorithm.time + relativedelta(months=self.holding_period), quantity, long_flag)
            self.symbols.append(managed_symbol)

    # Decrement holding period and liquidate symbols.
    def try_liquidate(self) -> None:
        symbols_to_delete: List[ManagedSymbol] = []
        for managed_symbol in self.symbols:
            # Liquidate.
            if self.algorithm.time >= managed_symbol.liquidation_date:
                symbols_to_delete.append(managed_symbol)
                    
                if managed_symbol.long_flag:
                    self.algorithm.market_order(managed_symbol.symbol, -managed_symbol.quantity)
                    self.long_len -= 1
                else:
                    self.algorithm.market_order(managed_symbol.symbol, managed_symbol.quantity)
                    self.short_len -= 1
                
        # Remove symbols from management.
        for managed_symbol in symbols_to_delete:
            self.symbols.remove(managed_symbol)
    
    def liquidate_ticker(self, ticker: str) -> None:
        symbol_to_delete: Union[None, Symbol] = None
        for managed_symbol in self.symbols:
            if managed_symbol.symbol.Value == ticker:
                if managed_symbol.long_flag:
                    self.algorithm.market_order(managed_symbol.symbol, -managed_symbol.quantity)
                    self.long_len -= 1
                else:
                    self.short_len -= 1
                    self.algorithm.market_order(managed_symbol.symbol, managed_symbol.quantity)
                symbol_to_delete = managed_symbol
                
                break
        
        if symbol_to_delete: self.symbols.remove(symbol_to_delete)
        else: self.algorithm.Debug("Ticker is not held in portfolio!")

@dataclass   
class ManagedSymbol():
    symbol: Symbol 
    liquidation_date: datetime 
    quantity: int 
    long_flag: bool

# Custom fee model.
class CustomFeeModel(FeeModel):
    def GetOrderFee(self, parameters):
        fee = parameters.Security.Price * parameters.Order.AbsoluteQuantity * 0.00005
        return OrderFee(CashAmount(fee, "USD"))

# https://quantpedia.com/strategies/informative-price-pressure-during-pre-fomc-days
# 
# The investment universe for this strategy primarily consists of high-beta stocks and broad market ETFs like SPY. (The selection is based on the research paper's 
# findings that high-beta stocks exhibit more substantial hedging-induced price pressure and return predictability.)
# (Detailed variable definitions are in Appendix A.)
# Fundamental Recapitulation: The strategy uses historical pre-FOMC-meeting returns as key indicator. The methodology involves calculating the average pre-FOMC-meeting 
# return over the past 24 months for each FOMC Day (-1).
# Input Variable Definition: The pre-FOMC-meeting return is the stock market return on the day before the start of a FOMC meeting. To reduce estimation noise, 
# incorporate all rescheduled FOMC meetings over the past two years. Specifically, for each month m, the key independent variable, Average Past Day (−1) Return[m−23,m], 
# is computed as the simple average of pre-FOMC-meeting returns over the previous 24 months in percentage, covering the period from month m−23 to month m.
# Trading Process: The strategy involves rebalancing positions on each FOMC Day (-1):
# If there is a one standard deviation decrease in Average Past Day (-1) Return[m−23,m], invest in the broad market for up to 2 years,
# else, liquidate the position or hold cash.
# Weighting & Rebalancing: Trade one asset (ETF representing stock index: SPY, QQQ, DIA) or a market-weighted high beta stock basket. Rebalanced is explained in the 
# previous paragraph.
# 
# QC Implementation changes:
#   - Entry mean return threshold is set to -5%.

# region imports
from AlgorithmImports import *
import data_tools
from dateutil.relativedelta import relativedelta
# endregion

class InformativePricePressureDuringPreFOMCDays(QCAlgorithm):

    def initialize(self) -> None:
        self.set_start_date(2000, 1, 1)
        self.set_cash(100_000)

        FOMC_offset_days: int = 1
        long_count: int = 16
        self._period: int = 24  # months

        # historical values
        # NOTE these value were obtained from in-sample statistics evaluation
        # FOMCd-1 2y averages
        hist_ret_mean: float = 0.0011076739323817675
        hist_ret_std: float = 0.0028424951741022274

        # FOMCd-1
        # hist_ret_mean: float = 0.0016810414297294828
        # hist_ret_std: float = 0.013461857412966283

        # hist_ret_std = -0.00173482124172046
        self._threshold: float = hist_ret_mean - hist_ret_std
        # self._threshold: float = -0.05

        self._FOMC_averages: List[float] = []
        self._FOMC_obj = data_tools.FOMC(self, FOMC_offset_days)
        self._trade_manager: data_tools.TradeManager = data_tools.TradeManager(self, long_count, 0, self._period)

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

        self._rebalance_flag: bool = False
        self.schedule.on(
            self.date_rules.on(self._FOMC_obj.Dates),
            self.time_rules.before_market_close(self._market),
            self._trade
        )

    def on_data(self, slice: Slice) -> None:
        self._trade_manager.try_liquidate()

        # Rebalance on -1 FOMC days.
        if not self._rebalance_flag:
            return
        self._rebalance_flag = False

        # Get history and average performance on -1 FOMC days.
        history: DataFrame = self.history(self._market, start=self.time - relativedelta(months=self._period), end=self.time).unstack(level=0)
        if history.empty:
            return

        returns: DataFrame = history.close.pct_change().dropna()
        recent_dates: List[datetime.date] = [
            d for d in self._FOMC_obj.Dates 
            if (self.time - relativedelta(months=self._period)).date() <= d <= self.time.date()
        ]

        df_observed: DataFrame = returns.loc[returns.index.normalize().isin(recent_dates)]
        
        traded_asset: Symbol = self._market if df_observed.mean()[0] <= self._threshold else self._cash

        # Trade execution.
        if slice.contains_key(traded_asset) and slice[traded_asset]:
            self._trade_manager.add(traded_asset, True, slice[traded_asset].price)

    def _trade(self) -> None:
        if all(self.securities[symbol].get_last_data() for symbol in [self._market, self._cash]):
            self._rebalance_flag = True