# region imports
from AlgorithmImports import *
# endregion

class SymbolData():
    def __init__(self, period: int, mom_period) -> None:
        self._prices: RollingWindow = RollingWindow[float](mom_period) 
        self._cpi_mom_change: RollingWindow = RollingWindow[float](period)
        self._last_inflation_change: int = 0
    
    def price_is_ready(self) -> bool:
        return self._prices.is_ready

    def is_ready(self) -> bool:
        return self._cpi_mom_change.is_ready
    
    def update_price(self, price: float) -> None:
        self._prices.add(price)

    def update(self, cpi_value: float) -> None:
        self._cpi_mom_change.add(cpi_value)

    def inflation_increase(self) -> bool:
        inflation_change: int = (
            1 
            if all(x > 0 for x in list(self._cpi_mom_change)) 
            else (
                -1 if all(x < 0 for x in list(self._cpi_mom_change)) 
                else self._last_inflation_change
            )
        )
        self._last_inflation_change = inflation_change
        return True if inflation_change == 1 else False

    def positive_momentum(self) -> bool:
        momentum: float = self._prices[0] / self._prices[self._prices.count - 1] - 1
        return True if momentum > 0 else False

    # def get_inflation_trend(self) -> List[int]:
    #     return [self.get_inflation_change(), self.get_momentum()]

class BLS_CPI(PythonData):
    _last_update_date: datetime.date = datetime(1,1,1).date()

    @staticmethod
    def get_last_update_date() -> datetime.date:
       return BLS_CPI._last_update_date

    def GetSource(self, config: SubscriptionDataConfig, date: datetime, isLiveMode: bool) -> SubscriptionDataSource:
        return SubscriptionDataSource(f'data.quantpedia.com/backtesting_data/economic/BLS_INFLATION_AS_REPORTED.csv', SubscriptionTransportMedium.RemoteFile, FileFormat.Csv)

    def Reader(self, config: SubscriptionDataConfig, line: str, date: datetime, isLiveMode: bool) -> BaseData:
        data = BLS_CPI()
        data.Symbol = config.Symbol

        if not line[0].isdigit(): return None
        split = line.split(';')
        
        # Parse the CSV file's columns into the custom data class
        data.Time = datetime.strptime(split[1], "%Y-%m-%d") #+ timedelta(days=1)
        if split[-1] != '':
            data.Value = float(split[-1])

        if data.Time.date() > BLS_CPI._last_update_date:
            BLS_CPI._last_update_date = data.Time.date()
        
        return data

# https://quantpedia.com/strategies/using-inflation-data-for-systematic-gold-and-treasury-investment-strategies/
# 
# The investment universe for this strategy includes gold, treasury bonds, and cash equivalents, explicitly using the GLD, IEF, and SHY ETFs.
# (Asset Selection Explanation: These instruments are selected based on their sensitivity to inflation dynamics as discussed in the research paper. Gold is chosen 
# for its traditional role as an inflation hedge, while treasury bonds are included due to their complex relationship with inflation and interest rates. The cash 
# equivalent (SHY) is used for capital preservation when no other assets show favorable performance.)
# (Data set includes the Consumer Price Index for All Urban Consumers: All Items in U.S. City Average (CPIAUCSL) from FRED (Federal Reserve Bank of St. Louis) and 
# the U.S. Bureau of Labor Statistics inflation data. Price data from Yahoo Finance.)
# Final Model Trading Strategy Variant: The strategy employs inflation data and momentum indicators to guide trading decisions. Inflation regimes are identified 
# using the rate of change in inflation metrics like the CPI. If the inflation rate increases, the regime is classified as Inflation UP; otherwise, it’s Inflation 
# DOWN. A 12-month momentum indicator is calculated for each asset to assess its performance trend.
# Buy and sell rules are based on these signals:
# In an Inflation UP and Momentum UP regime, buy GLD (if it has positive momentum),
# in an Inflation DOWN and Momentum UP regime, buy the treasury bond ETF IEF (it has positive momentum),
# if neither condition is fulfilled, allocate capital to SHY.
# Rebalancing & Weighting: The strategy involves monthly rebalancing to adjust positions based on the latest inflation and momentum signals. Position size assumes 
# full allocation to one asset only at the time.

# region imports
from AlgorithmImports import *
import data_tools
# endregion

class UsingInflationDataForSystematicGoldAndTreasuryInvestmentStrategies(QCAlgorithm):

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

        leverage: int = 4
        period: int = 2
        mom_period: int = 12

        tickers: List[str] = ['GLD', 'IEF', 'SHY']

        self._gld, self._ief, self._shy = [
            self.add_equity(ticker, Resolution.DAILY, leverage=leverage).symbol for ticker in tickers
        ]
        self._CPI: Symbol = self.add_data(data_tools.BLS_CPI, 'BLSCPI', Resolution.DAILY).symbol

        self._data: Dict[Symbol, data_tools.SymbolData] = {}

        for symbol in [self._gld, self._ief, self._CPI]:
            self._data[symbol] = data_tools.SymbolData(period, mom_period)

        self.settings.daily_precise_end_time = False
        self.settings.minimum_order_margin_portfolio_percentage = 0.

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

        self._rebalance_flag: bool = False
        self.schedule.on(
            self.date_rules.month_start(market),
            self.time_rules.at(0, 0),
            self._rebalance
        )

    def on_data(self, slice: Slice):
        # Monthly rebalance.
        if not self._rebalance_flag:
            return
        self._rebalance_flag = False

        # Check if data is still coming.
        if self.securities[self._CPI].get_last_data() and self.time.date() > data_tools.BLS_CPI.get_last_update_date():
            self.log('Data for CPI stopped coming.')
            return

        if not all(self.securities[symbol].get_last_data() for symbol in [self._gld, self._ief, self._shy, self._CPI]):
            return

        # Update data.
        for symbol, symbol_data in self._data.items():
            if symbol == self._CPI:
                symbol_data.update(self.securities[symbol].get_last_data().price)
            else:
                symbol_data.update_price(self.securities[symbol].get_last_data().price)

        if not all(self._data[symbol].price_is_ready() for symbol in [self._gld, self._ief]) and not self._data[self._CPI].is_ready():
            return

        # Trade asset decision.
        traded_asset: Symbol = (
            self._gld 
            if self._data[self._CPI].inflation_increase() 
            and self._data[self._gld].positive_momentum()
            else (
                self._ief 
                if not self._data[self._CPI].inflation_increase()
                and self._data[self._ief].positive_momentum()
                else self._shy
            )
        )

        # Trade execution.
        self.set_holdings(traded_asset, 1, True)

    def _rebalance(self) -> None:
        self._rebalance_flag = True