# region imports
from AlgorithmImports import *
from pandas.tseries.offsets import BDay
# 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

# 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/inflation-risk-premium-in-stocks-during-fomc-announcements/
# 
# The investment universe for this strategy consists of U.S. equities: NYSE, AMEX, and NASDAQ-listed common stocks, excluding financial and utility firms, to 
# avoid sector-specific biases. Furthermore, it excludes stocks with closing prices lower than $5 per share.
# (CRSP and Compustat are the primary data sources. Daily inflation swap data from Bloomberg are not needed for this variant.)
# Rationale: The trading rules are centered around profit margins and markups as price stickiness and inflation risk exposure indicators. Before pre-scheduled 
# FOMC announcement days, portfolios are adjusted to reflect anticipated inflation shocks. Specifically, firms in the lowest profit margin quintile, which are 
# more sensitive to inflation shocks, are overweighted. After FOMC announcements, changes in inflation swap rates are analyzed to adjust portfolio weights 
# accordingly. If inflation expectations rise, exposure to firms with low-profit margins is increased; if expectations fall, exposure is reduced and shifted 
# toward firms with high-profit margins.
# Ranking: Individual instruments are selected based on profit margins and markups, calculated using operating income and sales data from financial statements. 
# Based on these metrics, firms are ranked within their respective industries, defined by four-digit SIC codes. This ranking helps identify firms with varying 
# levels of price stickiness, which is crucial for assessing their exposure to inflation risk, as described later.
# Sorting: Before each pre-scheduled FOMC announcement day, we sort stocks into quintile portfolios based on their intra-industry profit margin. We hold these
# portfolios until two days before the following FOMC announcement.
# Strategy Execution: Construct a long-short (5-1) portfolio:
# Go long (buy) top quintile (5), and short (sell) the bottom one (1).
# As mentioned, hold portfolios until two days before the following pre-scheduled FOMC announcements.
# Rebalancing & Weighting: Maintain value-weighted portfolios. Repeat this whole procedure and add resort stocks to portfolios before each FOMC announcement day 
# for the following periods in which you intend to execute the strategy.
# 
# Implementation changes:
#   - Universe consists of 3000 largest stocks from NYSE, AMEX and NASDAQ.
#   - Profit margins are calculated relative to industry classification.

# region imports
from AlgorithmImports import *
from numpy import isnan
from functools import reduce
import data_tools
from typing import List, Dict, Set
# endregion

class InflationRiskPremiumInStocksDuringFOMCAnnouncements(QCAlgorithm):

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

        self._exchange_codes: List[str] = ['NYS', 'NAS', 'ASE']
        self._tickers_to_ignore: List[str] = ['KELYB']

        self._quantile: int = 5
        self._min_share_price: int = 5
        FOMC_offset_days: int = 4
        leverage: int = 5

        market: Symbol = self.add_equity('SPY', Resolution.DAILY).symbol
        self._FOMC_obj = data_tools.FOMC(self, FOMC_offset_days)

        self._weight: Dict[Symbol, float] = {}
        
        sector_code_flag: bool = False
        sector_code_str: str = 'asset_classification.morningstar_sector_code'
        industry_code_str: str = 'asset_classification.morningstar_industry_group_code'
        self._morningstar_code: str = sector_code_str if sector_code_flag else industry_code_str

        self._fundamental_count: int = 3_000
        self._fundamental_sorting_key = lambda x: x.market_cap

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

        self.schedule.on(
            self.date_rules.on(self._FOMC_obj.Dates),
            self.time_rules.before_market_close(market),
            self._trade
        )

    def fundamental_selection_function(self, fundamental: List[Fundamental]) -> List[Symbol]:
        if not self._selection_flag:
            return Universe.Unchanged

        selected: List[Fundamental] = [
            x for x in fundamental 
            if x.has_fundamental_data 
            and x.market == 'usa' 
            and x.market_cap != 0
            and x.price >= self._min_share_price
            and x.asset_classification.morningstar_sector_code != MorningstarSectorCode.FINANCIAL_SERVICES
            and not isnan(self._rgetattr(x, self._morningstar_code)) and self._rgetattr(x, self._morningstar_code) != 0
            and not isnan(x.financial_statements.income_statement.operating_income.three_months) and x.financial_statements.income_statement.operating_income.three_months != 0
            and not isnan(x.financial_statements.income_statement.total_revenue.three_months) and x.financial_statements.income_statement.total_revenue.three_months != 0
            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]]

        # Get mean of profit margin by industry code.
        industries: Set[int] = set([self._rgetattr(x, self._morningstar_code) for x in selected])
        industry_profit_margins: Dict[int, float] = {
            industry : np.mean([
                (
                    stock.financial_statements.income_statement.operating_income.three_months / 
                    stock.financial_statements.income_statement.total_revenue.three_months
                ) 
                for stock in selected 
                if self._rgetattr(stock, self._morningstar_code) == industry
            ]) 
            for industry in industries
        }

        long: List[Fundamental] = []
        short: List[Fundamental] = []

        # # Calculate relative profit margin.
        for industry in industries:
            relative_profit_margin: Dict[Fundamental, float] = {
                stock: (
                    stock.financial_statements.income_statement.operating_income.three_months / 
                    stock.financial_statements.income_statement.total_revenue.three_months - 
                    industry_profit_margins[industry]
                ) 
                for stock in selected
                if self._rgetattr(stock, self._morningstar_code) == industry
            }

            # Sort and divide stocks.
            if len(relative_profit_margin) >= self._quantile:
                sorted_by_relative_profit_margin: List[Tuple[Fundamental, float]] = sorted(relative_profit_margin.items(), key = lambda x: x[1], reverse=True)
                quintile: int = int(len(sorted_by_relative_profit_margin) / self._quantile)
                long.extend([x[0] for x in sorted_by_relative_profit_margin[:quintile]])
                short.extend([x[0] for x in sorted_by_relative_profit_margin[-quintile:]])

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

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

        # Trade execution.
        portfolio: List[PortfolioTarget] = [
            PortfolioTarget(symbol, w) for symbol, w in self._weight.items() 
            if slice.contains_key(symbol) 
            and slice[symbol]
        ]
        self.set_holdings(portfolio, True)
        self._weight.clear()

    def _trade(self) -> None:
        self._selection_flag = True

    # https://gist.github.com/wonderbeyond/d293e7a2af1de4873f2d757edd580288
    def _rgetattr(self, obj, attr, *args):
        def _getattr(obj, attr):
            return getattr(obj, attr, *args)
        return reduce(_getattr, [obj] + attr.split('.'))