function fns = GenLyapProb(A, M, LM, B, b, n, p, pre)

if nargin < 7
    error('error: Data related to problem ara in lack!')
elseif nargin == 7
    pre = 0;
end

    PI = perfectshuffle(p,p);
    fns.cost = @(x) f(A, M, LM, B, b, n, p, x);
    fns.grad = @(x) Grad(A, M, LM, B, b, n, p, x);
    fns.hess = @(x, v) HessEta(A, M, LM, B, b, n, p, x, v);
    fns.precon = @(x, v) preconditioner(x, v, A, M, LM, B, b, n, p, PI, pre);

end

%% Problem Related
% cost function
function [output, x] = f(A, M, LM, B, b, n, p, x)
    AV = A*x.V;
    MV = M*x.V;
    BV = b*(b'*x.V);
    VAV = x.V'*AV;
    VMV = x.V'*MV;
    
    output = trace(VAV*x.D*VMV*x.D) - trace(x.V'*BV*x.D);
end

% Riemannian gradient
function output = Grad(A, M, LM, B, b, n, p, x)
    AV = A*x.V;
    VAV = x.V'*AV;
    MV = M*x.V;
    VMV = x.V'*MV;
    BV = b*(b'*x.V);

    T = AV*x.D*VMV + MV*x.D*VAV - BV;
    output.S = sym(x.V'*T);
    output.Z = T-x.V*output.S;  
end

% Riemannian Hessian
function output = HessEta(A, M, LM, B, b, n, p, x, eta)
    AV = A*x.V;
    VAV = x.V'*AV;
    MV = M*x.V;
    VMV = x.V'*MV;
%     BV = b*(b'*x.V);

    ZMV = eta.Z'*MV;
    ZAV = eta.Z'*AV;
    T1 = AV*eta.S*VMV+MV*eta.S*VAV + ...
         AV*(ZMV) + MV*(ZAV) + ...
         (A*eta.Z)*VMV + (M*eta.Z)*VAV;
    T2 = AV*x.D*(ZMV') + MV*x.D*(ZAV') - b*(b'*eta.Z);
    output.S = sym(x.V'*T1);
    output.Z = T1 - x.V*(x.V'*T1) + (T2 - x.V*(x.V'*T2))/x.D;    
end

% preconditioning on A and M
function output = preconditioner(x, eta, A, MM, LM, B, b, n, p, PI, pre)
% Solving the following problem 
%      Approx(Hess(Y)[xi]) = eta for xi
% 
% preconditioner: integer
%   0   : no precontioner
%   1   : new preconditioer under the first metric
%   2   : preconditioner in Bart's paper but MM must be identity matrix


    if pre == 0     % no preconditioner
        output = eta;           

    elseif pre == 1 % new perconditioner under the first metric         
        dd = sqrt(diag(x.D));   %  dd = diag(R) = sqrt(sqrt(x.D))  
        AV = A*x.V;
        VAV = x.V'*AV;
        MV = MM*x.V;
        VMV = x.V'*MV;
        Y = x.V .* dd';         % note that A * diat(x) = A .* x', x a col vector
        if isempty(LM)   % O(p^3)
            YMY = dd .* (VMV .* dd');   
            L = chol(sym(YMY), 'lower');  % may be not numberically stable 
        else
            [~, L] = qr0(LM' * Y); % O(n^2p+np^2)
            L = L';
        end

        etamain = (0.5*x.V*eta.S + eta.Z) ./ dd';

%         LY = chol(YY, 'lower');  % may be not numberically stable
        YY = Y' * Y;    % O(n^2p)
        [~, LY] = qr0(Y);   % O(np^2)
        LY = LY';    

        LYAYL = L \ ((dd .* (VAV .* dd')) / L');  % O(np^2)
        LYAYL = sym(LYAYL);           % O(p^2)
        [Q, Lambda] = eigs(full(LYAYL), p);        % O(p^3)

        LQ = L' \ Q;          % O(p^2)
        Yeta = Y' * etamain;  % O(np^2)
        QLYA = (LQ)' * (AV .* dd')';            % O(np^2)     Q^T * L^{-1} * Y^T * A 
        [hatV, ~] = qr(full(MV .* dd'), 0);     % O(np^2)     hatV = orth(M * Y)
                 % O(np^2+np+p^3)
        tmp0 = etamain + Y * (LY' \ (LY \ (Yeta)));    % double, tmp0 = (I + P_Y) * eta                  
                 % O(np^2+np+p^3)
        tmp1 = mat2cell((tmp0 - hatV * (hatV' * tmp0)) * ((YY) * (LQ)), ...
                         [n], ones(1, p));             % cell,   tmp1 = (I - hatV * hatV^T) * (I + P_Y) * eta * (Y^T * Y) * L^{-T} * Q                       
                 % O(np)
        tmp2 = 2 * QLYA';                              % double, tmp2 = 2 * A * Y * L^{-T} * Q
                 % O(np^2+np)
        tmp2 = tmp2 - hatV * (hatV' * tmp2);           % double, tmp2 = (I - hatV * hatV^T) * 2 * A * Y * L^{-T} * Q         
                
        cLambda = num2cell(diag(Lambda)');   % cell         
                % O(p(np^c+np^2+np))        
        TMP = cellfun(@(x, y) subproblem1AM(A, MM, x, hatV, y, tmp2), cLambda, tmp1, 'UniformOutput', false);
                % O(np^2(p+1))
        U   = cell2mat(cellfun(@(x) QLYA * x(:,1), TMP, 'UniformOutput', false));
                % O(np^3(p+1))
        K   = cellfun(@(x, y) (2 * x) * speye(p) - QLYA * y(:, 2:end), cLambda, TMP, 'UniformOutput', false);            
                % O(p^2)
        tSxi = subproblem2(K, 2 * (LQ)' * (Yeta) * (YY) * (LQ), U, PI);
        tSxi = sym(tSxi);
        ctSxi = mat2cell(tSxi, [p], ones(1,p));
                % O()
        tZxi  = cell2mat(cellfun(@(x, y) (x(:, 1) - x(:, 2 : end) * y), TMP, ctSxi, 'UniformOutput', false));        
                % O(p^3)
        Sxi = LQ * tSxi * LQ';
                % O(np^3)
        Zxi = tZxi * LQ';
                % O(np^2)
        tmp = LY' \ (LY \ (Y' * Zxi));
                % O(np^2+np)
        Zxi1 = Y * sym(tmp);
                % O(np^2 + np)
        Zxi2 = Zxi - Y * tmp;
        
        tmp = (Y * Sxi + Zxi1 + Zxi2) .* dd'; 
        S = 2. * (x.V' * tmp);
        output.S  = sym(S); 
        output.Z = tmp - 0.5 * x.V * S;         
    elseif pre == 2   % preconditioner in Bart's paper; 
        AV = A*x.V;
        VAV = x.V'*AV;

        Seta = eta.S;
        Zeta = eta.Z;
        V = x.V;
%         VAV  = V' * (A * V);
        [Q, Lambda] = eig(full(sym(VAV)));        
        tV = V * Q;
        AtV = A * tV;
        tmp1 = AtV - tV * (tV' * AtV);
        tSeta = Q' * (Seta * Q);
        
        tZeta = mat2cell(Zeta * Q, [n], ones(1,p));
        cLambda = num2cell(diag(Lambda)');
        
        TMP = cellfun(@(x, y) subproblem1(A, x, tV, y, tmp1), cLambda, tZeta, 'UniformOutput', false);
        U   = cell2mat(cellfun(@(x) AtV' * x(:, 1), TMP, 'UniformOutput', false));
        K   = cellfun(@(x, y) x * speye(p) - AtV' * y(:, 2 : end), cLambda, TMP, 'UniformOutput', false);        

%         for i = 1 : p
%             tmp2(:,:,i) = subproblem1(A, Lambda(i, i), tV, [tZeta(:, i), AtV - tV * (tV' * AtV)]);
%             U(:, i) = AtV' * tmp2(:, 1, i);
%             K{i} = Lambda(i, i) * speye(p) - AtV' * tmp2(:, 2 : end, i);
%         end

        tSxi = subproblem2(K, tSeta, U, PI);
        tSxi = sym(tSxi);
        ctSxi = mat2cell(tSxi, [p], ones(1, p));
        tZxi  = cell2mat(cellfun(@(x, y) x(:, 1) - x(:, 2 : end) * y, TMP, ctSxi, 'UniformOutput', false));

        Sxi = Q * tSxi * Q';
        Zxi = tZxi * Q';       
        output.S = Sxi;
        output.Z = Zxi;        
    else
        error('error: this option does not exist for precond!')
    end
end

%% other useful tools
function X = sym(X)
    X = 0.5 * (X + X');
end

function [Q, R] = qr0(A)
% qr decomposition for A = Q * R, where Q^T * Q = I, R is upper triangular and
% its diagonal elements are not negative.
    [Q, R] = qr(A, 0);
    tmp = diag(R);
    tmp = sign(tmp);
    P = diag(tmp);
    R = P * R;
    Q = Q * P;
end

function output = invALtimesM(A, MM, lambda_i, B)  % O(np^c)
% global numSystems timeOfSystems
% numSystems = numSystems + size(B,2);
% tt = tic;
    ALtimesM = A + lambda_i * MM; 
    output = ALtimesM \ B;                     % ZWH, TODO: implement efficient computing inverse of a matrix
%     output = cholmod2(ALtimesM, B, -1);   % sometimes go wrong  
% timeOfSystems = timeOfSystems + toc(tt);
end


function output = subproblem1(A, lambda_i, tV, B1, B2)
% Solving the saddle-point problems
    n = size(tV,1);
    B = [B1, B2];
    tmp = invALtimesM(A, speye(n), lambda_i, [B, tV]);    
    Si = tV' * tmp(:, size(B, 2) + 1 : end);
    N = Si \ (tV' * tmp(:, 1 : size(B, 2)));
    output = tmp(:, 1 : size(B, 2)) - tmp(:, size(B, 2) + 1 : end) * N;
end

function output = subproblem1AM(A, MM, lambda_i, tV, B1, B2)  % O(np^c+np^2+np)
% Solving the saddle-point problems   
    B = [B1, B2];
    tmp = invALtimesM(A, MM, lambda_i, [B, tV]); % O(np^c)
    % (A + lambda * M)^{-1} * B
    AplamMB = tmp(:, 1 : size(B, 2)); 
    % (A + lambda * M)^{-1} * tV
    AplamMtV = tmp(:, size(B, 2) + 1 : end); 
    Si = tV' * AplamMtV;        % O(np^2)
    N = Si \ (tV' * AplamMB);          % O(np^2)
    output = AplamMB - AplamMtV * N;   % O(np^2+np)
end

function output = subproblem2(K, tSeta, U, PI)  % O(p^2)
% Solving the linear system:  (K + \Pi * K * \Pi) vec(\tilde{S}_\xi) = vec{R}    
    p = size(tSeta, 1);
    function output = LinearOper(S)
        S = reshape(S, p, p);
        tmp1 = zeros(p, p);
        tmp2 = zeros(p, p);
        for i = 1 : p
            tmp1(:, i) = K{i} * S(:, i);
            tmp2(i, :) = S(i, :) * K{i}';
        end
        output = reshape(tmp1 + tmp2, [], 1);
    end

    KK = sparse(p * p, p * p);
    for j = 1 : p
        KK((j - 1) * p + 1 : j * p, (j - 1) * p + 1 : j * p) = K{j};
    end

    L = ichol(sparse(KK + PI * KK * PI));
    b = tSeta - U - U';
    [output, ~] = pcg(@LinearOper, b(:), 1e-12, p * p, L, L');
    output = reshape(output, p, p);
    
%     b = tSeta - U - U';
%     output = reshape(myCG(@LinearOper, b(:), 1e-12, p * p), p, p);
    
end

% computing perfect shuffle matrix
function L = perfectshuffle(n, p)
    L = sparse(n * p, n * p);
    for h = 1 : p
       for k = 1 : n
           i = (k - 1) * p + h;
           j = (h - 1)* n + k;
           L(i, j) = 1;
       end
    end
end