clear all;
close all;

rng(1)
% 2-dim toy data generation
C=5;
samples_per_class=500;
sigma=1.0;
num_samples = C*samples_per_class;

% classes A-E with negative labels
noise_dim = 100;
feature_dim = 2;
dim = noise_dim+feature_dim;

cluster_centers = [2  -4; -3  3; 3   1; -5 -2; -4 -5];
colors_samples = {'og','sr','xb','+m','dy'};
colors_classifiers = {'g','r','b','m','y'};

figure
hold on
title('Training Data')
data = zeros(dim+1, num_samples);
for j=1:C
    data(:, (j-1)*samples_per_class+1:j*samples_per_class) = ...
        [repmat(cluster_centers(j, :),samples_per_class,1) + randn(samples_per_class,2)*sigma ...
        randn(samples_per_class, noise_dim)*3 ...
        j*ones(samples_per_class, 1)]';
    
    scatter(data(1, (j-1)*samples_per_class+1:j*samples_per_class),...
        data(2, (j-1)*samples_per_class+1:j*samples_per_class),colors_samples{j});
end

% Suffle data
shuffle = randperm(size(data, 2));
data = data(:, shuffle);

% Split data into training and test set
N = num_samples-0.1*num_samples;
train_data = data(:, 1:N);
test_data = data(:, N+1:end);




X = train_data(1:dim, :)';
%Normalize data
X = X/max(X(:));
y = train_data(end, :)';

lambda1 = 0.01;
lambda2 = 0.00;

L = norm(X,'fro')^2/N + lambda1;
tau = 1/L;


%%

max_iter=6000;



% compute classifier scores
classifier = @(W, b, X) X*W + repmat(b, size(X, 1), 1);

% evaluate loss given classifier scores
scores = @(classifier) exp(classifier) ./ repmat(sum(exp(classifier), 2), 1, C);
loss = @(scores) sum(-log(scores(sub2ind(size(scores), (1:N)', y))), 1);

objective = @(W, b) loss(scores(classifier(W, b, X)))/N ...
    + 0.5 * lambda1 * sum(W(:).^2)...
    + 0.5 * lambda1 * sum(b(:).^2)...
    + lambda2 * sum(W(:));

grad_W = @(W, scores, j) ...


grad_b = @(b, scores, j) ...


%% TODO implement proximal gradient here...

%% Plot objective value
figure;
plot(obj_val);


%% Plot classifiers
s = 0.1;
[x1Grid, x2Grid] = meshgrid(-4:s:4,...
    -4:s:4);
xGrid = [x1Grid(:),x2Grid(:)];
scores_train = classifier(W(1:2,:), b, xGrid);
figure;
hold on
for j=1:C
    scatter(data(1, data(end, :)==j),...
        data(2, data(end, :)==j),colors_samples{j});
    
    [~,h] = contour(x1Grid,x2Grid,reshape(scores_train(:, j),size(x1Grid)),[0 0],'k');
    h.LineColor = colors_classifiers{j};   
end
drawnow