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

# Source: https://www.finra.org/rules-guidance/key-topics/margin-accounts/margin-statistics
class MarginDebt(PythonData):
    _last_update_date:datetime.date = datetime(1,1,1).date()

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

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

    def Reader(self, config: SubscriptionDataConfig, line: str, date: datetime, isLiveMode: bool) -> BaseData:
        data = MarginDebt()
        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[0], "%Y-%m-%d") + relativedelta(months=1)
        if split[-1] != '.':
            data.Value = float(split[-1])

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

# https://quantpedia.com/strategies/margin-debt-market-timing-strategy
# 
# The investment universe for this strategy consists of two primary instruments: SPY ETF, which serves as a proxy for the stock market performance, and SHY ETF, 
# representing a low-risk, cash-like investment. SPY is selected based on its role as a widely recognized benchmark for the U.S. stock market, while SHY is chosen 
# for its stability and low risk, making it suitable for holding cash positions. The strategy also incorporates margin debt data as a key indicator, sourced from 
# FINRA, which provides insights into investor leverage and sentiment.
# The strategy employs moving averages as the primary tool for generating trading signals. Specifically, it uses 6 to 12-month moving averages of both SPY prices 
# and margin debt. The methodology involves calculating these moving averages monthly and comparing the latest SPY price and margin debt value to their respective 
# moving averages. A "Buy" signal is generated when both the SPY price and margin debt exceed their moving averages, indicating a favorable market condition. 
# Conversely, a "Hold Cash" signal is triggered when either the SPY price or margin debt falls below their moving averages, suggesting a more cautious approach.
# The strategy involves monthly rebalancing, where capital is allocated based on the latest signals. Positions are equally distributed across the different moving 
# average periods (6 to 12 months) to diversify risk and reduce the impact of any single period's underperformance. Risk management is achieved through this 
# diversification and the use of SHY ETF to hold cash positions during unfavorable market conditions. This approach helps mitigate potential losses and stabilizes 
# returns by ensuring that not all capital is exposed to market risk at any given time.

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

class MarginDebtMarketTimingStrategy(QCAlgorithm):

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

        self._SMA_range: range = range(6, 13)

        self._market: Symbol = self.add_equity('SPY', Resolution.DAILY).symbol
        self._cash: Symbol = self.add_equity('SHY', Resolution.DAILY).symbol
        self._margin_debt: Symbol = self.add_data(data_tools.MarginDebt, 'MARGIN_DEBT', Resolution.DAILY).symbol

        self._rebalance_flag: bool = False
        self.schedule.on(
            self.date_rules.month_start(self._market),
            self.time_rules.after_market_open(self._market),
            self._rebalance
        )

        self.settings.daily_precise_end_time = False

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

        # Chech if data is still coming.
        if self.securities[self._margin_debt].get_last_data() and self.time.date() > data_tools.MarginDebt.get_last_update_date():
            self.log('Margin Debt data stopped coming.')
            return

        # SPY and margin debt history.
        spy_history: DataFrame = self.history(self._market, start=self.time - relativedelta(months=max(self._SMA_range)), end=self.time).unstack(level=0)
        margin_debt_history: DataFrame = self.history(self._margin_debt, start=self.time - relativedelta(months=max(self._SMA_range) + 1), end=self.time).unstack(level=0)

        if any(df.empty for df in [spy_history, margin_debt_history]):
            self.log('Insufficient data for further calculation.')
            return

        spy_history_prices: DataFrame = spy_history.close.groupby(pd.Grouper(freq='MS')).first()
        margin_debt_history = margin_debt_history.value

        spy_allocation, shy_allocation = .0, .0

        # Evaluate signals using the SMA.
        for SMA_period in self._SMA_range:
            signals: List[bool] = []
            for df in [spy_history_prices, margin_debt_history]:
                cutoff_date: datetime = (self.time - relativedelta(months=SMA_period)).replace(day=1)
                observed_df: DataFrame = df.loc[cutoff_date:]
                signals.append(True if self.securities[df.columns[0]].get_last_data().price > observed_df.mean()[0] else False)
            if all(signals):
                spy_allocation += 1 / len(self._SMA_range)
            else:
                shy_allocation += 1 / len(self._SMA_range)

        # Trade execution.
        targets: List[PortfolioTarget] = []
        for symbol, allocation in [(self._market, spy_allocation), (self._cash, shy_allocation)]:
            if slice.contains_key(symbol) and slice[symbol]:
                targets.append(PortfolioTarget(symbol, allocation))
        
        self.set_holdings(targets, True)

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