function [label_img, label_img_soft, dist, centers ] = SLIC_Color_3D_Normals(lab, pcd, normals, user_params)
% beta_1: controls hardness of asg to compute weights (higher
% beta_1 -> more exclusive assignment)
% beta_2: controls how significant distance has to be for them to
% be similar (low beta_2 means when computing cluster centers many
% pixels will have similar influence)

%% SLIC Parameters
% Default Parameters
params.max_iter = 10; % Max iterations, 10 is usually more than enough
params.patch_sz = 32; % patch size parameter
params.max_c = 2.00;
params.max_s = 0.01;
params.max_n = 0.25;
params.beta_1 = 0.9;
params.beta_2 = 0.1;

% User Defined Parameters
if nargin > 2
    param_names = fieldnames(user_params);
    for i = 1:length(param_names)
        % Easy way to check that user_params has valid fields
        old_value = getfield(params, param_names{i});
        % Update params fields with the value from the user params
        param_value = getfield(user_params, param_names{i});
        params = setfield(params, param_names{i}, param_value);
    end
end

%% Initialization
n_row = size(lab, 1);
n_col = size(lab, 2);
grid_size = [n_row n_col];

% Make sure we have NaNs, not zeros
% for point cloud
pcd_zero = find(pcd(:,:,3) == 0);
pcd_zero = [pcd_zero; pcd_zero + prod(grid_size); pcd_zero + 2 * prod(grid_size)];
pcd(pcd_zero) = NaN;
% for normals
normal_zero = find(sum(normals,3) == 0);
normal_zero = [normal_zero; normal_zero + prod(grid_size); normal_zero + 2 * prod(grid_size)];
normals(normal_zero) = NaN;

% Need to know which pixels are valid
valid_pix = sum(isnan(lab) + isnan(pcd) + isnan(normals), 3) == 0;

% Sample pixels on a grid with stride psz,
% we use these indices to sample the initial cluster centers from data.
[grid_x, grid_y] = meshgrid((1+psz):psz:n_col, (1+psz):psz:n_row);
samp_inds = sub2ind(grid_size, grid_y(:), grid_x(:));
samp_inds_3D = cat(3, samp_inds, samp_inds + prod(grid_size), samp_inds + 2 * prod(grid_size));

%% Initialize the cluster centers
% Cluster center ids
n_centers = size(samp_inds,1);
% temp = num2cell(1:n_centers);
temp = num2cell(randperm(n_centers));
[centers(1:n_centers).id] = deal(temp{:});
% Cluster center (x, y) in image space
temp = num2cell(grid_x(:));
[centers(:).x] = deal(temp{:});
temp = num2cell(grid_y(:));
[centers(:).y] = deal(temp{:});
% Cluster center (x,y,z) in 3D space
temp = mat2cell(pcd(samp_inds_3D), ones(1,n_centers));
[centers(:).point] = deal(temp{:});
% Cluster center color
temp = mat2cell(lab(samp_inds_3D), ones(1,n_centers));
[centers(:).color] = deal(temp{:});
% Cluster center normal
temp = mat2cell(normals(samp_inds_3D), ones(1,n_centers));
[centers(:).normal] = deal(temp{:});

% We want to ignore invalid centers
center_is_valid = valid_pix(samp_inds)';

% Clear up data that is only used in initialization
clear temp samp_inds samp_inds_3D pcd_zero normal_zero

%% Initialize the labels and distances
n_soft = 4;
label_img = -1*ones(n_row, n_col, n_soft);
dist = inf(n_row, n_col, n_soft);

for iter = 1:max_iter
    % For each of the clusters:
    %   - calculate distances
    %   - update pixel assignments

    for c_ind = find(center_is_valid)
        %% Sample the neighborhood
        % Round cluster centers to image indices
        sam_r = round(centers(c_ind).y); % center row
        sam_c = round(centers(c_ind).x); % center column

        % Calculate neighborhood bounds
        sam_rt = max(1, sam_r - psz);
        sam_rb = min(n_row, sam_r + psz);
        sam_cl = max(1, sam_c - psz);
        sam_cr = min(n_col, sam_c + psz);

        % Sample neighborhood
        % lab color, position, normal, label, and distance
        sam_col = lab(sam_rt:sam_rb, sam_cl:sam_cr, :);
        sam_pos = pcd(sam_rt:sam_rb, sam_cl:sam_cr, :);
        sam_nor = normals(sam_rt:sam_rb, sam_cl:sam_cr, :);
        sam_labels = label_img(sam_rt:sam_rb, sam_cl:sam_cr, :);
        sam_dist = dist(sam_rt:sam_rb, sam_cl:sam_cr, :);

        % Find the distances from this center id.
        % Since we are updating this cluster center,
        % we want to reset these distances to inf and labels to -1.
        prev_inds = find(sam_labels == centers(c_ind).id);
        sam_labels(prev_inds) = -1;
        sam_dist(prev_inds) = Inf;

        % We compare against the maximum distances,
        % if the distance of this cluster center to a pixel is
        % less than this maximum, then we can replace it and assign
        % the pixel to this cluster center.
        [sam_xs, sam_ys] = meshgrid(1:size(sam_dist,2), 1:size(sam_dist,1));
        [min_dist, min_ind] = min(sam_dist, [], 3);
        min_ind = sub2ind(size(sam_dist), sam_ys, sam_xs, min_ind);
        % [max_dist, max_ind] = max(sam_dist, [], 3);
        % max_ind = sub2ind(size(sam_dist), sam_ys, sam_xs, max_ind);

        %% Compute the distances
        % Distance terms
        % Lab Color distance
        dist_color = sqrt(sum(bsxfun(@minus, sam_col, centers(c_ind).color).^2, 3));

        % Spatial distance
        dist_spatial = sqrt(sum(bsxfun(@minus, sam_pos, centers(c_ind).point).^2, 3));

        % Normal distance
        % sam_n normals and center_normal should be unit vectors
        % so the angle between the normals is acos(dot(n1,n2))
        % We also take min and max to make sure that
        % inputs to acos are in the range [-1 1].
        % This is necessary due to small numerical errors
        dist_normal = sum(bsxfun(@times, sam_nor, centers(c_ind).normal), 3);
        % ok_normal = ~isnan(dist_normal);
        % dist_normal(ok_normal) = min(1, dist_normal(ok_normal));
        % dist_normal(ok_normal) = max(-1, dist_normal(ok_normal));
        dist_normal = min(1, dist_normal);
        dist_normal = max(-1, dist_normal);

        % assert(all(dist_normal(:) > -1))
        % assert(all(dist_normal(:) <  1))
        dist_normal = acos(dist_normal);
        % assert(isreal(dist_normal))

        % Combine each of the distance terms
        % inside = cat(3, ...
        %                      dist_color/max_c, ...
        %                      dist_spatial/max_s, ...
        %                      dist_normal/max_n ...
        %                      );
        % sam_d2 = sum(inside.*inside,3).^(0.5);
        sam_d = sqrt(sum(cat(3, ...
                             dist_color/max_c, ...
                             dist_spatial/max_s, ...
                             dist_normal/max_n ...
                             ).^2,3));
        % assert(all(sam_d(~isnan(sam_d)) == sam_d2(~isnan(sam_d2))))

        %% Determine which pixel assignments should be updated
        % bwlabel ensures connected components
        % Note that if there are NaN terms they will not be updated
        % (this is what we want when there are NaNs and the center is not NaN).
        update = bwlabel(sam_d < min_dist,4);
        % update = bwlabel(sam_d < max_dist);
        % bwlabel can produce multiple connected components,
        % we want only the connected component overlapping the center
        update = (update == update((sam_r - sam_rt)+1, (sam_c - sam_cl)+1)) & update;

        % Update assignments in the sampled neighborhood
        % sam_labels(max_ind(update)) = centers(c_ind).id;
        % sam_dist(max_ind(update)) = sam_d(update);
        sam_labels(min_ind(update)) = centers(c_ind).id;
        sam_dist(min_ind(update)) = sam_d(update);


        % Update the main labels and distances
        label_img(sam_rt:sam_rb, sam_cl:sam_cr, :) = sam_labels;
        dist(sam_rt:sam_rb, sam_cl:sam_cr, :) = sam_dist;
    end

    %% Add clusters for any pixels that are unassigned
    % Find all lost labels and their indices
    lost_pix = all(label_img == -1, 3) & valid_pix;
    if any(lost_pix(:))
        [lost_y, lost_x] = find(lost_pix);
        lost_inds = sub2ind(grid_size, lost_y, lost_x);
        % Create new clusters on a psz step grid
        % find the unique clusters and use them as labels
        [unique_ids, ia, lost_ids] = unique(sub2ind(ceil(grid_size/psz), ...
                                                ceil(lost_y/psz), ceil(lost_x/psz)));

        % New cluster center ids
        max_id = numel(centers);
        new_ids = num2cell([1:length(unique_ids)] + max_id);
        lost_ids = lost_ids + max_id;

        % Update the image labels, distances, and centers
        label_img(lost_inds) = lost_ids;
        dist(lost_inds) = 0;
        [centers((max_id+1):(max_id+numel(new_ids))).id] = deal(new_ids{:});

        center_is_valid = [center_is_valid ones(1,numel(new_ids))];
    end

    % To deal with outliers having too much influecence, we weight
    % their influence to each cluster by finding soft assg to each
    % cluster (dist_ws), then using
    dist_exp_forws = exp(-beta_1*dist);
    dist_ws = bsxfun(@rdivide, dist_exp_forws, sum(dist_exp_forws, 3));
    dist_exp = exp(-beta_2*dist);

    %    center_active_inds = find(center_is_valid);
    %    for center_act_i = 1:numel(center_active_inds)
    %    center_i = center_active_inds(center_act_i); % Current cluster index
    %    center_id = center_ids(center_i); % cluster id

    for c_ind = find(center_is_valid)
        % Get points belonging to the cluster
        % NaNs should not be here at all.
        inside = (label_img == centers(c_ind).id);
        inside2 = bsxfun(@and, inside, valid_pix);

        % If this cluster contains at least ## pixels
        if any(inside2(:)) > 0
            cl_d_inds = find(inside2);
            [cl_y, cl_x, cl_d] = ind2sub(size(label_img), cl_d_inds);
            cl_inds = sub2ind(grid_size, cl_y, cl_x);
            cl_inds_3D = cat(3, cl_inds, cl_inds + prod(grid_size), cl_inds + 2 * prod(grid_size));

            % Compute the weight of each point to the cluster
            cluster_ws = dist_ws(cl_d_inds) .* dist_exp(cl_d_inds);
            cluster_ws = cluster_ws/sum(cluster_ws);

            % Get points belonging to the cluster
            cluster_c = lab(cl_inds_3D);
            cluster_s = pcd(cl_inds_3D);
            cluster_n = normals(cl_inds_3D);

            % Update cluster centers to the mean of its assigned members
            centers(c_ind).x = sum(cl_x.*cluster_ws);
            centers(c_ind).y = sum(cl_y.*cluster_ws);
            centers(c_ind).color  = sum(bsxfun(@times, cluster_ws, cluster_c),1);
            centers(c_ind).point  = sum(bsxfun(@times, cluster_ws, cluster_s),1);
            centers(c_ind).normal = sum(bsxfun(@times, cluster_ws, cluster_n),1);

            % If normal magnitude is less than threshold (.3)
            % then the cluster has high variance in its normals
            normal_norm = sqrt(sum(centers(c_ind).normal.^2,3));
            if normal_norm < .3
                center_is_valid(c_ind) = 0;
                %fprintf('cluster %d has too high of normal variance\n', c_ind);
            else
                centers(c_ind).normal = centers(c_ind).normal/normal_norm;
            end
        end
    end

    % Find any invalid centers
    center_is_valid = center_is_valid & ~sum(isnan([centers(:).color]), 3);
    center_is_valid = center_is_valid & ~sum(isnan([centers(:).point]), 3);
    center_is_valid = center_is_valid & ~sum(isnan([centers(:).normal]), 3);

    % Remove labels and distances of the inactive clusters
    for c_ind = find(~center_is_valid)
        % Get any points belonging to the cluster
        cl_d_inds = find(label_img == centers(c_ind).id);
        label_img(cl_d_inds) = -1;
        dist(cl_d_inds) = Inf;
    end

    % if DEBUG %& mod(iter,2)==0
    %     [grid_x, grid_y] = meshgrid(1:n_col, 1:n_row);
    %     [min_dist, min_ind] = min(dist, [], 3);
    %     min_ind = sub2ind(size(dist), grid_y, grid_x, min_ind);
    %     sfigure(imgid);
    %     imagesc(label_img(min_ind));
    %     title(sprintf('asgn @ iter %d', iter))
    %     drawnow;
    % end
end

[grid_x, grid_y] = meshgrid(1:n_col, 1:n_row);
[min_dist, min_ind] = min(dist, [], 3);
min_ind = sub2ind(size(dist), grid_y, grid_x, min_ind);

label_img_soft = label_img;
label_img = label_img(min_ind);

% if DEBUG
%     sfigure(imgid);
%     imagesc(label_img);colormap(hsv)
%     title(sprintf('final asg %g %g', beta_1, beta_2))
% end

% Remove the inactive clusters from bookkeeping
centers = centers(center_is_valid);

end
