function [mdiscr] = TImrf(emesh, selection, varargin)
% [mdiscr] = TImrf(emesh, selection, varargin)
% Runs a Pott's MRF over the mesh and returns discrete staining values im mdisc:
% 0 = no staining
% 0.5 = possible staining
% 1 = definite staining
% Arguments:
% emesh: emesh
% selection: indices in emesh or mesh itself
% varargin: 
%     'verbose' <number>
%         0 (or missing verbose) = no progress
%         1 text only progress
%         2 or 3: show figure of original mesh and ...
%             2 ... k-means and MRF 
%             3 ... MRF only
%     'nocorona'
%         Do not attempt to remove a corona (unstained triangles around circumference)
%         Useful with raw meshes but useless with DIC corrected meshes
%     'consensus' <n>
%         runs MRF n times and uses the consensus
%     'kmeans'
%         detects staining quick and dirty with k-means clustering

    verbose = 0;
    corona = [];
    consensus = 1;
    selection_mesh = 0;
    quick = [];

    if ~isempty(varargin)
        i = 1;
        
        % keep it compatible to old version meshdiscretize_emrf(emesh, selection, verbose)
        if isnumeric(varargin{1}),
            verbose = varargin{1};

        else
        
            while i <= length(varargin),
                switch(varargin{i}),
                    case 'verbose'
                        i = i+1;
                        verbose = varargin{i};
                    case 'nocorona'
                        corona = 1;
                    case 'single_mesh'
                        selection_mesh = 1;
                    case 'kmeans'
                        quick = 1;
                    case 'consensus'
                        i=i+1;
                        consensus = varargin{i};
                    otherwise
                        error('Argument %s not known', varagin{i});
                end
                i = i+1;
            end
            
        end
    end

    ini_clusters = 4;
    
    % to make it possible to use both idx or mesh as selection
    if size(selection, 2) ~= 1 || selection_mesh == 1,
        is_mesh = 1;
        sel_size = size(selection,2);
    else
        is_mesh = 0;
        sel_size = size(selection,1);
    end

    mdiscr = zeros(311, sel_size);
    t_conn = fully_connected(emesh);
    
    for embryo = 1:sel_size, 
        clusters = ini_clusters;
        
        if verbose > 0, fprintf('%d. Embryo, Number %d\n', embryo, selection(embryo)); end
        if is_mesh == 0, 
            stain = emesh.stain(:,selection(embryo));
        else
            stain = selection(:,embryo);
        end

        lvar = local_var(stain, t_conn);
        if verbose > 0, fprintf('Local variance = %f\n', lvar); end
        
        if isempty(corona) || verbose == 2, 
            if verbose > 0, fprintf('Initial clustering\n'); end
            [sk, sc] = safe_kmeans(stain, clusters);
                        
            if isempty(corona), 
                sortclus = sortrows([(1:clusters)' sc], -2);
                max_clus = sortclus(1,1);
                t_fedge = find(emesh.t_edge == 2); % full edge
                t_ftedge = find(emesh.t_edge > 0); % touching edge
                t_bri_clus = find(sk == max_clus);
                if verbose > 0, 
                    fprintf('Edge corona: %d > %d? or %d > %d?\n', ...
                        length(intersect(t_fedge, t_bri_clus)), length(t_bri_clus), ...
                        length(intersect(t_ftedge, t_bri_clus)), length(t_bri_clus));
                end
                t_maxedge = max( [ length(intersect(t_fedge, t_bri_clus)), length(intersect(t_ftedge, t_bri_clus))] );
                if t_maxedge > (length(t_bri_clus) * 0.8),
                    if verbose > 0, fprintf('Edge corona found, removing\n'); end
                    stain(sk == max_clus) = sortclus(2,2);
                end
            end
        else
            % no sk/sc assignment from k-means
            sk = zeros(311,1); sc=zeros(311,1); 
            if verbose > 0, fprintf('No corona removal\n'); end
        end

        if consensus > 1,
            if verbose > 0, fprintf('MRF segmenting, consensus of %d times\n', consensus); end
        else
            if verbose > 0, fprintf('MRF segmenting\n'); end
        end
        sk_mrf = ones(311,1);
        sk_mrf_a = zeros(311, consensus);
        for i=1:consensus,
            % sk_mrf = en_mrf(t_conn, sk_sort, stain, verbose);
            if isempty(quick),
                [sk_mrf_a(:, i), sc_mrf] = en_mrf(t_conn, 3, stain, verbose);
            else
                [sk_i, sc_mrf] = safe_kmeans(stain, 3);
                sortclus = sortrows([(1:length(sc_mrf))' sc_mrf], -2);
                for j=1:3,
                    sk_mrf_a(sk_i == sortclus(j,1),i) = j;
                end
            end
        end
        if consensus > 1, 
            % sk_var = var(sk_mrf_a, 0, 2);
            % sk_mrf(sk_var == 0) = sk_mrf_a(sk_var == 0, 1);
            sk_mrf = median(sk_mrf_a, 2);
        else
            sk_mrf = sk_mrf_a(:,1);
        end
        
        % mdiscr(:,embryo) = sk_mrf;

        disc_val = 1 / (length(unique(sk_mrf)) - 1);
        mdiscr(:,embryo) = (sk_mrf - 1) .* disc_val;

        if verbose > 1,
            d_clus = zeros(311,1);
            d_mrf = zeros(311,1);
            for i = 1:clusters,
                d_clus(sk == i) = uint8(sc(i));
    %             d_mrf(sk_mrf == i) = uint8(sortclus(i,2));
                d_mrf(sk_mrf == i) = 255 / clusters * (clusters - i);
            end
            
            figure;
            if verbose > 2,
                subplot(2,1,1); displaymesh(stain, emesh.p_scale, emesh.t);
                subplot(2,1,2); displaymesh(d_mrf, emesh.p_scale, emesh.t);
            else
                subplot(3,1,1); displaymesh(stain, emesh.p_scale, emesh.t);
                subplot(3,1,2); displaymesh(d_clus, emesh.p_scale, emesh.t);
                subplot(3,1,3); displaymesh(d_mrf, emesh.p_scale, emesh.t);
            end
        end
    end

end



function [sk, sc] = en_mrf(in_t_conn, no_clus, stain, verbose)

    % parameters
    gibbs_sampler =1;
    max_iter = 200;
    update_centroids = 1;
    lambda = 1;
    alpha = 5;
    gamma = 0.1;
    epsilon = 10;
    
    
    sk = fix(rand(311,1) .* no_clus) + 1;
    sc = rand(no_clus,1) .* 255;
    sv = ones(no_clus,1);

    t_conn = in_t_conn .* ~eye(size(in_t_conn,1));
    t_conn = logical(t_conn);
    iter = 1;
    
    en = zeros(311,1);
    for i=1:size(in_t_conn,1),
        en(i) = g_energy(i, sk(i));
    end

    % essentially a MRF sampler with edge detection from p. 55 Stan Li book
    % using the more understandable formula 1.5 from G. Winkler book
    % The thing is either greedy (ICM) or Gibbs Sampler with simulated
    % annealing (set variable gibbs_sampler = 1)
    
    T = 1;
    
    while 1,
        tri_change = 0;

        if update_centroids == 1,
            for i=1:length(sc),
                sc(i) = mean(stain(sk == i));
                sv(i) = std(stain(sk == i));
            end
        end
        
        for i=1:size(in_t_conn,1),
            
            new_lab = fix(no_clus * rand(1,1)) + 1;
            h = g_energy(i, new_lab);
            
            if h < en(i)
                sk(i) = new_lab;
                en(i) = h;
                tri_change = tri_change + 1;
            elseif (rand(1,1) > 1-T) && gibbs_sampler == 1,
                sk(i) = new_lab;
                en(i) = h;
                tri_change = tri_change + 1;
            end
        end

        if gibbs_sampler == 1,
            T = T /log(max_iter+iter) * log(max_iter);
        end
        
        if verbose > 1, fprintf('Iteration %d -> %d changes\n', iter, tri_change); end
        
        iter = iter + 1;
        if tri_change == 0 || iter > max_iter,
            break;
        end
    end

    if update_centroids == 1,
        sk_sort = zeros(length(sk),1);
        sc_sort = sortrows([(1:length(sc))' sc], -2);
        for i=1:length(sc),
            sk_sort(sk == sc_sort(i,1)) = i;
            if verbose > 1, fprintf('%d=%.1f, ', i, sc(i)); end
        end
        if verbose > 1, fprintf('\n'); end
        sk = sk_sort;
    end
    
    
    function h = g_energy(i, nlabel)
        dc = stain(i);
        dn = stain(t_conn(i,:));
        fc = sc(nlabel);
        fn = sc(sk(t_conn(i,:)));
        
        edge = (abs(dn - dc) > epsilon) > 0;
        
        d_f = sum(([dn; dc] - [fn; fc]).^2);
        f = sum(((fn - fc).^2) .* (1-edge) + gamma .* edge);
        
        h = alpha * d_f + lambda * f;
    end


end




function [k, c] = safe_kmeans(data, clusters)

    iter = 1;
    max_iter = 100;
    while iter < max_iter
        try
            %[k, c] = kmeans(data, clusters, 'replicates', 5);
            [k, c] = kmeans(data, clusters, 'replicates', 5, 'emptyaction', 'drop', 'start', 'uniform');
            %[k, c] = kmeans(data, clusters, 'replicates', 5, 'distance', 'correlation');
            break;
        catch
            iter = iter+1;
            er = lasterror;
            fprintf('%d. Iteration: %s\n', iter, er.message);
        end
    end
    if iter == max_iter,
        error('No kmeans clustering found');
    end
end


function [var] = local_var(data, t_conn)
    var = 0;
    for i=1:size(t_conn,1),
        var = var + std(data(t_conn(2,:)));
    end
    var = var /size(t_conn,1);
end

% from find_restricted
function [t_conn] = fully_connected(emesh)

    for i=1:length(emesh.t)
        tnfr{i} = [];   
        tnf{i} = []; 
        for j=1:size(emesh.t,2)       
            [prow, pcol] = find(emesh.t==emesh.t(i,j));
            tnfr{i} = [tnfr{i}, prow(prow ~= i)'];
        end
        for j=1:length(emesh.tn{i})
           if length(find(tnfr{i} == emesh.tn{i}(j))) > 1,
            tnf{i} = [tnf{i}, emesh.tn{i}(j)];
          end
        end
    end

    t_conn=zeros(311,311);
    for i=1:length(tnf), 
        for j=1:length(tnf{i}), 
            t_conn(i,tnf{i}(j)) = 1; 
            t_conn(i,i) = 1; 
        end;
    end

    t_conn = logical(t_conn);
end