import numpy as np


def hmm_forward(pi, A, B, obs):
    T, N = len(obs), len(pi)
    alpha = np.zeros((T, N))
    alpha[0] = pi * B[:, obs[0]]
    for t in range(1, T):
        alpha[t] = (alpha[t - 1] @ A) * B[:, obs[t]]
    return alpha, alpha[-1].sum()


def hmm_viterbi(pi, A, B, obs):
    T, N = len(obs), len(pi)
    delta = np.zeros((T, N))
    psi = np.zeros((T, N), dtype=int)
    delta[0] = pi * B[:, obs[0]]
    for t in range(1, T):
        for j in range(N):
            trans = delta[t - 1] * A[:, j]
            psi[t, j] = np.argmax(trans)
            delta[t, j] = trans[psi[t, j]] * B[j, obs[t]]
    path = [int(np.argmax(delta[-1]))]
    for t in range(T - 1, 0, -1):
        path.append(int(psi[t, path[-1]]))
    path.reverse()
    return path, delta[-1].max()


def gmm_em(X, k, n_iters=100, seed=0):
    rng = np.random.default_rng(seed)
    n, d = X.shape
    idx = rng.choice(n, k, replace=False)
    means = X[idx].copy()
    covs = np.array([np.eye(d) for _ in range(k)])
    weights = np.ones(k) / k

    def gaussian_pdf(x, mean, cov):
        diff = x - mean
        inv = np.linalg.inv(cov)
        exponent = -0.5 * diff @ inv @ diff
        norm = 1.0 / np.sqrt(
            (2 * np.pi) ** d * np.linalg.det(cov)
        )
        return norm * np.exp(exponent)

    for _ in range(n_iters):
        # E-step
        resp = np.zeros((n, k))
        for i in range(n):
            for j in range(k):
                resp[i, j] = weights[j] * gaussian_pdf(
                    X[i], means[j], covs[j]
                )
            resp[i] /= resp[i].sum()

        # M-step
        Nk = resp.sum(axis=0)
        for j in range(k):
            means[j] = (resp[:, j] @ X) / Nk[j]
            diff = X - means[j]
            covs[j] = (resp[:, j] * diff.T) @ diff / Nk[j]
            covs[j] += 1e-6 * np.eye(d)  # numerical stability
        weights = Nk / n

    return means, covs, weights, resp
