fix the bug in code below
"""
File: semisupervised.py
Project: autograder_test_files
File Created: September 2020
Author: Shalini Chaudhuri (<PRESIDIO_ANONYMIZED_EMAIL_ADDRESS>)
Updated: September 2022, Arjun Agarwal
"""

import numpy as np
from sklearn.metrics import accuracy_score
from sklearn.naive_bayes import GaussianNB
from tqdm import tqdm

SIGMA_CONST = 1e-06
LOG_CONST = 1e-32


def complete_(data):
    """
    Args:
        data: N x (D+1) numpy array where the last column is the labels
    Return:
        labeled_complete: n x (D+1) array (n <= N) where values contain both complete features and labels
    """
    return data[~np.isnan(data).any(axis=1)]


def incomplete_(data):
    """
    Args:
        data: N x (D+1) numpy array where the last column is the labels
    Return:
        labeled_incomplete: n x (D+1) array (n <= N) where values contain incomplete features but complete labels
    """
    return data[np.isnan(data[:, :-1]).any(axis=1)]


def unlabeled_(data):
    """
    Args:
        data: N x (D+1) numpy array where the last column is the labels
    Return:
        unlabeled_complete: n x (D+1) array (n <= N) where values contain complete features but incomplete labels
    """
    return data[np.isnan(data[:, -1]) & ~np.isnan(data[:, :-1]).any(axis=1)]


class CleanData(object):

    def __init__(self):
        pass

    def pairwise_dist(self, x, y):
        """
        Args:
            x: N x D numpy array
            y: M x D numpy array
        Return:
            dist: N x M array, where dist[i, j] is the euclidean distance between
            x[i, :] and y[j, :]
        """
        x_squared = np.sum(x**2, axis=1, keepdims=True)
        y_squared = np.sum(y**2, axis=1)
        xy = np.dot(x, y.T)
        dist = np.sqrt(x_squared + y_squared - 2 * xy)
        return dist

    def __call__(self, incomplete_points, complete_points, K, **kwargs):
        """
        Function to clean or "fill in" NaN values in incomplete data points based on
        the average value for that feature for the K-nearest neighbors in the complete data points.

        Args:
            incomplete_points: N_incomplete x (D+1) numpy array, the incomplete labeled observations
            complete_points:   N_complete   x (D+1) numpy array, the complete labeled observations
            K: integer, corresponding to the number of nearest neighbors you want to base your calculation on
            kwargs: any other args you want
        Return:
            clean_points: (N_complete + N_incomplete) x (D+1) numpy array, containing both the complete points and recently filled points

        Notes:
            (1) The first D columns are features, and the last column is the class label
            (2) There may be more than just 2 class labels in the data (e.g. labels could be 0,1,2 or 0,1,2,...,M)
            (3) There will be at most 1 missing feature value in each incomplete data point (e.g. no points will have more than one NaN value)
            (4) You want to find the k-nearest neighbors, from the complete dataset, with the same class labels;
            (5) There may be missing values in any of the features. It might be more convenient to address each feature at a time.
            (6) Do NOT use a for-loop over N_incomplete; you MAY use a for-loop over the M labels and the D features (e.g. omit one feature at a time)
            (7) You do not need to order the rows of the return array clean_points in any specific manner
        """
        filled_incomplete_points = []
        complete_features = complete_points[:, :-1]
        complete_labels = complete_points[:, -1]

        for point in incomplete_points:
            nan_index = np.where(np.isnan(point))[0][0]
            incomplete_features = point[:-1]
            incomplete_label = point[-1]
            same_label_mask = complete_labels == incomplete_label
            same_label_complete_features = complete_features[same_label_mask]

            # use dict func
            distances = self.pairwise_dist(
                np.delete(incomplete_features.reshape(1, -1), nan_index, axis=1),
                np.delete(same_label_complete_features, nan_index, axis=1),
            )
            nearest_neighbor_indices = np.argsort(distances.flatten())[:K]
            avg_value = np.nanmean(
                same_label_complete_features[nearest_neighbor_indices, nan_index]
            )
            filled_point = point.copy()
            filled_point[nan_index] = avg_value
            filled_incomplete_points.append(filled_point)
        clean_points = np.vstack((complete_points, np.array(filled_incomplete_points)))

        return clean_points


def median_clean_data(data):
    """
    Args:
        data: N x (D+1) numpy array where only last column is guaranteed non-NaN values and is the labels
    Return:
        median_clean: N x (D+1) numpy array where each NaN value in data has been replaced by the median feature value
    Notes:
        (1) When taking the median of any feature, do not count the NaN value
        (2) Return all values to max one decimal point
        (3) The labels column will never have NaN values
    """
    median_clean = data
    for feature_idx in range(data.shape[1] - 1):
        feature_median = np.nanmedian(data[:, feature_idx])
        median_clean[:, feature_idx] = np.where(
            np.isnan(median_clean[:, feature_idx]),
            feature_median,
            median_clean[:, feature_idx],
        )

    return np.round(median_clean)


class SemiSupervised(object):

    def __init__(self):
        pass

    def softmax(self, logit):
        """
        Args:
            logit: N x D numpy array
        Return:
            prob: N x D numpy array where softmax has been applied row-wise to input logit
        """
        logit_sub = logit - logit.max(axis=1, keepdims=True)
        exp_logit = np.exp(logit_sub)
        sum_exp_logit = np.sum(exp_logit, axis=1, keepdims=True)
        return exp_logit / sum_exp_logit

    def logsumexp(self, logit):
        """
        Args:
            logit: N x D numpy array
        Return:
            s: N x 1 array where s[i,0] = logsumexp(logit[i,:])
        """
        logit_max = np.max(logit, axis=1, keepdims=True)
        return logit_max + np.log(
            np.sum(np.exp(logit - logit_max), axis=1, keepdims=True)
        )

    def normalPDF(self, logit, mu_i, sigma_i):
        """
        Args:
            logit: N x D numpy array
            mu_i: 1xD numpy array, the center for the ith gaussian.
            sigma_i: 1xDxD numpy array, the covariance matrix of the ith gaussian.
        Return:
            pdf: 1xN numpy array, the probability distribution of N data for the ith gaussian

        Hint:
            np.diagonal() should be handy.
        """
        D = logit.shape[1]
        det_sigma_i = np.linalg.det(sigma_i)
        inv_sigma_i = np.linalg.inv(sigma_i)
        exp_term = np.exp(
            -0.5 * np.sum((logit - mu_i) * np.dot(logit - mu_i, inv_sigma_i), axis=1)
        )
        pdf = (1 / np.sqrt((2 * np.pi) ** D * det_sigma_i)) * exp_term
        return pdf.reshape(1, -1)

    def _init_components(self, points, K, **kwargs):
        """
        Args:
            points: Nx(D+1) numpy array, the observations
            K: number of components
            kwargs: any other args you want
        Return:
            pi: numpy array of length K; contains the prior probabilities of each class k
            mu: KxD numpy array, the center for each gaussian.
            sigma: KxDxD numpy array, the diagonal standard deviation of each gaussian.
        Hint:
            As explained in the algorithm, you need to calculate the values of mu, sigma and pi based on the labelled dataset
        """
        labeled_complete = complete_(points)

        X_labeled = labeled_complete[:, :-1]
        y_labeled = labeled_complete[:, -1].astype(int)

        n_features = X_labeled.shape[1]
        mu = np.zeros((K, n_features))
        sigma = np.zeros((K, n_features, n_features))
        pi = np.zeros(K)

        for k in range(K):
            class_points = X_labeled[y_labeled == k]
            class_size = class_points.shape[0]

            if class_size > 0:
                mu[k] = np.mean(class_points, axis=0)
                sigma[k] = (
                    np.cov(class_points, rowvar=False)
                    + np.eye(n_features) * SIGMA_CONST
                )
                pi[k] = class_size / X_labeled.shape[0]

        return pi, mu, sigma

    def _ll_joint(self, points, pi, mu, sigma, **kwargs):
        """
        Args:
            points: NxD numpy array, the observations
            pi: np array of length K, the prior of each component
            mu: KxD numpy array, the center for each gaussian.
            sigma: KxDxD numpy array, the diagonal standard deviation of each gaussian.
        Return:
            ll(log-likelihood): NxK array, where ll(i, j) = log pi(j) + log NormalPDF(points_i | mu[j], sigma[j])
        """
        K = len(pi)
        ll = np.zeros((points.shape[0], K))
        for k in range(K):
            pdf = self.normalPDF(points, mu[k], sigma[k])
            ll[:, k] = np.log(pi[k]) + np.log(pdf).flatten()
        return ll

    def _E_step(self, points, pi, mu, sigma, **kwargs):
        """
        Args:
            points: NxD numpy array, the observations
            pi: np array of length K, the prior of each component
            mu: KxD numpy array, the center for each gaussian.
            sigma: KxDxD numpy array, the diagonal standard deviation of each gaussian.
        Return:
            gamma: NxK array, the posterior distribution (a.k.a, the soft cluster assignment) for each observation.

        Hint: You should be able to do this with just a few lines of code by using _ll_joint() and softmax() defined above.
        """
        ll_joint = self._ll_joint(points, pi, mu, sigma)
        gamma = self.softmax(ll_joint)
        return gamma

    def _M_step(self, points, gamma, **kwargs):
        """
        Args:
            points: NxD numpy array, the observations
            gamma: NxK array, the posterior distribution (a.k.a, the soft cluster assignment) for each observation.
        Return:
            pi: np array of length K, the prior of each component
            mu: KxD numpy array, the center for each gaussian.
            sigma: KxDxD numpy array, the diagonal standard deviation of each gaussian.

        Hint:  There are formulas in the slide.
        """
        N, K = gamma.shape
        pi = np.sum(gamma, axis=0)
        pi /= np.sum(pi)

        mu = np.zeros((K, points.shape[1]))
        for k in range(K):
            mu[k] = np.mean(points * gamma[:, k].reshape(-1, 1), axis=0)

        sigma = np.zeros((K, points.shape[1], points.shape[1]))
        for k in range(K):
            sigma[k] = np.cov(points, rowvar=False, aweights=gamma[:, k])

        return pi, mu, sigma

    def __call__(
        self, points, K, max_iters=100, abs_tol=1e-16, rel_tol=1e-16, **kwargs
    ):
        """
        Args:
            points: N x (D+1) numpy array, where
                - N is # points,
                - D is the number of features,
                - the last column is the point labels (when available) or NaN for unlabeled points
            K: integer, number of clusters
            max_iters: maximum number of iterations
            abs_tol: convergence criteria w.r.t absolute change of loss
            rel_tol: convergence criteria w.r.t relative change of loss
            kwargs: any additional arguments you want
        Return:
            pi, mu, sigma: (1xK np array, KxD numpy array, KxDxD numpy array)
        """
        N, D_plus_1 = points.shape
        D = D_plus_1 - 1
        pi, mu, sigma = self._init_components(points, K, **kwargs)
        labeled_points = complete_(points)
        unlabeled_points = unlabeled_(points)
        X_u = unlabeled_points[:, :-1]

        log_likelihood = -np.inf

        for i in range(max_iters):
            gamma = self._E_step(X_u, pi, mu, sigma)
            pi, mu, sigma = self._M_step(
                np.vstack((labeled_points[:, :-1], X_u)), gamma
            )
            ll_new = np.sum(np.log(np.sum(self._ll_joint(X_u, pi, mu, sigma), axis=1)))
            ll_change = ll_new - log_likelihood
            log_likelihood = ll_new

            if abs(ll_change) < abs_tol or abs(ll_change / log_likelihood) < rel_tol:
                break

        return pi, mu, sigma


class ComparePerformance(object):

    def __init__(self):
        pass

    @staticmethod
    def accuracy_semi_supervised(training_data, validation_data, K: int) -> float:
        """
        Train a classification model using your SemiSupervised object on the training_data.
        Classify the validation_data using the trained model
        Return the accuracy score of the model's predicted classification of the validation_data

        Args:
            training_data: N_t x (D+1) numpy array, where
                - N_t is the number of data points in the training set,
                - D is the number of features, and
                - the last column represents the labels (when available) or a flag that allows you to separate the unlabeled data.
            validation_data: N_v x(D+1) numpy array, where
                - N_v is the number of data points in the validation set,
                - D is the number of features, and
                - the last column are the labels
            K: integer, number of clusters for SemiSupervised object
        Return:
            accuracy: floating number

        Note: (1) validation_data will NOT include any unlabeled points
              (2) you may use sklearn accuracy_score: https://scikit-learn.org/stable/modules/generated/sklearn.metrics.accuracy_score.html
        """
        semisupervised = SemiSupervised()
        pi, mu, sigma = semisupervised(training_data, K)
        validation_features = validation_data[:, :-1]
        ll_joint = semisupervised._ll_joint(validation_features, pi, mu, sigma)
        predicted_labels = np.argmax(ll_joint, axis=1)
        true_labels = validation_data[:, -1].astype(int)
        accuracy = accuracy_score(true_labels, predicted_labels)
        return accuracy

    @staticmethod
    def accuracy_GNB(training_data, validation_data) -> float:
        """
        Train a Gaussion Naive Bayes classification model (sklearn implementation) on the training_data.
        Classify the validation_data using the trained model
        Return the accuracy score of the model's predicted classification of the validation_data

        Args:
            training_data: N_t x (D+1) numpy array, where
                - N is the number of data points in the training set,
                - D is the number of features, and
                - the last column represents the labels
            validation_data: N_v x (D+1) numpy array, where
                - N_v is the number of data points in the validation set,
                - D is the number of features, and
                - the last column are the labels
        Return:
            accuracy: floating number

        Note: (1) both training_data and validation_data will NOT include any unlabeled points
              (2) use sklearn implementation of Gaussion Naive Bayes: https://scikit-learn.org/stable/modules/generated/sklearn.naive_bayes.GaussianNB.html
        """
        gnb = GaussianNB()
        training_features = training_data[:, :-1]
        training_labels = training_data[:, -1].astype(int)
        gnb.fit(training_features, training_labels)
        validation_features = validation_data[:, :-1]
        predicted_labels = gnb.predict(validation_features)
        true_labels = validation_data[:, -1].astype(int)
        accuracy = accuracy_score(true_labels, predicted_labels)
        return accuracy

{
	"name": "ValueError",
	"message": "operands could not be broadcast together with shapes (19020,10) (232,1) ",
	"stack": "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m\n\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)\nCell \u001b[0;32mIn[24], line 52\u001b[0m\n\u001b[1;32m     47\u001b[0m accuracy_median_cleaned_data \u001b[38;5;241m=\u001b[39m ComparePerformance\u001b[38;5;241m.\u001b[39maccuracy_GNB(\n\u001b[1;32m     48\u001b[0m     median_cleaned_data, validation\n\u001b[1;32m     49\u001b[0m )\n\u001b[1;32m     50\u001b[0m \u001b[38;5;66;03m# =========================================================================\u001b[39;00m\n\u001b[1;32m     51\u001b[0m \u001b[38;5;66;03m# SEMI SUPERVISED GNB WITH ALL DATA (your implementation)\u001b[39;00m\n\u001b[0;32m---> 52\u001b[0m accuracy_semi_supervised \u001b[38;5;241m=\u001b[39m ComparePerformance\u001b[38;5;241m.\u001b[39maccuracy_semi_supervised(\n\u001b[1;32m     53\u001b[0m     cleaned_and_unlabeled, validation, \u001b[38;5;241m2\u001b[39m\n\u001b[1;32m     54\u001b[0m )\n\u001b[1;32m     55\u001b[0m \u001b[38;5;66;03m# ==========================================================================\u001b[39;00m\n\u001b[1;32m     56\u001b[0m \u001b[38;5;66;03m# COMPARISON\u001b[39;00m\n\u001b[1;32m     57\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\"\"\u001b[39m\u001b[38;5;124m===COMPARISON===\u001b[39m\u001b[38;5;124m\"\"\"\u001b[39m)\n\nFile \u001b[0;32m~/Desktop/4641/HW2/student_files/semisupervised.py:357\u001b[0m, in \u001b[0;36maccuracy_semi_supervised\u001b[0;34m(training_data, validation_data, K)\u001b[0m\n\u001b[1;32m    351\u001b[0m \u001b[38;5;129m@staticmethod\u001b[39m\n\u001b[1;32m    352\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21maccuracy_GNB\u001b[39m(training_data, validation_data) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m \u001b[38;5;28mfloat\u001b[39m:\n\u001b[1;32m    353\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    354\u001b[0m \u001b[38;5;124;03m    Train a Gaussion Naive Bayes classification model (sklearn implementation) on the training_data.\u001b[39;00m\n\u001b[1;32m    355\u001b[0m \u001b[38;5;124;03m    Classify the validation_data using the trained model\u001b[39;00m\n\u001b[1;32m    356\u001b[0m \u001b[38;5;124;03m    Return the accuracy score of the model's predicted classification of the validation_data\u001b[39;00m\n\u001b[0;32m--> 357\u001b[0m \n\u001b[1;32m    358\u001b[0m \u001b[38;5;124;03m    Args:\u001b[39;00m\n\u001b[1;32m    359\u001b[0m \u001b[38;5;124;03m        training_data: N_t x (D+1) numpy array, where\u001b[39;00m\n\u001b[1;32m    360\u001b[0m \u001b[38;5;124;03m            - N is the number of data points in the training set,\u001b[39;00m\n\u001b[1;32m    361\u001b[0m \u001b[38;5;124;03m            - D is the number of features, and\u001b[39;00m\n\u001b[1;32m    362\u001b[0m \u001b[38;5;124;03m            - the last column represents the labels\u001b[39;00m\n\u001b[1;32m    363\u001b[0m \u001b[38;5;124;03m        validation_data: N_v x (D+1) numpy array, where\u001b[39;00m\n\u001b[1;32m    364\u001b[0m \u001b[38;5;124;03m            - N_v is the number of data points in the validation set,\u001b[39;00m\n\u001b[1;32m    365\u001b[0m \u001b[38;5;124;03m            - D is the number of features, and\u001b[39;00m\n\u001b[1;32m    366\u001b[0m \u001b[38;5;124;03m            - the last column are the labels\u001b[39;00m\n\u001b[1;32m    367\u001b[0m \u001b[38;5;124;03m    Return:\u001b[39;00m\n\u001b[1;32m    368\u001b[0m \u001b[38;5;124;03m        accuracy: floating number\u001b[39;00m\n\u001b[1;32m    369\u001b[0m \n\u001b[1;32m    370\u001b[0m \u001b[38;5;124;03m    Note: (1) both training_data and validation_data will NOT include any unlabeled points\u001b[39;00m\n\u001b[1;32m    371\u001b[0m \u001b[38;5;124;03m          (2) use sklearn implementation of Gaussion Naive Bayes: https://scikit-learn.org/stable/modules/generated/sklearn.naive_bayes.GaussianNB.html\u001b[39;00m\n\u001b[1;32m    372\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[1;32m    373\u001b[0m     gnb \u001b[38;5;241m=\u001b[39m GaussianNB()\n\u001b[1;32m    374\u001b[0m     training_features \u001b[38;5;241m=\u001b[39m training_data[:, :\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m]\n\nFile \u001b[0;32m~/Desktop/4641/HW2/student_files/semisupervised.py:315\u001b[0m, in \u001b[0;36m__call__\u001b[0;34m(self, points, K, max_iters, abs_tol, rel_tol, **kwargs)\u001b[0m\n\u001b[1;32m      0\u001b[0m <Error retrieving source code with stack_data see ipython/ipython#13598>\n\nFile \u001b[0;32m~/Desktop/4641/HW2/student_files/semisupervised.py:279\u001b[0m, in \u001b[0;36m_M_step\u001b[0;34m(self, points, gamma, **kwargs)\u001b[0m\n\u001b[1;32m    273\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__call__\u001b[39m(\n\u001b[1;32m    274\u001b[0m     \u001b[38;5;28mself\u001b[39m, points, K, max_iters\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m100\u001b[39m, abs_tol\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1e-16\u001b[39m, rel_tol\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1e-16\u001b[39m, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs\n\u001b[1;32m    275\u001b[0m ):\n\u001b[1;32m    276\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    277\u001b[0m \u001b[38;5;124;03m    Args:\u001b[39;00m\n\u001b[1;32m    278\u001b[0m \u001b[38;5;124;03m        points: N x (D+1) numpy array, where\u001b[39;00m\n\u001b[0;32m--> 279\u001b[0m \u001b[38;5;124;03m            - N is # points,\u001b[39;00m\n\u001b[1;32m    280\u001b[0m \u001b[38;5;124;03m            - D is the number of features,\u001b[39;00m\n\u001b[1;32m    281\u001b[0m \u001b[38;5;124;03m            - the last column is the point labels (when available) or NaN for unlabeled points\u001b[39;00m\n\u001b[1;32m    282\u001b[0m \u001b[38;5;124;03m        K: integer, number of clusters\u001b[39;00m\n\u001b[1;32m    283\u001b[0m \u001b[38;5;124;03m        max_iters: maximum number of iterations\u001b[39;00m\n\u001b[1;32m    284\u001b[0m \u001b[38;5;124;03m        abs_tol: convergence criteria w.r.t absolute change of loss\u001b[39;00m\n\u001b[1;32m    285\u001b[0m \u001b[38;5;124;03m        rel_tol: convergence criteria w.r.t relative change of loss\u001b[39;00m\n\u001b[1;32m    286\u001b[0m \u001b[38;5;124;03m        kwargs: any additional arguments you want\u001b[39;00m\n\u001b[1;32m    287\u001b[0m \u001b[38;5;124;03m    Return:\u001b[39;00m\n\u001b[1;32m    288\u001b[0m \u001b[38;5;124;03m        pi, mu, sigma: (1xK np array, KxD numpy array, KxDxD numpy array)\u001b[39;00m\n\u001b[1;32m    289\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[1;32m    290\u001b[0m     N, D_plus_1 \u001b[38;5;241m=\u001b[39m points\u001b[38;5;241m.\u001b[39mshape\n\u001b[1;32m    291\u001b[0m     D \u001b[38;5;241m=\u001b[39m D_plus_1 \u001b[38;5;241m-\u001b[39m \u001b[38;5;241m1\u001b[39m\n\n\u001b[0;31mValueError\u001b[0m: operands could not be broadcast together with shapes (19020,10) (232,1) "
}