import numpy as np


def softmax(predictions):
    """
    Computes probabilities from scores

    Arguments:
      predictions, np array, shape is either (N) or (batch_size, N) -
        classifier output

    Returns:
      probs, np array of the same shape as predictions - 
        probability for every class, 0..1
    """

    preds = np.array(predictions.copy())

    if preds.ndim > 1:
        preds -= np.max(preds, axis=1).reshape(-1, 1)
        proba = np.exp(preds) / np.sum(np.exp(preds), axis=1).reshape(-1, 1)
    else:
        preds -= np.max(preds)
        proba = np.exp(preds) / np.sum(np.exp(preds))

    
    return proba


def cross_entropy_loss(probs, target_index):
    """
    Computes cross-entropy loss

    Arguments:
      probs: np array, shape is either (N) or (batch_size, N) -
        probabilities for every class
      target_index: np array of int, shape is (1) or (batch_size) -
        index of the true class for given sample(s)

    Returns:
      loss: single value
    """

    if probs.ndim > 1:
        true_probs = probs[np.arange(len(target_index)), target_index.reshape(-1)]
    else:
        true_probs = probs[target_index]

    loss = np.sum(-np.log(true_probs))
    return loss


def softmax_with_cross_entropy(predictions, target_index):
    """
    Computes softmax and cross-entropy loss for model predictions,
    including the gradient

    Arguments:
      predictions: np array, shape is either (N) or (batch_size, N) -
        classifier output
      target_index: np array of int, shape is (1) or (batch_size) -
        index of the true class for given sample(s)

    Returns:
      loss: single value - cross-entropy loss
      dprediction: np array same shape as predictions - gradient of predictions by loss value
    """

    # TODO implement softmax with cross-entropy

    if predictions.ndim > 1:
        mask = np.zeros(predictions.shape)
        mask[np.arange(len(target_index)), target_index.reshape(-1)] = 1
    else:
        mask = np.zeros(predictions.shape)
        mask[target_index] = 1

    probs = softmax(predictions.copy())
    loss = cross_entropy_loss(probs, target_index)
    dpredictions = probs - mask

    return loss, dpredictions


def l2_regularization(W, reg_strength):
    """
    Computes L2 regularization loss on weights and its gradient

    Arguments:
      W: np array, weights
      reg_strength: float value

    Returns:
      loss: single value - l2 regularization loss
      gradient: np.array same shape as W - gradient of weight by l2 loss
    """

    w = W.copy()
    loss = np.sum(w * w) * reg_strength
    grad = w * 2 * reg_strength

    return loss, grad
    

def linear_softmax(X, W, target_index):
    """
    Performs linear classification and returns loss and gradient over W

    Arguments:
      X, np array, shape (num_batch, num_features) - batch of images
      W, np array, shape (num_features, classes) - weights
      target_index, np array, shape (num_batch) - index of target classes

    Returns:
      loss, single value - cross-entropy loss
      gradient, np.array same shape as W - gradient of weight by loss
    """

    predictions = np.dot(X, W)
    loss, dpred = softmax_with_cross_entropy(predictions, target_index)
    dW = X.T.dot(dpred)
    return loss, dW


class LinearSoftmaxClassifier():
    def __init__(self):
        self.W = None

    def fit(self, X, y, batch_size=100, learning_rate=1e-7, reg=1e-5, epochs=1):
        """
        Trains linear classifier
        
        Arguments:
          X, np array (num_samples, num_features) - training data
          y, np array of int (num_samples) - labels
          batch_size, int - batch size to use
          learning_rate, float - learning rate for gradient descent
          reg, float - L2 regularization strength
          epochs, int - number of epochs
        """

        num_train = X.shape[0]
        num_features = X.shape[1]
        num_classes = np.max(y)+1
        if self.W is None:
            self.W = 0.001 * np.random.randn(num_features, num_classes)

        loss_history = []
        for epoch in range(epochs):
            loss = 0
            shuffled_indices = np.arange(num_train)
            np.random.shuffle(shuffled_indices)
            sections = np.arange(batch_size, num_train, batch_size)
            batches_indices = np.array_split(shuffled_indices, sections)

            for batch in batches_indices:
                loss_softmax, dw_softmax = linear_softmax(X[batch], self.W, y[batch])
                loss_reg, dw_reg = l2_regularization(self.W, reg)
                loss += (loss_softmax + loss_reg)
                self.W -= learning_rate * (dw_softmax + dw_reg)
            loss_history.append(loss)
            # end
            if epoch % 100 == 0:
                print("Epoch %i, loss: %f" % (epoch, loss))

        return loss_history

    def predict(self, X):
        """
        Produces classifier predictions on the set
       
        Arguments:
          X, np array (test_samples, num_features)

        Returns:
          y_pred, np.array of int (test_samples)
        """

        y_pred = np.zeros(X.shape[0], dtype=np.int)

        # TODO Implement class prediction
        # Your final implementation shouldn't have any loops
        pred = np.argmax(softmax(X.dot(self.W)), axis=1)


        return pred
