if __name__ == "__main__":
    # Tiny HMM: 2 hidden states (Noun=0, Verb=1), 3 observation
    # symbols ("dog"=0, "runs"=1, "fast"=2)
    pi = np.array([0.6, 0.4])
    A = np.array([[0.3, 0.7], [0.6, 0.4]])
    B = np.array([
        [0.5, 0.1, 0.1],   # Noun emits "dog" most often
        [0.1, 0.6, 0.4],   # Verb emits "runs"/"fast" most often
    ])
    obs = [0, 1, 2]  # "dog runs fast"

    alpha, prob = hmm_forward(pi, A, B, obs)
    print("P(obs sequence):", prob)
    path, best_score = hmm_viterbi(pi, A, B, obs)
    tags = ["Noun" if s == 0 else "Verb" for s in path]
    print("Most likely tags:", tags)

    rng = np.random.default_rng(0)
    true_means = np.array([[0, 0], [6, 6]])
    X = np.vstack([
        rng.normal(loc=m, scale=0.8, size=(100, 2)) for m in true_means
    ])
    means, covs, weights, resp = gmm_em(X, k=2)
    print("\nRecovered GMM means:")
    print(np.round(means, 2))
