function [ci, csize, fassign, T, Qmax, id_fuzzy] = assign_optim(X, gradFlowEQ, nodeInfo, opt)
% Optimal assignment of macrostates by transition probability
%	[ci, csize, fassign, T, Qmax, id_fuzzy] = assign_optim(X, gradFlowEQ, nodeInfo, opt)
%

%if nargin<4,
%	opt.fname = 'fuzzy_assign.data';
%end

MAXITER = 10;

numMicro = max(size(X));
numNode = length(nodeInfo);

MacroStates = find(diag(gradFlowEQ)==1);
numMacro = length(MacroStates);

% Compute the implied equilibrium distribution from each node. 
distEQ = gradFlowEQ(MacroStates,:);

% cluster index
ci = zeros(numMicro,1);

% deterministic assignment on basins of local maxima
for i=1:numMacro,
	macroBasin = find(distEQ(i,:)==1);
	for j=1:length(macroBasin), 
		if j==1,
			id_macro{i} = nodeInfo{macroBasin(j)}.set;
		else
			id_macro{i} = union(id_macro{i},nodeInfo{macroBasin(j)}.set);
		end
	end
	ci(id_macro{i}) = i;
end

% fuzzy microstates on barrier
id_fuzzy = find(ci==0);
disp(sprintf(' Number of barrier microstates: %d',length(id_fuzzy)));

% Construct new transition count matrix from X
T = speye(numMacro+length(id_fuzzy));

for i=1:numMacro,
	row_id = find(ci==i);
	for j=i:numMacro,
		col_id = find(ci==j);
		T(i,j) = sum(sum(X(row_id,col_id)));
	end
	for j=1:length(id_fuzzy),
		T(i,j+numMacro) = sum(X(row_id,id_fuzzy(j)));
	end
end
T = T + transpose(triu(T,1));
T(numMacro+1:numMacro+length(id_fuzzy),numMacro+1:numMacro+length(id_fuzzy)) = X(id_fuzzy,id_fuzzy);

% Construct column Markov transition matrix.
d = sum(T,2);
n = length(d);
M = T*sparse(1:n,1:n,1./d);
Mp = M;

eps = 1e-4; % small for tie-breaker

% fuzzy assignment probability matrix
fass = [eye(numMacro), zeros(numMacro, length(id_fuzzy))];
fass (:,numMacro+1:end) = full(Mp(1:numMacro,numMacro+1:end));

iter = 1;
fassign{iter} = fass;
[fass_sort, id_fass_sort]=sort(fass, 1, 'descend');
id_domore{iter} = find( fass_sort(1,:)<eps | (fass_sort(1,:)-fass_sort(2,:))<eps);
disp(sprintf(' Number of empty assignment: %d', length(id_domore{iter})));

% find the optimal assignment
CI{iter} = ci;
CI{iter}(id_fuzzy) = id_fass_sort(1,numMacro+1:end);
CSIZ{iter} = histc(CI{iter},unique(CI{iter}));

Q(iter) = metastab(X,CI{iter});
disp(sprintf(' Metastability Q = %f',Q(iter)));

Qmax = 0;

if Qmax<Q(iter),
	Qmax = Q(iter);
	iter_max = iter;
	ci = CI{iter};
	csize = CSIZ{iter};
end

while (~isempty(id_domore{iter}) & iter<=MAXITER),
	ci(id_domore{iter})=0;
	
	% reconstruct new transition count matrix from X
	T = speye(numMacro+length(id_domore{iter}));
	for i=1:numMacro,
		row_id = find(ci==i);
		for j=i:numMacro,
			col_id = find(ci==j);
			T(i,j) = sum(sum(X(row_id,col_id)));
		end
		for j=1:length(id_domore{iter}),
			T(i,j+numMacro) = sum(X(row_id,id_domore{iter}(j)));
		end
	end
	T = T + transpose(triu(T,1));
	T(numMacro+1:numMacro+length(id_domore{iter}),numMacro+1:numMacro+length(id_domore{iter})) = X(id_domore{iter},id_domore{iter});

	% reconstruct column Markov transition matrix.
	d = sum(T,2);
	n = length(d);
	M = T*sparse(1:n,1:n,1./d);
	Mp = M;

	iter = iter + 1;
	disp(sprintf(' Iteration %d.',iter));	

	Mp = Mp^(iter);
	fass = full(Mp(1:numMacro,numMacro+1:end));
	[fass_sort, id_fass_sort]=sort(fass, 1, 'descend');

	id_domore{iter} = id_domore{iter-1}(find( fass_sort(1,:)<eps | (fass_sort(1,:)-fass_sort(2,:))<eps));
	disp(sprintf(' Number of empty assignment: %d', length(id_domore{iter})));

	% find the optimal assignment
	CI{iter} = ci;
	CI{iter}(id_domore{iter-1}) = id_fass_sort(1,:);
	CSIZ{iter} = histc(CI{iter},unique(CI{iter}));
	Q(iter) = metastab(X,CI{iter});
	disp(sprintf(' Metastability I (first-reach): %f',Q(iter)));
	
	if Qmax<Q(iter),
		Qmax = Q(iter);
		iter_max = iter;
		ci = CI{iter};
		csize = CSIZ{iter};
	end

	% rearrange the fuzzy assignment
%	fassign{iter} = fass;
%	fassign{iter}(:,numMacro+1:end) = full(Mp(1:numMacro,numMacro+1:end));
%	fassign{iter}(:,id_domore{iter}) = ones(numMacro,length(id_domore{iter}))/numMacro;
%	
%	% find the optimal assignment
%	[F_rowsort, id_rowsort]=sort(fassign{iter},1,'descend');
%	CI{iter}(id_fuzzy) = id_rowsort(1,numMacro+1:end);
%	CSIZ{iter} = histc(CI{iter},unique(CI{iter}));
%
%	Q(iter) = metastab(X,CI{iter});
%	disp(sprintf(' Metastability II (all-iteration): %f',Q(iter)));
%	
%	if Qmax<Q(iter),
%		Qmax = Q(iter);
%		ci = CI{iter};
%		csize = CSIZ{iter};
%	end
end

% [Qmax,iter_max]=max(Q);
% ci = CI{iter_max};
% csize = CSIZ{iter_max};

disp(sprintf('---- Maximal metastability reached at iteration %d: %f',iter_max,Qmax));

%% write to the fuzzy assign files
%if nargin == 4,
%	fname = opt.fname; % 'fuzzy_assign.data';
%		
%	ci_fuzzy = ci;
%	id_fuzzy = id_domore{iter_max};
%	
%	ci_fuzzy(id_fuzzy) = 0;
%	
%	% deterministic new macro assignment
%	for i=1:numMacro,
%		id_macro{i} = find(ci_fuzzy==i);
%	end

%	% Construct new transition count matrix from X
%	T = speye(numMacro+length(id_fuzzy));

%	for i=1:numMacro,
%		for j=i:numMacro,
%			T(i,j) = sum(sum(X(id_macro{i},id_macro{j})));
%		end
%		for j=1:length(id_fuzzy),
%			T(i,j+numMacro) = sum(X(id_macro{i},id_fuzzy(j)));
%		end
%	end
%	T = T + transpose(triu(T,1));

%	T(numMacro+1:numMacro+length(id_fuzzy),numMacro+1:numMacro+length(id_fuzzy)) = X(id_fuzzy,id_fuzzy);

%	fid = fopen(fname,'w');
%	
%	fprintf(fid,'%d\t%d\n',size(T,1),numMacro);
%	for i=1:size(T,1),
%		if i<=numMacro,
%			fprintf(fid,'%d\t%d\t%d\t%d\n',i,1,i,1);
%		else
%			fprintf(fid,'%d\t%d',i,numMacro);
%			fprintf(fid,'\t%d\t%0.4f',[1:numMacro;ones(1,numMacro)./numMacro]);
%			fprintf(fid,'\n');
%		end,
%	end

%	fclose(fid);
%end

