Two-state Markov-switching model: Hamilton filter and Kim smoother

"""M13 · Two-state Markov-switching model: Hamilton filter and Kim smoother. Read notes §13.5.

Model:
    y_t | s_t = j  ~  N(mu_j, sigma_j^2),   j in {0, 1}
    P[i, j] = Pr(s_t = j | s_{t-1} = i)      (rows sum to 1)

Convention: arrays are indexed t = 0..n-1.
    filtered[t]  = Pr(s_t | y_0..y_t)
    predicted[t] = Pr(s_t | y_0..y_{t-1})   (predicted[0] = initial distribution)
    smoothed[t]  = Pr(s_t | y_0..y_{n-1})
Each is an (n, 2) array whose rows sum to 1.

Run:  cd exercises && .venv/bin/python -m pytest modules/M13-forecasting/tests/test_m13_hamilton.py
"""

from __future__ import annotations

import numpy as np


def ergodic_probs(P: np.ndarray) -> np.ndarray:
    """Stationary distribution pi (pi' P = pi', sum pi = 1). For 2 states pi_0 = (1 - p11) / (2 - p00 - p11)."""
    raise NotImplementedError


def simulate_ms(n: int, mus, sigmas, P, seed: int = 0) -> tuple[np.ndarray, np.ndarray]:
    """Simulate (y, s) of length n, s_0 drawn from the ergodic distribution. (Given; uses ergodic_probs.)"""
    rng = np.random.default_rng(seed)
    P = np.asarray(P, dtype=float)
    mus = np.asarray(mus, dtype=float)
    sigmas = np.asarray(sigmas, dtype=float)
    s = np.empty(n, dtype=int)
    s[0] = rng.choice(2, p=ergodic_probs(P))
    u = rng.random(n)
    for t in range(1, n):
        s[t] = int(u[t] > P[s[t - 1], 0])
    y = mus[s] + sigmas[s] * rng.standard_normal(n)
    return y, s


def hamilton_filter(y, mus, sigmas, P, pi0=None) -> tuple[np.ndarray, np.ndarray, float]:
    """Hamilton (1989) filter, notes eqs. 13.23-13.25. Return (filtered, predicted, loglik).

    predicted[0] = pi0 (default: ergodic_probs(P));
    filtered[t]  = predicted[t] * f(y_t | s_t) / sum_j predicted[t, j] f(y_t | s_t = j);
    predicted[t+1] = P' filtered[t];
    loglik = sum_t ln sum_j predicted[t, j] f(y_t | s_t = j), f the normal density.
    filtered[t] must depend on y_0..y_t only (a test rewrites the future and checks).
    """
    raise NotImplementedError


def kim_smoother(filtered: np.ndarray, predicted: np.ndarray, P: np.ndarray) -> np.ndarray:
    """Kim (1994) backward recursion, notes eq. 13.27:

        smoothed[n-1] = filtered[n-1]
        smoothed[t, i] = filtered[t, i] * sum_j P[i, j] smoothed[t+1, j] / predicted[t+1, j]
    """
    raise NotImplementedError


def fit_ms2(y) -> dict:
    """MLE of (mu0, mu1, sigma0, sigma1, p00, p11) by maximising hamilton_filter's loglik.

    Return dict(mus, sigmas, P, loglik) with states ordered so that sigmas[0] <= sigmas[1].
    Reparametrise: sigma = exp(.), p = logistic(.). Start the two vols apart
    (e.g. 0.6 and 1.5 times the sample std) to break the label symmetry.
    """
    raise NotImplementedError