import numpy as np
import time

def loss_i(x_i, y, W):
    delta = 1
    scores_i = np.dot(W, x_i)

    margin = 0
    for j in range(scores_i.shape[0]):
        if j == y:
            continue
        else:
            margin += max(0, scores_i[j] - scores_i[y] + delta)

    return margin

def loss_i_v(x_i, y, W):
    delta = 1
    scores_i = np.dot(W, x_i.T)
    margins_i = np.maximum(0, scores_i - scores_i[y] + delta)
    margins_i[y] = 0
    margin_i = np.sum(margins_i, axis=0)
    return margin_i

def loss(X, Y, W):
    delta = 1

    scores = np.dot(W, X.T)

    rows = np.arange(Y.shape[0])
    margins = np.maximum(0, scores - scores[Y, rows] + delta)

    # Find indices that need to be set to zero.
    # Case when j == y in loss_i_v.
    subtract_Y_indices = np.zeros(shape=margins.shape, dtype=bool)
    subtract_Y_indices[Y, rows] = True
    # Set proper indices to 0.
    margins[subtract_Y_indices] = 0

    return np.sum(margins) / X.shape[0] + np.sum(W ** 2)

def loss_iterative(X, Y, W):
    loss = 0

    for i in range(X.shape[0]):
        loss += loss_i(X[i, :], Y[i], W)

    return loss / X.shape[0] + np.sum(W ** 2)

"""
Example:
"""
X = np.array([
    [0, 1, 2, 0, 1, 1, 1],
    [3, 1, 1, 0, 2, 2, 2],
    # [1, 2, 3, 0, 3, 3, 3],
])

Y = np.array([1, 2])

labels = np.array([0, 1, 2, 3, 4])

W = np.zeros((labels.shape[0], X.shape[1]))

t = time.time()
loss1 = loss_iterative(X, Y, W)
e1 = time.time() - t
print("loss_iterative: {}, elapsed: {}".format(loss1, e1))

t = time.time()
loss2 = loss(X, Y, W)
e2 = time.time() - t
print("loss_vectorized: {}, elapsed: {}".format(loss2, e2))