% Estimation function for the application
% selects and estimates a regularization method for the coefficent of the
% regression X = A*Z + eps
% selecs and estimates a regularization method for the covariance matrix of
% the residuals eps
%
% Input: 
% -X: dependet variable of size NxT
% -Z: independent variable of size N2xT
% -RegA: method for the coefficient estimation. Can take values OLS
% (fast) or elastic-net, which performs full CV 
% -RegSigma: method for covariance estimation. Can take values
% 'sample-cov', 'Ledoit-Wolf', 'Threshold', 'allCV', 'allCV_AIC'. The
% latter two perform cross validation and select the model either on the
% MSE or the AIC, respectively.
% -deltas: Threshold hyperparameter grid
% -Kfold: number of folds for Cross validation
% -globalRes: to be used if one wants to superimpose a residual series
% (needed for mixed static-dynamic estimation)
% -globaldfA: number of degrees of freedom needed for the coefficient of
% the residuals are super imposed
%
% Output:
% -Ahat, Areg: OLS and regularization estimate of the coefficient
% -Sighat, SigReg: Sample-Cov and regularized covariance estimate
% -RegInfo: structure containing info about the regression, such as degrees
% of freedom, partial R2, AICs, and CV losses


function [Ahat, Areg, Sighat, SigReg, RegInfo]=RegularizeApp(X,Z,RegA,RegSigma,deltas,Kfold,globalRes,globaldfA)
[N,T] = size(X);
[Ahat, ~, Sighat]=VARols(X,Z,T);

Xa = (X-mean(X,2))./std(X')'+mean(X,2);
Za = Z;
Za(2:end,:) = (Z(2:end,:)-mean(Z(2:end,:),2))./std(Z(2:end,:)')';
AhatStd = VARols(Xa,Za,T);

switch RegA
    % Estimation of the coefficient for the regression X = A*Z + eps
    case 'OLS'
        Areg = Ahat;
        [AICcCoef, partR2, partR2spillover, partR2OLS, partR2spilloverOLS, minMSE, dfAcoef] = deal(0);

    case 'elastic-net'
        tic
        nLambdas = 1000;
        a = 1/3;  % 1/3 is equal to 0.5 in the adaptive enet paper
        lambdas = fliplr(logspace(-3,1,nLambdas));
        
        baseCVols = KfoldCV(Xa,Za, Kfold,0,a);            % Calculates the base MSE loss
        MSE = KfoldCV(Xa,Za, Kfold,lambdas,a);            % Runs CV for a first enet guess
        [~,minMSEx] = min(MSE); lambdaMin = lambdas(minMSEx);
        if minMSEx == 1
            disp('LambdaMax')
        elseif minMSEx == nLambdas
            disp('LambdaMin')
        end
        
        % Value for initializiation of adapative elastic net        
        Aini(:,:,1) = AhatStd;                          % OLS weights
        Aini(:,:,2) = adaptiveElasticNet(Xa,Za,lambdaMin,a);    % enet weights
        Aini(:,:,3) = ones(size(AhatStd));              % uniform weights
        alphas = [0 logspace(-5,-.3,10) 0.66 0.75 0.9 1];   % alpha grid
        
        pp = plot(MSE);
        hold on
        sc = scatter(minMSEx,MSE(minMSEx));
        hold off
        
        for j = 1:3
            k = 0;
            AIni = Aini(:,:,j);
            AIni(:,1) = inf(size(AIni,1),1);
            weights = ( abs(AIni(:))).^(-1);
            
            for a = alphas
                k = k+1;
                lambdaUpper =  2.5*a +5*(1-a)- 3*(j==2);
                lambdaLower = -4*a -2.5*(1-a) - 3*(j==2);
                lambdas = fliplr( logspace(lambdaLower,lambdaUpper, nLambdas) );
                
                MSE2(:,k) = KfoldCV(Xa, Za, Kfold, lambdas, a, weights);
                [minMSE(k,j), I3] = min(MSE2(:,k));
                lambdaMin(k,j) = lambdas(I3);

                pp.YData = MSE2(:,k);
                ylim([minMSE(k,j) 0.5*(minMSE(k,j)+mean(MSE2(:,k)))])
                sc.XData = I3; sc.YData = minMSE(k,j);
                title(['Run j = ' num2str(j) ',alpha = ' num2str(a) ', lambdamin = ' num2str(lambdas(I3))])
                drawnow;
                if I3 == 1
                    disp(['LambdaMax for alpha=' num2str(a) ', lambda=' num2str(lambdaMin(k,j)) ' and j=' num2str(j)] )
                elseif I3 == nLambdas
                    disp(['LambdaMin for alpha=' num2str(a) ', lambda=' num2str(lambdaMin(k,j))  ' and j=' num2str(j)])
                end
            end
            
            [minminMSE(j),I(j)] = min(minMSE(:,j));
            aIni(j) = alphas(I(j));
        end
        
        [~,initial] = min(minminMSE);
%         aMin = aIni(initial);
        for j = 1:3
            AIni = Aini(:,:,j);
            AIni(:,1) = inf(size(AIni,1),1);
            weights = abs(AIni(:)).^(-1);
            
            [AregFinal(:,:,j)] = adaptiveElasticNet(X ,Z ,lambdaMin(I(j),j), aIni(j), weights);
            
            dfAcoef(j) = sum(AregFinal(:,:,j)~=0,[1 2]);
            res = X-AregFinal(:,:,j)*Z;
            res = res-mean(res,2);
            SigenetAssisted = (1/(T))*(res*res');
            LL = -N*T*0.5*log(2*pi) - T*0.5*log(det(SigenetAssisted)) -0.5*trace(res'/SigenetAssisted*res);
            K =  dfAcoef(j) + N*(N+1)/2 -N;
            penaltyTerm = 2*K*T*N/( T*N - K -N) ;
            
            penaltyTerm = 2*K ;
            
            AICcCoef(j) = -2*LL + penaltyTerm;
        end
        Areg = AregFinal(:,:,initial);
        minMSE = minMSE/baseCVols;
        disp(['Coefficient estimation took ' num2str(toc/60) ' minutes.'])
        
        % Calculate partial R2s
        % regularized
        resAux = X-Areg*Z;
        SSEfull = sum((resAux - mean(resAux,2)).^2,2);
        
        resAux = X-Areg(:,1)*Z(1,:) - Areg(:,end-N+1:end)*Z(end-N+1:end,:);
        SSEred = sum((resAux- mean(resAux,2)).^2,2);
        partR2 = (SSEred- SSEfull)./SSEred;
        
        resAux = X- Areg(:,1:end-N)*Z(1:end-N,:);
        SSEred2 = sum((resAux - mean(resAux,2)).^2,2);
        partR2spillover = (SSEred2- SSEfull)./SSEred2;   
         
        % OLS
        resAux = X-Ahat*Z;
        SSEfull = sum((resAux- mean(resAux,2)).^2,2);
        
        resAux = X-Ahat(:,1)*Z(1,:) - Ahat(:,end-N+1:end)*Z(end-N+1:end,:);
        SSEred = sum((resAux-mean(resAux,2)).^2,2);
        partR2OLS = (SSEred- SSEfull)./SSEred;
        
        resAux = X- Ahat(:,1:end-N)*Z(1:end-N,:);
        SSEred2 = sum((resAux-mean(resAux,2)).^2,2);
        partR2spilloverOLS = (SSEred2- SSEfull)./SSEred2;  
end

%% Regularization of Sigma
% Select between Ledoit-Wolf, Threshold, CV-w

switch RegSigma
    
    case 'Ledoit-Wolf'
        res=X-Areg*Z;
        res=res-mean(res,2);
        SigenetAssisted=(1/(T-1))*(res*res');
        SigReg=shrinkLedoit(SigenetAssisted,res);
        loglikelihood=0;
        
    case 'Threshold'
        res=X-Areg*Z;
        res=res-mean(res,2);
        SigenetAssisted=(1/(T-1))*(res*res');
        if length(deltas)>1
            Kfold=10;
            parfor f=1:length(deltas)
                loglikelihood(f)=KfoldCVsigma(res, Kfold,deltas(f),Rule,'Threshold');
            end
            [~,I]=min(real(loglikelihood));
            plot(real(loglikelihood))
            SigReg=adaptiveThresh(res,SigenetAssisted,Rule,deltas(I));
        else
            SigReg(:,:)=adaptiveThresh(res,SigenetAssisted,Rule,deltas);
        end
        
        
    case 'allCV'
        Methods={'Threshold'; ...
            'GLASSO'; 'Ledoit-Wolf'; 'Ledoit-Wolf'};
        etas = [0:0.05:2 2.05:0.25:4]; 
        
        res = X-Areg*Z;
        res = res-mean(res,2);
        SigenetAssisted = (1/(T))*(res*res');
        
        [~,alpha] = shrinkLedoit(SigenetAssisted,res);
        delta = [deltas;... Threshold
            logspace(0,3,length(deltas));... Glasso
            repmat(alpha,1,length(deltas));... Ledoit-Wolf
            0:1/(length(deltas)-1):1];%deltas];
        
        tic
        baseCV=KfoldCVsigma(res, Kfold,1,1,'Sample-Cov');
        for l=1:size(Methods,1)
            parfor f=1:size(delta,2)
                [loglikelihood(f,l), minEta(f,l)] = KfoldCVsigma(res, Kfold,delta(l,f),etas,Methods{l});
            end
        end
        [~,Itemp] = min(real(loglikelihood(:,1)));
        minEta(:,1) = minEta(Itemp,1);
        parfor f=1:size(delta,2)
            [loglikelihood(f,1), ~] = KfoldCVsigma(res, Kfold,delta(1,f),minEta(Itemp,1),Methods{1});
            %             minEta(f,1) = minEta(Itemp,1);
        end
        
        loglikelihood = loglikelihood/baseCV;
        plot(loglikelihood(1:end,:),'LineWidth',2)
        legend(Methods); ylim([min(min(loglikelihood)) loglikelihood(1,1)])
        
        [minlikelis,I] = min(real(loglikelihood));
        [~,I2] = min(minlikelis);
        %                 I2 = 2; %Temporary  selects GLASSO
        
        deltaOpt = delta(I2,I(I2));
        etaOpt = minEta(I(I2),I2);
        switch Methods{I2}
            case 'Threshold'
                SigReg = adaptiveThresh(res,SigenetAssisted,etaOpt,deltaOpt);
            case 'GLASSO'
                ThetaHat = (abs(inv(SigenetAssisted))).^(-0.75);
                deltaOpt = deltaOpt*ThetaHat/mean(mean(ThetaHat));
                SigReg = GraphicalLasso(res',deltaOpt);
            case 'Ledoit-Wolf'
                SigReg = shrinkLedoitmanual(SigenetAssisted,res,deltaOpt);
        end
        
        
        disp(['Estimation of Covariance matrix took ' num2str(toc/60) ' minutes.'])
        
    case 'allCV_AIC'
        Methods = {'Threshold'; ...
            'normalized GLASSO';...
            'Ledoit-Wolf'};
        etas = [1:0.05:2 2.05:0.15:3]; 
        if ~isempty(globalRes)
            res = globalRes;
        else
            res = X-Areg*Z;
        end
        res = res-mean(res,2);
        SigenetAssisted = (1/(T))*(res*res');
        
        [~,alpha] = shrinkLedoit(SigenetAssisted,res);
        delta = [deltas;... Threshold
            logspace(-1.5,-0.5,length(deltas));... normalized Glasso
            repmat(alpha,1,length(deltas))]; % Ledoit-Wolf
        
        tic
        baseCV = KfoldCVsigma(res, Kfold,1,1,'Sample-Cov');

        invSigIni = inv(SigenetAssisted); ones(N);
        
        for l=1:size(Methods,1)
            parfor f=1:size(delta,2)
                [loglikelihood(f,l), minEta(f,l)] = KfoldCVsigma(res, Kfold, delta(l,f), etas, Methods{l}, invSigIni);
            end
        end
        [~,Itemp] = min(real(loglikelihood(:,1)));
        minEta(:,1) = minEta(Itemp,1);
        parfor f=1:size(delta,2)
            [loglikelihood(f,1), ~] = KfoldCVsigma(res, Kfold,delta(1,f),minEta(Itemp,1),Methods{1});
        end
        
        
        [minlikelis,I] = min(real(loglikelihood));
        if I(1)==1
            disp(['min hyperparameter for ' Methods{1}])
        elseif I(2)==1
            disp(['min hyperparameter for ' Methods{2}])
        elseif I(1)==length(delta)
            disp(['max hyperparameter for ' Methods{1}])
        elseif I(2)==length(delta)
            disp(['max hyperparameter for ' Methods{2}])
        end
        

        for l=1:size(Methods,1)
            
            deltaOpt = delta(l,I(l));
            etaOpt = minEta(I(1),1);
            switch Methods{l}
                case 'Threshold'
                    SigReg(:,:,l) = adaptiveThresh(res,SigenetAssisted,etaOpt,deltaOpt);
                    dfSig(l) = sum(triu(SigReg(:,:,l))~=0,[1 2]);
                    
                case 'GLASSO'
                    
                    deltaOpt = deltaOpt*(ones(N)-eye(N));  % only penalizes off diagonals
                    [SigReg(:,:,l), invSig] = GraphicalLasso(res',deltaOpt);
                    
                    dfSig(l) = sum(triu(invSig)~=0,[1 2]);
                    
                case 'normalized GLASSO'
                    deltaOpt = deltaOpt*(ones(N)-eye(N));  % only penalizes off diagonals
                    
                    SaveStd = std(res');
                    resStd = SaveStd'.\res;
                    [SigRegA, invSig]=GraphicalLasso(resStd',deltaOpt);
                    SigReg(:,:,l) = diag(SaveStd)*SigRegA*diag(SaveStd);
                    dfSig(l) = sum(triu(invSig)~=0,[1 2]);

                case 'adaptive normalized GLASSO'
                    
                    invCorrIni = inv(corr(res'));
                    ThetaHat = abs(invCorrIni).^(-1);
                    deltaOpt = deltaOpt*ThetaHat/mean(mean(ThetaHat));           
                    SaveStd = std(res');
                    resStd = SaveStd'.\res;
                    [SigRegA, invSig]=GraphicalLasso(resStd',deltaOpt);
                    SigReg(:,:,l) = diag(SaveStd)*SigRegA*diag(SaveStd);
                    dfSig(l) = sum(triu(invSig)~=0,[1 2]);

                case 'adaptive GLASSO'

                    ThetaHat = abs(invSigIni).^(-1);
                    deltaOpt = deltaOpt*ThetaHat/mean(mean(ThetaHat));   
                    [SigReg(:,:,l), invSig] = GraphicalLasso(res',deltaOpt);
                    dfSig(l) = sum(triu(invSig)~=0,[1 2]);

                case 'Ledoit-Wolf'
                    SigReg(:,:,l) = shrinkLedoitmanual(SigenetAssisted,res,deltaOpt);
                    dfSig(l) = sum(triu(SigReg(:,:,l))~=0,[1 2]);

            end
            
            if ~isempty(globaldfA)
                dfA = globaldfA;
            else
                dfA = sum(Areg~=0,[1 2]);
            end
            
            LLreg(l) = -N*T*0.5*log(2*pi) - T*0.5*log(det(SigReg(:,:,l))) -0.5*trace(res'/SigReg(:,:,l)*res);
            K =  dfA + dfSig(l) -N;
            penaltyTerm = 2*K*T*N/( T*N - K -N) ;  % use this if AIC should use the corrected version
            penaltyTerm = 2*K ;

            AICcReg(l) = -2*LLreg(l) + penaltyTerm;% + 2*df + (2*df^2 + 2*df)/(numel(X) - df -1);
            
        end
        
        if any( isinf(AICcReg) ) % if AIC doesnt give a number, use CV instead
            [~, minAICindex] = min(AICcReg);
        else         % 
            [~,minAICindex] = min(minlikelis);
%             minAICindex = find(contains(Methods,'normalized GLASSO')); 
        end
        SigReg = SigReg(:,:,minAICindex);
        LLsample = -N*T*0.5*log(2*pi) - T*0.5*log(det(SigenetAssisted)) -0.5*trace(res'/SigenetAssisted*res);
        dfSig(l+1) = sum(triu(SigenetAssisted)~=0,[1 2]);
        K = dfSig(l+1) -N + dfA;
        penaltyTerm = 2*K*T*N/( T*N - K -N) ; % use this if AIC should use the corrected version
        penaltyTerm = 2*K;
        AICcReg(l+1) = -2*LLsample + penaltyTerm;
        
        disp(['  Estimation of Covariance matrix took ' num2str(toc/60) ' minutes.'])
        dfSig = dfSig(minAICindex);
        
        
    case 'Sample-Cov'
        res = X-Areg*Z;
        res = res-mean(res,2);
        SigReg(:,:,1,1) = (1/(T-1))*(res*res');
        [loglikelihood, dfSig, etaOpt, AICcReg ]= deal(0);
        
end

RegInfo = struct('AICcSig',AICcReg,'AICcCoef',AICcCoef,'CVlossCoef',minMSE,'CVlossSigma',loglikelihood,'dfA',dfAcoef,'dfSig',dfSig,...
    'eta',etaOpt,'partR2',partR2,'partR2spillover',partR2spillover,'partR2OLS',partR2OLS,'partR2spilloverOLS',partR2spilloverOLS);

end