import numpy as np
from scipy.stats import multivariate_normal

class GMM:
    def __init__(self, n_components: int, n_iter: int) -> None:
        self.n_components = n_components
        self.n_iter = n_iter
        self.means = None
        self.covariances = None
        self.pi = None
        self.reg_covar = 1e-6

    def initialize_parameters(self, X: np.ndarray) -> None:
        np.random.seed(0)
        random_idx = np.random.permutation(X.shape[0])
        self.means = X[random_idx[:self.n_components]]
        self.covariances = [np.cov(X.T) + self.reg_covar * np.eye(X.shape[1]) for _ in range(self.n_components)]
        self.pi = np.ones(self.n_components) / self.n_components

    def e_step(self, X: np.ndarray) -> np.ndarray:
        responsibilities = np.zeros((X.shape[0], self.n_components))
        for i in range(self.n_components):
            rv = multivariate_normal(self.means[i], self.covariances[i])
            responsibilities[:, i] = self.pi[i] * rv.pdf(X)
        responsibilities /= responsibilities.sum(axis=1, keepdims=True)
        return responsibilities

    def m_step(self, X: np.ndarray, responsibilities: np.ndarray) -> None:
        Nk = responsibilities.sum(axis=0)
        self.means = np.dot(responsibilities.T, X) / Nk[:, np.newaxis]
        for i in range(self.n_components):
            x_minus_mean = X - self.means[i]
            self.covariances[i] = np.dot(responsibilities[:, i] * x_minus_mean.T, x_minus_mean) / Nk[i]
            self.covariances[i] += self.reg_covar * np.eye(X.shape[1])
            self.pi[i] = Nk[i] / X.shape[0]

    def fit(self, X: np.ndarray) -> None:
        self.initialize_parameters(X)
        for _ in range(self.n_iter):
            responsibilities = self.e_step(X)
            self.m_step(X, responsibilities)

    def predict(self, X: np.ndarray) -> np.ndarray:
        responsibilities = self.e_step(X)
        return np.argmax(responsibilities, axis=1)