function [fST,ST,ST2,varargout] = supgMethod(X,T,av,nu,sigma,s,tau,dt,referenceElement,type,varargin)
% [fST,ST,ST2,varargout] = supgMethod(X,T,av,nu,sigma,s,tau,dt,referenceElement,type,varargin)
% Creates stabilization terms for SUPG method.
% Takes as inputs system problems and tau term.
% varagin are the values of W and w when implicit scheme.
% Outputs fST vector of stabilization to be added to f and ST and ST2 matrix of
% stabilization terms to be added at K. varargout ST3 for the explicit scheme.

% Setting variable arguments
if type == 2 || type == 3
    W = varargin{1};
    w = varargin{2};
end

% Number of nodes and elements
nPt = size(X,1); 
nElem = size(T,1); 

if type == 3
    ST  = zeros(2*nPt);
    ST2 = zeros(2*nPt);
    fST = zeros(2*nPt,1);
else
    ST  = zeros(nPt);
    ST2 = zeros(nPt);
    ST3 = zeros(nPt);
    fST = zeros(nPt,1);
end

% Loop on elements
for ielem=1:nElem
    Te = T(ielem,:); 
    Xe = X(Te,:); 
    ae = av(Te,:);
    se = s(Te,:);
    if type == 1
        [fSTe,STe,STe2,STe3] = EleMat(Xe,ae,nu,sigma,se,tau,dt,referenceElement,type);
        ST3(Te,Te)           = ST3(Te,Te) + STe3;
    elseif type == 2 || type == 3
        [fSTe,STe,STe2]      = EleMat(Xe,ae,nu,sigma,se,tau,dt,referenceElement,type,W,w);
    end
    
    % Assembly
    if type == 1 || type == 2
        ST(Te,Te)  = ST(Te,Te)  + STe;
        ST2(Te,Te) = ST2(Te,Te) + STe2;
        fST(Te)    = fST(Te)    + fSTe;
    elseif type == 3
        Te = [T(ielem,:) T(ielem,:)+nPt];
        ST(Te,Te)  = ST(Te,Te)  + STe;
        ST2(Te,Te) = ST2(Te,Te) + STe2;
        fST(Te)    = fST(Te)    + fSTe;
    end
end

if type == 1;   varargout{1} = ST3;     end

function [fSTe,STe,STe2,varargout] = EleMat(Xe,ae,nu,sigma,se,tau,dt,referenceElement,type,varargin)

if type == 2 || type == 3
    W = varargin{1};
    w = varargin{2};
end

% reference element information
nen = referenceElement.nen; 
ngaus = referenceElement.ngaus; 
wgp = referenceElement.GaussWeights; 
N = referenceElement.N; 
Nxi = referenceElement.Nxi; 
Neta = referenceElement.Neta; 
N2xi = referenceElement.N2xi; 
N2eta = referenceElement.N2eta;
N2xieta = referenceElement.N2xieta;
N2etaxi = referenceElement.N2etaxi;

if type == 3
    STe  = zeros(2*nen);
    STe2 = zeros(2*nen);
    fSTe = zeros(2*nen,1);
else
    STe  = zeros(nen);
    STe2 = zeros(nen);
    STe3 = zeros(nen);
    fSTe = zeros(nen,1);
end

% Loop on Gauss points 
for ig = 1:ngaus
    N_ig = N(ig,:);
    Nxi_ig = Nxi(ig,:);
    Neta_ig = Neta(ig,:);
    N2xi_ig = N2xi(ig,:);
    N2eta_ig = N2eta(ig,:);
    N2xieta_ig = N2xieta(ig,:);
    N2etaxi_ig = N2etaxi(ig,:);
    
    J     = [Nxi_ig*Xe(:,1)  Nxi_ig*(Xe(:,2))
             Neta_ig*Xe(:,1) Neta_ig*(Xe(:,2))];
         
    Jacob = [(Nxi_ig*Xe(:,1))^2 (Nxi_ig*Xe(:,2))*(Nxi_ig*Xe(:,1)) N2xi_ig*Xe(:,1) (Nxi_ig*Xe(:,1))*(Nxi_ig*Xe(:,2)) (Nxi_ig*Xe(:,2))^2 N2xi_ig*Xe(:,2);
             (Neta_ig*Xe(:,1))*(Nxi_ig*Xe(:,1)) (Neta_ig*Xe(:,2))*(Nxi_ig*Xe(:,1)) N2xieta_ig*Xe(:,1) (Neta_ig*Xe(:,1))*(Nxi_ig*Xe(:,2)) (Neta_ig*Xe(:,2))*(Nxi_ig*Xe(:,2)) N2xieta_ig*Xe(:,2)
             0 0 Nxi_ig*Xe(:,1) 0 0 Nxi_ig*(Xe(:,2))
             (Nxi_ig*Xe(:,1))*(Neta_ig*Xe(:,1)) (Nxi_ig*Xe(:,2))*(Neta_ig*Xe(:,1)) N2etaxi_ig*Xe(:,1) (Nxi_ig*Xe(:,1))*(Neta_ig*Xe(:,2)) (Nxi_ig*Xe(:,2))*(Neta_ig*Xe(:,2)) N2etaxi_ig*Xe(:,2)
             (Neta_ig*Xe(:,1))^2 (Neta_ig*Xe(:,2))*(Neta_ig*Xe(:,1)) N2eta_ig*Xe(:,1) (Neta_ig*Xe(:,1))*(Neta_ig*Xe(:,2)) (Neta_ig*Xe(:,2))^2 N2eta_ig*Xe(:,2);
             0 0 Neta_ig*Xe(:,1) 0 0 Neta_ig*(Xe(:,2))];
    
    dvolu = wgp(ig)*det(J);
    res   = Jacob\[N2xi_ig;N2xieta_ig;Nxi_ig;N2etaxi_ig;N2eta_ig;Neta_ig];
    d2N   = res(1,:)+res(5,:);
    a     = N_ig*ae; 
    a_dN  = a(1)*res(3,:) + a(2)*res(6,:);
    L     = a_dN-nu*d2N+sigma*N_ig;
    
	% Choose according to time scheme.
    if type == 1
        STe  = STe  + tau*(a_dN)'*N_ig*dvolu;
        STe2 = STe  + tau*(a_dN)'*N_ig*dvolu;
        STe3 = STe  + dt*tau*(a_dN)'*L*dvolu;
        fSTe = fSTe + dt*tau*(a_dN)'*mean(se)*dvolu;
    elseif type == 2
        STe  = STe  + tau*(W*a_dN)'*(1/dt*N_ig+W*L)*dvolu;
        STe2 = STe2 + tau*(W*a_dN)'*(w*L)*dvolu;
        fSTe = fSTe + tau*(W*a_dN)'*(w*mean(se))*dvolu;
    elseif type == 3
        tW  = tau*W;
        tWW = tau*W*W;
        tWw = tau*W*w;
        
        STe  = STe  + [tW(1,1)*(a_dN)'*(1/dt*N_ig) + tWW(1,1)*(a_dN)'*L ...
                       tW(1,2)*(a_dN)'*(1/dt*N_ig) + tWW(1,2)*(a_dN)'*L;
                       tW(2,1)*(a_dN)'*(1/dt*N_ig) + tWW(2,1)*(a_dN)'*L ...
                       tW(2,2)*(a_dN)'*(1/dt*N_ig) + tWW(2,2)*(a_dN)'*L]*dvolu;
        STe2 = STe2 + [tWw(1)*(a_dN)'*(a_dN-nu*d2N+sigma*N_ig) zeros(nen);
                       zeros(nen) tWw(2)*(a_dN)'*(a_dN-nu*d2N+sigma*N_ig)]*dvolu;
        fSTe = fSTe + [tWw(1)*(a_dN)'*mean(se);
                       tWw(2)*(a_dN)'*mean(se)]*dvolu;
    end
end

if type == 1;   varargout{1} = STe3;   end
