function [fail sols] = solve_stack(sols0, param, theta, final)

sols = sols0;
fail = 0;

M = param.M;
N = param.N;
h = param.h;
iave = param.iave;
Ctot = param.Ctot;
F = param.F;
Qn = param.Qn;
Qna = param.Qna;
lambda = param.lambda;
lfact = lambda*h*h;
% fact = 1/4/F/Qn/Ctot;

ivec = zeros(N,M);
diduvec = zeros(N,M);
didcvec = zeros(N,M);
Cave = zeros(N,M);
Us = sols.Us;
Cs = sols.Cs;

tol = 1e-10;
residual = 2*tol;
maxit = 20;
nit = 0;

% Order unknowns U then C
Ntot = 2*N*M;
res = zeros(Ntot,1);

fprintf('\n Computing theta = %d \n',theta)

while (residual > tol) && (fail == 0)
    J = sparse(Ntot,Ntot);
    % Find the local currents
    for m=1:M
        Cave(:,m) = (Cs(1:N,m)+Cs(2:(N+1),m))/2;
        for j=1:N;
            [I dIdU dIdC] = fni(Us(j,m),Cave(j,m),param);
            ivec(j,m) = I;
            diduvec(j,m) = dIdU;
            didcvec(j,m) = dIdC;
        end
    end
    
    % equations for U
    for m=1:m
        for j=1:N
            % special case j=1, m=1 - match average current
            if j==1 && m==1
                res(1) = h*sum(ivec(:,1)) - iave;
                for j=1:N
                    J(1,j) = h*diduvec(j,1);
                end
                J(1,N*M+1) = 0.5*h*didcvec(1,1);
                for j=2:N
                    J(1,N*M+j) = 0.5*h*(didcvec(j,1) + didcvec(j-1,1));
                end
            else
                % fundamental voltage equation
                indxu = (m-1)*N+j;
                indxc = indxu+N*M;
                % D2 U part 
                if j==1
                    res(indxu) = Us(2,m) - Us(1,m);
                    J(indxu,indxu) = -1;
                    J(indxu,indxu+1) = 1;
                elseif j==N
                    res(indxu) = Us(N-1,m) - Us(N,m);
                    J(indxu,indxu) = -1;
                    J(indxu,indxu-1) = 1;
                else
                    res(indxu) = Us(j-1,m) - 2*Us(j,m)+Us(j+1,m);
                    J(indxu,indxu) = -2;
                    J(indxu,indxu-1) = 1;
                    J(indxu,indxu+1) = 1;
                end
                % Now i part 
                if m == 1
                    res(indxu) = res(indxu) + 2*lfact*(ivec(j,1) - ivec(j,2));
                    J(indxu,indxu) = J(indxu,indxu) + 2*lfact*diduvec(j,1);
                    J(indxu,indxu+N) = J(indxu,indxu+N) - 2*lfact*diduvec(j,2);
                    if j > 1
                        J(indxu,indxc-1) = lfact*didcvec(j-1,1);
                        J(indxu,indxc-1+N) = -lfact*didcvec(j-1,2);
                    end
                    J(indxu,indxc) = lfact*didcvec(j,1);
                    J(indxu,indxc+N) = -lfact*didcvec(j,2);
                elseif m==M
                    res(indxu) = res(indxu) + lfact*(ivec(j,M) - ivec(j,M-1));
                    J(indxu,indxu) = J(indxu,indxu) + lfact*diduvec(j,M);
                    J(indxu,indxu-N) = J(indxu,indxu-N) - lfact*diduvec(j,M-1);
                    if j > 1
                        J(indxu,indxc-1) = lfact*didcvec(j-1,M)/2;
                        J(indxu,indxc-1-N) = -lfact*didcvec(j-1,M-1)/2;
                    end
                    J(indxu,indxc) = lfact*didcvec(j,M)/2;
                    J(indxu,indxc-N) = -lfact*didcvec(j,M-1)/2;
                else
                    res(indxu) = res(indxu) + lfact*(2*ivec(j,m) - ...
                        ivec(j,m-1)-ivec(j,m+1));
                    J(indxu,indxu) = J(indxu,indxu) + 2*lfact*diduvec(j,m);
                    J(indxu,indxu-N) = J(indxu,indxu-N) - lfact*diduvec(j,m-1);
                    J(indxu,indxu+N) = J(indxu,indxu+N) - lfact*diduvec(j,m+1);
                    if j > 1
                        J(indxu,indxc-1) = 2*lfact*didcvec(j-1,m)/2;
                        J(indxu,indxc-1-N) = -lfact*didcvec(j-1,m-1)/2;
                        J(indxu,indxc-1+N) = -lfact*didcvec(j-1,m+1)/2;
                    end
                    J(indxu,indxc) = 2*lfact*didcvec(j,m)/2;
                    J(indxu,indxc-N) = -lfact*didcvec(j,m-1)/2;
                    J(indxu,indxc+N) = -lfact*didcvec(j,m+1)/2;
                end
            end
        end
    end
                    
    
    % equations for C
    for m=1:M
        if m == 1
            Q = Qn + theta*(Qna-Qn);
        else
            Q = Qn;
        end
        
        fact = 1/4/F/Q/Ctot;

        for j=1:N
            indxu = (m-1)*N+j;
            indxc = indxu+N*M;
            res(indxc) = Cs(j+1,m)-Cs(j,m) + h*(Ctot-Cave(j,m))^2*ivec(j,m)*fact;
            J(indxc,indxu) = h*(Ctot-Cave(j,m))^2*diduvec(j,m)*fact;
            if j > 1
                J(indxc,indxc-1) = -1 + h*(Ctot-Cave(j,m))^2*didcvec(j,m)*fact/2 - ...
                    h*(Ctot-Cave(j,m))*ivec(j,m)*fact;
            end
            J(indxc,indxc) = 1 + h*(Ctot-Cave(j,m))^2*didcvec(j,m)*fact/2 - ...
                h*(Ctot-Cave(j,m))*ivec(j,m)*fact;
        end
    end
    
    % update U and C
    residual = max(abs(res));
    disp(sprintf('Residual: %d',residual))
    nit = nit+1;
    if nit >maxit
        fail = 1;
    end
    
    update = J\res;
    for m=1:M
        Us(:,m) = Us(:,m) - update((m-1)*N+1:m*N);
        Cs(2:(N+1),m) = Cs(2:(N+1),m) - update((m-1+M)*N+1:(m+M)*N);
    end
    
end

sols.Us = Us;
sols.Cs = Cs;

if final ==1
    for m=1:M
        Cave(:,m) = (Cs(1:N,m)+Cs(2:(N+1),m))/2;
        for j=1:N;
            [I dIdU dIdC] = fni(Us(j,m),Cave(j,m),param);
            ivec(j,m) = I;
        end
    end
    sols.is = ivec;
end
