"""Support Vector Machine (SVM) model."""

import numpy as np


class SVM:
    def __init__(self, n_class: int, lr: float, epochs: int, reg_const: float):
        """Initialize a new classifier.

        Parameters:
            n_class: the number of classes
            lr: the learning rate
            epochs: the number of epochs to train for
            reg_const: the regularization constant
        """
        self.w = None  # TODO: change this
        self.alpha = lr
        self.epochs = epochs
        self.reg_const = reg_const
        self.n_class = n_class

    def calc_gradient(self, X_train: np.ndarray, y_train: np.ndarray) -> np.ndarray:
        """Calculate gradient of the svm hinge loss.

        Inputs have dimension D, there are C classes, and we operate on
        mini-batches of N examples.

        Parameters:
            X_train: a numpy array of shape (N, D) containing a mini-batch
                of data
            y_train: a numpy array of shape (N,) containing training labels;
                y[i] = c means that X[i] has label c, where 0 <= c < C

        Returns:
            the gradient with respect to weights w; an array of the same shape
                as w
        """
        # TODO: implement me for sure
        num_train, dim = X_train.shape
        self.w = 0.0001 * np.random.randn(dim, self.n_class)
        
        dW = np.zeros(self.w.shape)
        loss = 0.0
        for i in range(num_train):
            scores = X_train[i].dot(self.w)
            correct_class_score = scores[y_train[i]]
            for j in range(self.n_class):
                if j == y_train[i]:
                    continue
                margin = scores[j] - correct_class_score + 1
                if margin > 0:
                    loss += margin
                    dW[:, y_train[i]] = dW[:, y_train[i]] - X_train[i]
                    dW[:,j] = dW[:,j] + X_train[i]
        loss /= num_train
        dW = dW / num_train
        loss += self.reg_const * np.sum(self.w * self.w)
        dW = dW + self.reg_const * 2 * self.w
        
        return dW

    def train(self, X_train: np.ndarray, y_train: np.ndarray):
        """Train the classifier.

        Hint: operate on mini-batches of data for SGD.

        Parameters:
            X_train: a numpy array of shape (N, D) containing training data;
                N examples with D dimensions
            y_train: a numpy array of shape (N,) containing training labels
        """
        # TODO: implement me
        batch_size=200
        num_train, dim = X_train.shape
        num_classes = self.n_class
        if self.w is None:
            self.w = 0.001 * np.random.randn(dim, num_classes)
        loss_history = []
        for it in range(self.epochs):
            X_batch = None
            y_batch = None
            batch_indices = np.random.choice(num_train, batch_size, replace=False)
            X_batch = X_train[batch_indices]
            y_batch = y_train[batch_indices]
            grad = self.calc_gradient(X_batch, y_batch)
            self.w = self.w - self.alpha * grad
        return

    def predict(self, X_test: np.ndarray) -> np.ndarray:
        """Use the trained weights to predict labels for test data points.

        Parameters:
            X_test: a numpy array of shape (N, D) containing testing data;
                N examples with D dimensions

        Returns:
            predicted labels for the data in X_test; a 1-dimensional array of
                length N, where each element is an integer giving the predicted
                class.
        """
        # TODO: implement me
        y_pred = np.zeros(X_test.shape[0])
        scores = X_test.dot(self.w)
        y_pred = scores.argmax(axis=1)
        return y_pred
