import numpy as np
from numpy.typing import NDArray


def move_towards_mean(mu: NDArray, mu_mean: NDArray, t: float) -> NDArray:
    '''
    Returns a new point that is moved towards the mean according to the rule:
    point(t) = point + t * (mean - point), 0 ≤ t ≤ 1
    '''
    new_mu = mu + t * (mu_mean - mu)
    return new_mu


def split_points_by_clusters(n_points: int, n_clusters: int, probability: float) -> list[int]:
    '''
    Get the number of points to generate and the number of clusters as input
    returns the list with number of point in each cluster
    '''
    n_list = []
    if n_clusters == 2:
        n_1 = int(probability * n_points)
        n_list = [n_1, n_points - n_1]
    elif n_clusters == 3:
        n_1 = int(probability * n_points)
        n_2 = int(probability * n_points)
        n_3 = n_points - n_1 - n_2
        n_list = [n_1, n_2, n_3]
    return n_list


# pylint: disable= too-many-locals
def generate_mix_distribution(
    probability: float,
    mu_list: list[NDArray],
    cov_matrices: list[NDArray],
    n_samples: int,
    t: float
    ) -> tuple[NDArray, NDArray, NDArray]:
    '''
    Returns points from gaussian distributions.
    generated by distribution probability, covariance matrix and means vector.
    '''
    n_clusters = len(mu_list)

    mu_mean = np.mean(mu_list, axis=0)
    modified_mu_list = [move_towards_mean(mu, mu_mean, t) for mu in mu_list]

    splitted_points = split_points_by_clusters(n_samples, n_clusters, probability)

    distributions = []
    for n, mu, covariance_matrix in zip(splitted_points, modified_mu_list, cov_matrices):
        distribution = np.random.multivariate_normal(
            np.squeeze(mu, axis=0), covariance_matrix, n)
        distributions.append(distribution)

    samples = np.concatenate(distributions, axis=0)

    labels = np.array([])
    for label_id, n in enumerate(splitted_points):
        samples_labels = np.full(n, label_id)
        labels = np.concatenate((labels, samples_labels))

    # Shuffle the samples and labels
    permutation = np.random.permutation(n_samples)
    samples = samples[permutation]
    labels = labels[permutation]

    return samples, labels, np.mean(modified_mu_list, axis=1)
