function [all_theta] = mlr_optim(X, y, num_labels, lambda)

% Some useful variables
m = size(X, 1);
n = size(X, 2);

% You need to return the following variables correctly 
% unrolling trick
all_theta = zeros(num_labels * (n + 1), 1);

% Add ones to the X data matrix
X = [ones(m, 1) X];
X = X(:);

% MLR code
options = optimset('GradObj', 'on', 'MaxIter', 50);
initial_theta = all_theta;
tic
all_theta = fmincg(@(t)(costFunction(t, X, y, num_labels, ...
    lambda)), initial_theta, options)';
toc
all_theta = reshape(all_theta, num_labels, n + 1);
end

function [J, grad] = costFunction(theta, X, y, num_labels, lambda)
    m = size(y, 1);
    n = size(X, 1) / m;
    theta = reshape(theta, num_labels, n);
    X = reshape(X, m, n);
    J = 0;
    grad = zeros(size(theta));
    for nex = 1:m
        J = J - log(exp(theta(y(nex),:) * X(nex,:)') / ...
            sum(exp(theta * X(nex,:)')));
        Z = sum(exp(theta * X(nex,:)'));
        for c = 1:(num_labels - 1)
            grad(c,:) = grad(c,:) + (X(nex,:) * ...
                (exp(theta(c,:) * X(nex,:)')/Z - (c == y(nex))) + ...
                lambda / m / m * theta(c,:));
        end
    end
    grad = reshape(grad, num_labels * n, 1);
    % Regularisation issue
    for c = 1:(num_labels - 1)
        J = J + (lambda / (2 * m)) * sum(sum(theta.^2));
    end
end

