function [A,b,Aeq,beq,idxbin,sol2part,sos] = util_formulate_graph_connect( adjmat, varargin )
% formulate graph connectivity constraint
% x=[x1,x2,...xn,  = node selection
%    y1,y2,...ym,  = edge selection, yk is for selection of tril(adjmat)(k)
%    f1,f2...f_2m, = flow on edge, (f_2k-1, f_2k) are for yk edge
%    g1,g2..gn] = flow source, gk connects to node xk
% Ax<=b, Aeq*x=beq, variables are binary
% idxbin = indices of variables that are binary, while other variables are
% continuous
% sos = special order constraint, following gurobi format, can be used with
% gurobi model. If this output is used, the constraints will not include
% symmetry breaking constraints
% g and f can take continous value, x and y are binary
% sol2part = a function that interprets the solution into a structure

pp = inputParser;

%number of variables before/after all variables, this is used when you are
%going to append/prepend this constraint to another problem
pp.addParamValue('nvar_before',0,@(x)x>=0); 
pp.addParamValue('nvar_after',0,@(x)x>=0);
pp.parse(varargin{:});

% nx = nnz(any(adjmat,1));
nx = size(adjmat,1);
ny = nnz(tril(adjmat));
nf = 2*ny;
ng = nx;

nvar = nx+ny+nf+ng;
idxbin = 1:nx+ny;

% function that converts x,y,f,g to variable index
x2varidx = @(x)x;
y2varidx = @(y)nx+y;
f2varidx = @(f)nx+ny+f;
g2varidx = @(g)nx+ny+nf+g;

% index table, convert from edge to node
% y2node(:,y) = [i,j] means edge y connects xi and xj
% this also defines a direction for each edge, then f=2*y-1 is flowing from
% i to j, f=2*y from j to i
% nodeidx = find(any(adjmat,1));
nodeidx = 1:nx;
[is,js] = find(tril(adjmat));
[~,loc1] = ismember(is,nodeidx);
[~,loc2] = ismember(js,nodeidx);
y2node = cat(1,loc1(:)',loc2(:)');

% node selection constraint: an edge is selectable only when both end are
% selected
% coefA_nodeselect = zeros(2*ny,nvar);
ijks = zeros(3,0);
for i=1:ny
    y = i;
    x1 = y2node(1,y);
    x2 = y2node(2,y);
    
    yy = y2varidx(y);
    xx1 = x2varidx(x1);
    xx2 = x2varidx(x2);
    
    ents = [2*i-1 yy 1; 2*i-1 xx1 -1; 2*i yy 1; 2*i xx2 -1]';
    ijks = cat(2,ijks,ents);
%     coefA_nodeselect(2*i-1,[yy,xx1]) = [1 -1];
%     coefA_nodeselect(2*i,[yy,xx2]) = [1 -1];
end
coefA_nodeselect = sparse(ijks(1,:),ijks(2,:),ijks(3,:),max(ijks(1,:)),nvar);
b_nodeselect = zeros(2*ny,1);

% flow capacity constraint
% coefA_flowcap = zeros(nf,nvar);
ijks = zeros(3,0);
for i=1:nf
    f = i;
    y = ceil(f/2);
    
    ff = f2varidx(f);
    yy = y2varidx(y);
    
    ents = [i,ff,1;i,yy,-2*nx]';
    ijks = cat(2,ijks,ents);
%     coefA_flowcap(i,[ff,yy]) = [1,-2*nx];
end
coefA_flowcap = sparse(ijks(1,:),ijks(2,:),ijks(3,:),nf,nvar);
b_flowcap = zeros(nf,1);

% flow balance constraint
% coefAeq_flowbalance = zeros(nx,nvar);
beq_flowbalance = zeros(nx,1);
ijks = zeros(3,0);
for i=1:nx
    g = i;
    x = i;
    
    yout = find(y2node(1,:)==x);
    yin = find(y2node(2,:)==x);
    
    fin = cat(2,2*yin-1,2*yout);
    fout = cat(2,2*yin,2*yout-1);
    
    gg = g2varidx(g);
    xx = x2varidx(x);
    ffin = f2varidx(fin);
    ffout = f2varidx(fout);
    
    ents = zeros(3,length(gg)+length(xx)+length(ffin)+length(ffout));
    ents(1,:) = i;
    ents(2,:) = [gg,xx,ffin,ffout];
    ents(3,:) = [1,-1,ones(1,length(ffin)),-ones(1,length(ffout))];
    ijks = cat(2,ijks,ents);
    
%     coefAeq_flowbalance(i,[gg,xx]) = [1,-1];
%     coefAeq_flowbalance(i,ffin) = 1;
%     coefAeq_flowbalance(i,ffout) = -1;
end
coefAeq_flowbalance = sparse(ijks(1,:),ijks(2,:),ijks(3,:),max(ijks(1,:)),nvar);

% source ordering constraint, to break symmetry of the problem
coefA_srcorder = zeros(0,nvar);
b_srcorder = [];
sos = [];
if nargout<7
    for i=1:nx
        g = i;
        gg = g2varidx(g);
        for j=1:i-1
            thisrow = zeros(1,nvar);
            x = j;
            xx = x2varidx(x);
            thisrow([gg,xx]) = [1 nx];
            coefA_srcorder(end+1,:) = thisrow;
            b_srcorder(end+1) = nx;
        end
    end
else
    % use SOS1 constraint instead
    sos.type = 1;
    sos.index = nx+ny+nf+(1:ng);
    sos.weights = ones(1,length(sos.index));
end
b_srcorder = b_srcorder(:);

%% range constraint, everything should be non-negative
coefA_range = -speye(nvar);
b_range = zeros(nvar,1);

%%

% output the model
A = cat(1,coefA_nodeselect,coefA_flowcap,coefA_srcorder,coefA_range);
b = cat(1,b_nodeselect,b_flowcap,b_srcorder,b_range);
Aeq = coefAeq_flowbalance;
beq = beq_flowbalance;

% prepend/append variables
nvar_before = pp.Results.nvar_before;
if isempty(nvar_before)
    nvar_before = 0;
end
nvar_after = pp.Results.nvar_after;
if isempty(nvar_after)
    nvar_after = 0;
end

nrow = size(A,1);
dummy_before = sparse(nrow,nvar_before);
dummy_after = sparse(nrow,nvar_after);
A = cat(2,dummy_before,A,dummy_after);

nrow = size(Aeq,1);
dummy_before = sparse(nrow,nvar_before);
dummy_after = sparse(nrow,nvar_after);
Aeq = cat(2,dummy_before,Aeq,dummy_after);

idxbin = idxbin + nvar_before;
if ~isempty(sos)
    sos.index = sos.index + nvar_before;
end

sol2part = @(x)explain_solution(x,adjmat,y2node,nvar_before,nvar_after);

end

function solpart = explain_solution(linx,adjmat,y2node,nvar_before,nvar_after)
linx = linx(nvar_before+1:end-nvar_after);
nx = size(adjmat,1);
ny = nnz(tril(adjmat));
nf = 2*ny;
ng = nx;

assert(length(linx) == nx+ny+nf+ng);
thres = 1e-6;

% node selection
x = linx(1:nx);
solpart.idxcompts = find(abs(x)>thres);

% edge selection
y = linx(nx+(1:ny));
idxedge = find(abs(y)>thres);
edgeidxmap = zeros(size(adjmat));
edgeidxmap(tril(adjmat)) = 1:ny;
newadjmat = ismember(edgeidxmap,idxedge);
newadjmat = newadjmat | newadjmat';
solpart.adjmat = newadjmat;

% flow assignment
f = linx(nx+ny+(1:nf));
flowmat = zeros(size(adjmat));
flowfrom = f(1:2:end);
flowto = f(2:2:end);
idxfrom = sub2ind(size(adjmat),y2node(1,:),y2node(2,:));
idxto = sub2ind(size(adjmat),y2node(2,:),y2node(1,:));
flowmat(idxfrom) = flowfrom;
flowmat(idxto) = flowto;
solpart.flowmat = flowmat;

% flow source
g = linx(nx+ny+nf+(1:ng));
idxsrc = find(abs(g)>thres);
solpart.idxcompts_flowsource = idxsrc;

end

function adjmat = sol_to_adjmat(sol,adjmat)
ny = nnz(tril(adjmat));
nx = size(adjmat,1);
idx = find(tril(adjmat));

ys = abs(sol(nx+1:nx+ny))>1e-6;
idxuse = idx(ys);
adjmat = false(size(adjmat));
adjmat(idxuse) = 1;
adjmat = adjmat | adjmat';
end


