import numpy as np
from typing import Tuple
from sklearn import datasets
import matplotlib.pyplot as plt


def get_dataset(num_centers: int) -> Tuple[np.ndarray, np.ndarray]:
    np.random.seed(0)
    X = []
    Y = []
    for cpt in range(num_centers):
        X += (
            np.random.multivariate_normal(
                mean=(np.random.normal(scale=2, size=(2,))),
                cov=np.diag([0.1, 0.1]),
                size=100,
            )
        ).tolist()
        Y += (cpt * np.ones(shape=(100,))).tolist()
    X = np.array(X)
    Y = np.array(Y).astype(np.int32)
    b = np.zeros((Y.size, Y.max() + 1))
    b[np.arange(Y.size), Y] = 1
    Y = b
    return X, Y


def show_dataset(X: np.ndarray, Y: np.ndarray, step: int):
    for cpt in range(Y.shape[-1]):
        plt.scatter(
            X[np.where(np.argmax(Y, axis=-1) == cpt), 0],
            X[np.where(np.argmax(Y, axis=-1) == cpt), 1],
        )
    if step is not None:
        plt.savefig(f"figs/{step}.png")
    else:
        plt.savefig(f"dataset.png")
    plt.close()