function cluster_manifold(X, k, varargin)
% Manifold clustering based on local fit of polynomial functions.
%
% Input arguments:
% X:
%    data set with dimensions [n,N]
% k:
%    number of clusters
%
% References:
% Levente Hunyadi and Istvan Vajk, "Modeling by fitting a union of polynomial functions to data",
%    International Journal of Pattern Recognition and Artificial Intelligence, 27(2), 2013,
%    http://dx.doi.org/10.1142/S0218001413500043

% Copyright 2012 Levente Hunyadi

% Tested with MatLab 7.14.0.739 (R2012a) 64-bit

    if nargin < 2
        narginchk(0, 0);
        
        if 1
            X = example_petals;
            k = 4;
            cluster_manifold(X, k);
            return;
        elseif 0
            X = example_ellipses;
            k = 2;
            cluster_manifold(X, k, ...
                'fit_function', @ellipsefit_koopmans, ...
                'neighborhood_size', 50, ...
                'affinity_attenuation', 5 ...
            );
            return;
        end
    end

    validateattributes(X, {'numeric'}, {'nonempty','real','2d'});
    assert(size(X,1) >= 2 && size(X,1) <= 3, 'Expected data in two or three dimensions.');
    validateattributes(k, {'numeric'}, {'positive','integer','scalar'});
    options = struct( ...
        'neighborhood_size', 50, ...
        'affinity_attenuation', 1 ...
    );
    switch size(X,1)
        case 2
            options = merge_options(options, 'fit_function', @quad2dfit_koopmans);
        case 3
            options = merge_options(options, 'fit_function', @quad3dfit_koopmans);
    end
    options = merge_options(options, varargin{:});
    
    validateattributes(options.neighborhood_size, {'numeric'}, {'positive','integer','scalar'});
    validateattributes(options.fit_function, {'function_handle'}, {'scalar'});
    validateattributes(options.affinity_attenuation, {'numeric'}, {'positive','scalar'});
    N = size(X,2);
    d = size(X,1);  % data dimension

    % normalize data to unit noise magnitude
    sigma = repmat(sqrt(d)/d, 1, d);
    %sigma = [ cosd(45) sind(45) ];  % for 2D
    Sigma = diag(sigma);
    X = Sigma \ X;
    
    if 0
        dataplot(X);
        return;
    end
    
    fitfunc = options.fit_function;

    % obtain parameter estimates for each point based on its locality (neighborhood)
    Mu = zeros(1,N);
    switch size(X,1)
        case 2
            Theta = zeros(6,N);
        case 3
            Theta = zeros(10,N);
    end
    neighborhood = cell(1,N);
    for pix = 1 : N  % point index
        [~,nix] = optimum_neighborhood(X, pix, options.neighborhood_size, fitfunc);
        [Theta(:,pix),Mu(pix)] = fitfunc(X(:,nix));  % estimate made with unit noise magnitude assumption for each direction
        neighborhood{pix} = nix;
    end

    if 0
        neighborhood_figure(X, Theta, neighborhood);
        return;
    end

    % mean noise parameter estimate
    mu = median(Mu);

    D = zeros(N,N);  % distance of each point to an estimated hyperplane, rows = points, cols = hyperplanes
    for pix = 1 : N  % point index
        D(:,pix) = project_hyperplane(X, Theta(:,pix));
    end

    if 0
        dataplot(X);
        return;
    end
    
    if 0
        %ranged_distance_figure(X, Theta, D);
        distance_figure(X, Theta);
        return;
    end
    
    % scatter matrix
    if 0
        S = exp(-options.affinity_attenuation * sqrt(D' .* D));
        f = cluster_spectral(S, k);
    else
        S = exp(-options.affinity_attenuation * D);
        f = cluster_bestwcut(S, k);
    end
    
    if 1
        grouping_figure(X, f, fitfunc);
        switch size(X,1)
            case 2
                %axis off
                lb = min(min(X));
                ub = max(max(X));
                xlim([lb,ub]);
                ylim([lb,ub]);
        end
    end
    
    f = kplanes(X, f, fitfunc);

    if 1
        grouping_figure(X, f, fitfunc);
    end
end

function [count,nix] = optimum_neighborhood(X, pix, count, fitfunc)

    nix = point_neighborhood(X, pix, count);  % nearest neighbors
    [p,mu] = fitfunc(X(:,nix));
    if 1
        e = mu / numel(nix);  % normalized error
    else
        e = project_hyperplane(X, p) / numel(nix);
    end
    e_last = Inf;

    while e < e_last
        count = count + 5;
        nix = point_neighborhood(X, pix, count);  % neighbors
        [p,mu] = fitfunc(X(:,nix));

        e_last = e;
        if 1
            e = mu / numel(nix);  % normalized error
        else
            e = project_hyperplane(X, p) / numel(nix);
        end
    end

    if 0
        hold('on');
        dataplot(X(:,nix), 'Color', 'b');
        switch size(X,1)  % dimensionality
            case 2
                theta = quad2dfit_koopmans(X(1,nix), X(2,nix));
                imconic(theta);
        end
        hold('off');
    end
end

function nix = point_neighborhood(X, pix, count)
% Finds the nearest neighborhood of a data point.
%
% Input arguments:
% X:
%    data set with dimensions [n,N]
% pix:
%    the index of the key point in the data set from which to measure distance
% count:
%    the (minimum) number of points to comprise the neighborhood around the key point
%
% Output arguments:
% nix:
%    the indices of the nearest neighbors

    validateattributes(X, {'numeric'}, {'nonempty','real','2d'});
    validateattributes(pix, {'numeric'}, {'positive','integer','scalar','<=',size(X,2)});

    % the neighborhood size is upper-bounded by the number of points
    count = min(count, size(X,2));

    % distance of other points from key point
    dist = sum(bsxfun(@minus, X, X(:,pix)).^2, 1).^0.5;
    
    [~,dist_ord] = sort(dist);
    nix = dist_ord(1 : count);  % indices of nearest neighbors

    if 0
        figure;
        hold('on');
        dataplot(X, 'Color', 'b');
        dataplot(X(:,nix), 'Color', 'r');
        dataplot(X(:,pix), 'Color', 'k', 'Marker', 'o');
        hold('off');
    end
end

function [d, Xp] = project_hyperplane(X, theta)
% Fits a hyperplane to a set of points.
%
% Input arguments:
% X:
%    data set with dimensions [n,N]
%
% Output arguments:
% p:
%    a parameter vector with dimensions [d,1]
% d:
%    a vector of point distances to the hyperplane

    validateattributes(X, {'numeric'}, {'nonempty','real','2d'});
    validateattributes(theta, {'numeric'}, {'nonempty','real','vector'});

    switch size(X,1)  % dimensionality
        case 2
            if is_ellipse(theta)
                x = X(1,:);
                y = X(2,:);
                
                [c1,c2,a,b,phi] = ellipse_im2ex(theta);
                Xp = ellipseproj([x;y], [c1,c2,a,b,phi]);
                Xp = real(Xp);  % cancel imaginary components
                d = realsqrt(sum((X - Xp).^2, 1));
            else
                [d, Xp] = pqdist(X, theta);
            end
        case 3
            [d, Xp] = pqdist(X, theta);
    end

    if 0
        figure;
        hold('all');
        dataplot(X, 'Color', 'b');
        imconic(theta);
        dataplot(Xp, 'Color', 'r');
        hold('off');
    end
    
    validateattributes(d, {'numeric'}, {'nonempty','real','vector'});
end

function f = kplanes(X, f, fitfunc)
% K-planes algorithm.
%
% Input arguments:
% X:
%    data set with dimensions [n,N]
% f:
%    an initial clustering
% fitfunc:
%    a fit function that expects two vector parameters
    
    N = size(X, 2);
    k = max(f);

    iter = 0;
    while iter < 4
        iter = iter + 1;

        switch size(X,1)
            case 2
                Theta = zeros(6,N);
            case 3
                Theta = zeros(10,N);
        end
        D = zeros(N,k);  % distance of each point to an estimated hyperplane, rows = points, cols = hyperplanes
        
        for i = 1 : k
            % estimate curve parameters for groups
            Theta(:,i) = fitfunc(X(:,f==i));  % estimate made with unit noise magnitude assumption for each direction

            % assign close points to curves
            D(:,i) = project_hyperplane(X, Theta(:,i));
        end
        
        [~,f] = min(D,[],2);
    end
end
