%%%%
% 人工智能的数学基础 -- 2019秋
% 华东师范大学 数学科学学院 黎芳（教授）
%
% # Homework2 -- 最近邻
%
% > 截止时间：2019年10月6日23:59
%
% 你需要利用以下任意一种方式实现`KNN_impl`函数
%
% - 朴素KNN，即两两比较
% - KD树（奖励分1）
%
% MIT license 陈久宁 <johnnychen94@hotmail.com>
%%%%

[train_input, train_label] = readData_iris();
train_input = train_input(1:2, :); % 为了方便绘图，只使用前两组特征

input = train_input;
Verbose = true;

K = 1;
preds_1 = KNN(input, train_input, train_label, K, Verbose);

K = 5;
preds_2 = KNN(input, train_input, train_label, K, Verbose);


% 你只需要修改该函数即可让代码正常运行
function preds = KNN_impl(input, train_input, train_label, K)
% inplementation of KNN
% Inputs:
%   input: array of shape [2, n1]
%          test input data
%   train_input: array of shape [2, n2]
%          training input data
%   train_label: array of shape [1, n2]
%          training label
%   K: integer
%      number of neighborhood elements

    error("Not Implemented")
end


function preds = KNN(input, train_input, train_label, K, Verbose)
% K-nearest neighborhood classification algorithm
%
% Inputs:
%   input: array of shape [2, n1]
%          test input data
%   train_input: array of shape [2, n2]
%          training input data
%   train_label: array of shape [1, n2]
%          training label
%   K: integer
%      number of neighborhood elements
%   Verbose: default: true
%      true to show results

    assert(size(input,1)==2, "input is array of shape [2, n].")

    preds = KNN_impl(input, train_input, train_label, K);

    if Verbose
        n = 50;
        x_min = min([input(1,:), train_input(1,:)]);
        x_max = max([input(1,:), train_input(1,:)]);
        y_min = min([input(2,:), train_input(2,:)]);
        y_max = max([input(2,:), train_input(2,:)]);

        [X, Y] = meshgrid(linspace(x_min,x_max, n), linspace(y_min, y_max, n));
        Z = KNN([X(:)';Y(:)'], train_input, train_label, K, false);
        Z = reshape(Z, n, n);

        figure; hold on;
        title(sprintf("%d-Nearest Neighborhood", K));
        view(0, 90)
        colormap(winter);

        mesh(X, Y, Z, "LineWidth", 10);

        offset = max(Z(:)) + 1; % make all marker visible
        plot3(input(1,preds==0), input(2,preds==0), preds(preds==0)+offset, "ro", "LineWidth", 20)
        plot3(input(1,preds==1), input(2,preds==1), preds(preds==1)+offset, "b+", "LineWidth", 20)
        plot3(input(1,preds==2), input(2,preds==2), preds(preds==2)+offset, "kx", "LineWidth", 20)
    end
end

function [X,Y] = readData_iris()
% readData_iris reads iris dataset
    dataurl    = 'https://archive.ics.uci.edu/ml/machine-learning-databases/iris/iris.data';
    datadir    = 'data';
    filename   = fullfile(datadir,'Iris.txt');

    if ~isfolder(datadir)
        mkdir(datadir);
    end
    if ~isfile(filename)
        websave(filename, dataurl);
    end

    assert(isfile(filename), [filename,' not found']);

    data = readtable(filename);
    X = table2array(data(:,1:4))';
    Y = table2cell(data(:,5)); % convert table to cell
    Y = cellfun(@labelclass, Y)'; % convert string to scalar

    % shuffle data
    idx = randperm(size(X,2));
    X = X(:, idx);
    Y = Y(:, idx);
    return

    function num = labelclass(str) % onehot encoding
        switch str
            case 'Iris-setosa'
                num = 0;
            case 'Iris-versicolor'
                num = 1;
            case 'Iris-virginica'
                num = 2;
            otherwise
                error(['unrecognized iris class', str])
        end
    end
end
