% Function to estimate regularized and unregularized covariance matrix.
% Output also includes the coefficient matrix used/estimated
%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.
%   Fa: A vector of length Na containing some values.
%   deltaSig: Hyperparameters for the regularisation.
%   N: A scalar value representing the number of agents.
%   T: 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.
%   EstLast: if provided, the last estimation of the coefficient matrix
%           (speeds up the code)
function [Ahat, Areg, Sighat, SigReg]=RegularizeSig(X,Z,RegSigma,deltaSig,N,T,RuleA,real,Estlast)
[Ahat, ~, Sighat]=VARols(X,Z,T);


%% Coefficient estimation/regularisation with AENET only

Xa=(X-mean(X,2))./std(X')';
Za=(Z-mean(Z,2))./std(Z')';
AhatStd=VARols(Xa,Za,T);
if nargin > 9
    Areg = Estlast;
else
    
    weights = (abs(AhatStd(:))).^(-1);
    kk=0;
    alphas2 = [0 logspace(-5,-.3,5) 0.75];
    Fsmaller=100;
    Areg=[];
    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
    
    if (nargin > 7)
        for l = 1:size(Areg,4)
            testnorm(l)=norm(real-Areg(:,:,1,l),'fro');
        end
        [~, l]=min(testnorm);
        Areg=squeeze(Areg(:,:,1,l));
    end
    
end
%% Covariance regularisation for different targets

switch RegSigma
    
    case 'Ledoit-Wolf'
        res=X-Areg*Z;
        res=res-mean(res,2);
        SigenetAssisted=(1/(size(res,2)-1))*(res*res');
        SigReg(:,:,1,1,1)=shrinkLedoit(SigenetAssisted,res);
        
    case 'Threshold'
        res=X-Areg*Z;
        res=res-mean(res,2);
        SigenetAssisted=(1/(size(res,2)-1))*(res*res');
        deltaSig = logspace(-1,1,length(deltaSig));
        switch RuleA
            case 'multiple Eta'
                Eta = [1:0.25:8];
            case 'soft-threshold'
                Eta = 1;
        end
        k = 0;
        for f=1:length(deltaSig)
            for g = 1:length(Eta)
                k = k+1;
                SigReg(:,:,1,1,k) = adaptiveThresh(res,SigenetAssisted,Eta(g),deltaSig(f));
            end
        end
        
        
    case 'Sample-Cov'
        res=X-Areg*Z;
        res=res-mean(res,2);
        SigReg=(1/(size(res,2)-1))*(res*res');
        
    case 'GLASSO'
        res=X-Areg*Z;
        res=res-mean(res,2);
        
        SaveStd=std(res');
        res=SaveStd'.\res;
        
        alpha = N/T;
        minGrid = -15*(1-alpha)+1*alpha;
        maxGrid = -3*(1-alpha)+5*alpha;
        
        deltaSig = [0 logspace(minGrid,maxGrid,99)] ;

        w = cov(res');
        theta =inv(w);
        for f=1:length(deltaSig)
            [w, theta] = glasso(N, cov(res'), 0, deltaSig(f).*ones(N), 0, 1, 0, 0, 1e-4, 1e4, w, theta);
            SigReg(:,:,1,1,f) =  diag(SaveStd)*w*diag(SaveStd);
        end
        

end

end