plot_monthly_holdings_example.py

from __future__ import annotations

import argparse
import os

import matplotlib.pyplot as plt
import pandas as pd

from Backtest_utils import BacktestConfig
from Dynamic_stock_selection import DynamicStockSelector
from Optimize_Portfolio import PortfolioOptimizer
from Regression import FF5Regression
from main_update import (
    COV_LOOKBACK,
    CLUSTER_METHOD,
    HMM_COVARIANCE_FLOOR,
    HMM_N_ITER,
    HMM_STICKINESS,
    HMM_TOL,
    IDIO_VAR_FLOOR,
    MIN_COV_OBS,
    RANDOM_STATE,
    REALIZED_SHIFT_MONTHS,
    RISK_MODEL,
    SIGNAL_SHIFT_MONTHS,
    build_regime_model,
    load_dual_clock_inputs,
)

DEFAULT_MONTH = None
SAVE_FIGURE = False
FACTOR_COLS = ["Mkt-RF", "SMB", "HML", "RMW", "CMA"]


def build_month_snapshot(holding_month: pd.Timestamp, config: BacktestConfig) -> tuple[pd.Series, int, pd.Timestamp]:
    price_df, returns_df, ff5_realized_df, ff5_signal_df, _ = load_dual_clock_inputs(
        realized_shift_months=REALIZED_SHIFT_MONTHS,
        signal_shift_months=SIGNAL_SHIFT_MONTHS,
    )
    selector = DynamicStockSelector()
    optimizer = PortfolioOptimizer(returns_df, factor_df=ff5_realized_df)

    common_dates = returns_df.index.intersection(ff5_realized_df.index)
    common_dates = common_dates.intersection(ff5_signal_df.index)
    common_dates = common_dates.sort_values()
    holding_month = pd.to_datetime(holding_month)

    if holding_month not in common_dates:
        raise ValueError(f"Holding month {holding_month.date()} is not in the common monthly index.")

    next_idx = common_dates.get_loc(holding_month)
    if next_idx <= config.beta_lookback:
        raise ValueError(
            f"Holding month {holding_month.date()} is too early. "
            f"Please choose a month after {common_dates[config.beta_lookback + 1].date()}."
        )

    signal_date = common_dates[next_idx - 1]
    cluster_train_dates = common_dates[: next_idx - 1]
    beta_train_dates = common_dates[next_idx - 1 - config.beta_lookback : next_idx - 1]

    cluster_train_ff5_df = ff5_signal_df.loc[cluster_train_dates]
    beta_train_ff5_df = ff5_realized_df.loc[beta_train_dates]
    beta_train_returns_df = returns_df.loc[beta_train_dates]

    cluster_model = build_regime_model(cluster_train_ff5_df, config)
    train_regime_series = cluster_model.factor_data[cluster_model.label_col].copy()
    train_regime_series.name = "Regime"

    regression_model = FF5Regression(price_df, beta_train_ff5_df)
    regression_model.returns = beta_train_returns_df
    regression_model.ff5_df = beta_train_ff5_df
    regression_model.ff5_clean = beta_train_ff5_df.copy()
    regression_result = regression_model.run_all_regression(ridge_alpha=config.ridge_alpha)
    if regression_result.empty:
        raise ValueError("Regression result is empty for the selected month.")

    beta_df = regression_result[
        ["ticker", "beta_mkt", "beta_smb", "beta_hml", "beta_rmw", "beta_cma"]
    ].copy().set_index("ticker")

    regime_factor_means = selector.compute_regime_factor_means(
        ff5_df=cluster_train_ff5_df,
        regime_series=train_regime_series,
    )

    current_factor = ff5_signal_df.loc[[signal_date], FACTOR_COLS]
    current_regime = cluster_model.predict_regime(
        current_factor,
        history_df=cluster_train_ff5_df,
    )
    top_stocks = selector.select_top_k_stocks(
        beta_df=beta_df,
        regime_factor_means=regime_factor_means,
        current_regime=current_regime,
        top_k=config.top_k,
    )
    if top_stocks.empty:
        raise ValueError("Top stock selection is empty for the selected month.")

    weights = optimizer.solve_min_risk_with_signal_tilt(
        stock_list=top_stocks.index.tolist(),
        raw_signal=top_stocks["score"],
        current_date=signal_date,
        lookback=config.cov_lookback,
        w_max=config.w_max,
        lambda_signal=config.lambda_signal,
        risk_model=config.risk_model,
        beta_df=beta_df,
        min_obs=config.min_cov_obs,
        idio_var_floor=config.idio_var_floor,
    )

    all_top_weights = top_stocks["score"].copy()
    all_top_weights[:] = 0.0
    all_top_weights.loc[weights.index] = weights
    all_top_weights.name = "weight"
    return all_top_weights.sort_values(ascending=True), current_regime, signal_date


def plot_month_snapshot(
    holding_month: pd.Timestamp,
    weights: pd.Series,
    regime: int,
    signal_date: pd.Timestamp,
    save_path: str | None = None,
):
    plt.figure(figsize=(13, 14))
    colors = ["#1f77b4" if weight > 0 else "#d9d9d9" for weight in weights.values]
    plt.barh(weights.index, weights.values, color=colors)
    plt.xlabel("Portfolio Weight")
    plt.ylabel("Top 50 Selected Stocks")
    plt.title(
        f"Signal Tilt Holdings for {holding_month.strftime('%Y-%m')}\n"
        f"Signal Date: {signal_date.strftime('%Y-%m-%d')} | Regime: {regime}"
    )
    plt.grid(axis="x", alpha=0.25)
    plt.tight_layout()

    if save_path is not None:
        plt.savefig(save_path, dpi=200, bbox_inches="tight")

    plt.show()


def parse_args():
    parser = argparse.ArgumentParser(
        description="Plot one month's regime and top-50 portfolio weights."
    )
    parser.add_argument(
        "--month",
        type=str,
        default=None,
        help="Holding month in YYYY-MM or YYYY-MM-DD format. Default uses the latest available month.",
    )
    parser.add_argument(
        "--save",
        action="store_true",
        help="Save the chart to the data folder as a PNG file.",
    )
    return parser.parse_args()


def choose_default_month(config: BacktestConfig) -> pd.Timestamp:
    _, returns_df, ff5_realized_df, ff5_signal_df, _ = load_dual_clock_inputs(
        realized_shift_months=REALIZED_SHIFT_MONTHS,
        signal_shift_months=SIGNAL_SHIFT_MONTHS,
    )
    common_dates = returns_df.index.intersection(ff5_realized_df.index)
    common_dates = common_dates.intersection(ff5_signal_df.index)
    common_dates = common_dates.sort_values()
    return common_dates[-1]


def main(month: str | None = None, save: bool | None = None):
    args = parse_args()
    config = BacktestConfig(
        cov_lookback=COV_LOOKBACK,
        risk_model=RISK_MODEL,
        min_cov_obs=MIN_COV_OBS,
        idio_var_floor=IDIO_VAR_FLOOR,
        cluster_method=CLUSTER_METHOD,
        hmm_n_iter=HMM_N_ITER,
        hmm_tol=HMM_TOL,
        hmm_stickiness=HMM_STICKINESS,
        hmm_covariance_floor=HMM_COVARIANCE_FLOOR,
        random_state=RANDOM_STATE,
    )

    month_value = month if month is not None else args.month
    save_value = SAVE_FIGURE if save is None else save
    if month is None:
        save_value = args.save or save_value

    if month_value is None:
        holding_month = choose_default_month(config)
    else:
        holding_month = pd.to_datetime(month_value)

    weights, regime, signal_date = build_month_snapshot(holding_month, config)

    save_path = None
    if save_value:
        filename = f"monthly_holdings_example_{holding_month.strftime('%Y_%m')}.png"
        save_path = os.path.join("data", filename)

    plot_month_snapshot(
        holding_month=holding_month,
        weights=weights,
        regime=regime,
        signal_date=signal_date,
        save_path=save_path,
    )

    print(f"Holding month: {holding_month.strftime('%Y-%m-%d')}")
    print(f"Signal date:   {signal_date.strftime('%Y-%m-%d')}")
    print(f"Regime:        {regime}")
    if save_path is not None:
        print(f"Saved plot to: {save_path}")


if __name__ == "__main__":
    main(DEFAULT_MONTH, SAVE_FIGURE)