% Function to estimate regularized and unregularized coefficents.
% Output also includes the respective sample covariance matrix for the
% residuals.
%Inputs
%   X: Dependent variable matrix of dimensions N x T.
%   Z: Independent variable matrix of dimensions N x T.
%   RegA: A string with the regularization target for the coef.
%   RegSigma: A string with the regularization cov target.
%   F: A vector of length Na containing some values.
%   deltas: A vector of length Na.
%   Na: A scalar value representing the number of agents.
%   Ta: A scalar value representing the number of time steps.
%   realCof: structure with true parameters.
%   target: String for FEVD target.
%   H: A scalar value for the horizon of the FEVD.
function [Ahat, Areg, Sighat, SigReg] = RegularizeA(X,Z,RegA,RegSigma,Fa,deltasA,N,T)
% Calculate the unregularized version first
[Ahat, ~, Sighat] = VARols(X,Z,T);

% Calculate standardized OLS
Xa = (X-mean(X,2))./std(X')';
Za = (Z-mean(Z,2))./std(Z')';
AhatStd = VARols(Xa,Za,T);

%% Perform regularization for different RegA.
switch RegA
    
    case 'OLS'
        for l=1:1
            Areg(:,:,1,l)=Ahat;
        end

    case 'AENET'
        % Main adaptive elastic net specification using the standardized
        % approach
        weights = (abs(AhatStd(:))).^(-1);
        alphas2 = [0 logspace(-5,-.3,5) 0.75]; % Alpha grid for the simulation
        Fsmaller = 100; %number of lambda values to pick from.
        Areg = [];
        
        kk = 0;
        for alpha = alphas2
            kk=kk+1;
            opts = struct('nlambda',Fsmaller,'alpha', alpha,'intr',0,'penalty_factor',weights,'lambda_min',max(0.001,alpha*0.1),'standardize',0);
            AenetAss = glmnet(sparse(kron(Za',speye(N))), Xa(:),[],opts);
            if Fsmaller-size(AenetAss.beta,2)>0
                AenetAss.beta(:,size(AenetAss.beta,2)+1:Fsmaller) = repmat(AenetAss.beta(:,end),1,Fsmaller-size(AenetAss.beta,2));
            end
            for l=1:size(AenetAss.beta,2)
                AregGrid(:,:,1,l)=diag(std(X'))*reshape(AenetAss.beta(:,l),N,N)/diag(std(Z'));
            end
            Areg = cat(4,Areg,AregGrid);
        end
        
    case 'LASSO'
        opts=struct('nlambda',Fa,'alpha', 1,'intr',0,'lambda_min',max(0.001,0.1));
        AenetAss= glmnet(sparse(kron(Za',speye(N))), Xa(:),[],opts);
        if Fa-size(AenetAss.beta,2)>0
            AenetAss.beta(:,size(AenetAss.beta,2)+1:Fa)=repmat(AenetAss.beta(:,end),1,Fa-size(AenetAss.beta,2));
        end
        for l=1:size(AenetAss.beta,2)
            Areg(:,:,1,l)=diag(std(X'))*reshape(AenetAss.beta(:,l),N,N)/diag(std(Z'));
        end
        
    case 'RIDGE'
        opts=struct('nlambda',Fa,'alpha',0,'intr',0,'lambda_min',max(0.001,0*0.1));
        AenetAss= glmnet(sparse(kron(Za',speye(N))), Xa(:),[],opts);
        if Fa-size(AenetAss.beta,2)>0
            AenetAss.beta(:,size(AenetAss.beta,2)+1:Fa)=repmat(AenetAss.beta(:,end),1,Fa-size(AenetAss.beta,2));
        end
        for l=1:size(AenetAss.beta,2)
            Areg(:,:,1,l)=diag(std(X'))*reshape(AenetAss.beta(:,l),N,N)/diag(std(Z'));
        end
        
    case 'Variational Garrotte'
        tic
        AhatRegAssistant=VGnoCVversionRuben(sparse(kron(Z',speye(N)))',X(:)', 10);
        toc
        if Fa-size(AhatRegAssistant,2)>0
            AhatRegAssistant(:,size(AhatRegAssistant,2)+1:Fa)=repmat(AhatRegAssistant(:,end),1,Fa-size(AhatRegAssistant,2));
        end
        for ii=1:size(AhatRegAssistant,2)
            Areg(:,:,1,ii)=reshape(AhatRegAssistant(:,ii),N,N);
        end
        
    case 'Threshold'
        dX=X-mean(X,2); dZ=Z-mean(Z,2);
        aux=(1/(size(dX,2)-1))*[dX;dZ]*[dX' dZ'];
        for l=1:length(deltasA)
            aux2=adaptiveThresh([dX; dZ],aux,Rule,deltasA(l));
            Areg(:,:,1,l)=aux2(1:N,N+1:2*N)/aux2(N+1:2*N,N+1:2*N);
        end
        
    case 'Geometric'
        diffY=X-Z;
        opts=struct('nlambda',Fa,'alpha', 1,'intr',0);
        AenetAss= glmnet(sparse(kron(diffY',speye(N))), Z(:),[],opts);
        if Fa-size(AenetAss.beta,2)>0
            AenetAss.beta(:,size(AenetAss.beta,2)+1:Fa)=repmat(AenetAss.beta(:,end),1,Fa-size(AenetAss.beta,2));
        end
        for l=1:size(AenetAss.beta,2)
            Sp=-reshape(AenetAss.beta(:,l),N,N);
            if det(Sp)==0
                Areg(:,:,1,l)=ones(size(Sp));
            else
                Areg(:,:,1,l)= eye(size(Sp))-inv(Sp);
            end
        end
        
end

%% Perform Estimation on the covariance function of the residuals
switch RegSigma
    
    case 'Sample-Cov'
        for l=1:size(Areg,4)
            res=X-Areg(:,:,1,l)*Z;
            res=res-mean(res,2);
            SigReg(:,:,1,l)=(1/(size(res,2)-1))*(res*res');
        end
end

end