%% clear before start
clear

%% Import Train File
disp('Read training data...')

%Open training file
trainingFile = fopen('usps.train');

%Read training data
K = fscanf(trainingFile , '%i' , 1);
D = fscanf(trainingFile , '%i' , 1);
numbers = fscanf(trainingFile , '%i' , inf);
N=size(numbers,1)/(D+1);
numbers=reshape(numbers, D+1, N);
train_correct_class(:,1)=numbers(1,:);
train_samples=numbers(2:end,:)';
clear numbers;

fclose(trainingFile);
disp('...done.')

%% Prepare training data
disp('Prepare train data...')
prior = zeros(K:1);
train = cell(K,1);
for k=1:K
    selector=(train_correct_class==k);
    prior(k,1)=sum(selector)/N;
    train{k,1} = train_samples(selector,:);
end
disp('...done.')

%% Calculate gauss distribution
disp('Calculate gauss...')
mu=zeros(K,D);
sigma=zeros(D,D);
for k=1:K
    Nk=size(train{k,1},1);
    mu(k,:)=sum(train{k,1})/Nk;
    
    for i=1:Nk
        update=train{k,1}(i,:)-mu(k,:);
        sigma=sigma+update'*update;
    end
end
sigma=sigma/N;
diagsigma = diag(diag(sigma));
disp('...done.')

%% Import test data
disp('Read test data...')

%Open training file
testFile = fopen('usps.test');

numbers = fscanf(testFile , '%i' , inf);
Ntest=size(numbers,1)/(D+1);
numbers=reshape(numbers, D+1, Ntest);
test_correct_class(:,1)=numbers(1,:);
test_samples=numbers(2:end,:)';
clear numbers;

fclose(testFile);
disp('...done.')

%% Test on test data
estimated=zeros(Ntest,1);
for i=1:Ntest
    logprob=zeros(K,1);
    x=test_samples(i,:);
    for k=1:K
        xmu=x-mu(k,:);
        logprob(k,1)=log(prior(k,1))-1/2*xmu*diagsigma*xmu';
    end
    [p,estimated(i)]=max(logprob);
%     if (estimated(i)~=test_correct_class(i,1))
%         disp(i)
%         disp(c)
%         disp(test_correct_class(i,1))
%         contourf(reshape(x,16,16),'DisplayName','y');figure(gcf)
%         pause
%     end
end
error=1-sum(estimated(:,1)==test_correct_class(:,1))/Ntest;
fprintf('Error rate is %.2f%%\n',error)

confusion=zeros(K,K);
for i=1:Ntest
    confusion(test_correct_class(i,1),estimated(i,1))=confusion(test_correct_class(i,1),estimated(i,1))+1;
end
disp('Confusion matrix:')
disp(confusion)
contourf(confusion,'DisplayName','y');figure(gcf)
