% ARMA(2,2) model
% x_n = a_1 x_{n-1} + a_2 x_{n-2} + b_0 w_n + b_1 w_{n-1}
% x_i = 0 of i<0
% w_n are iid Gaussian mean zero and std sigma

clear all

randn('state',0)
rand('twister',5489)


N=5;
sigma=1;

% an AR(2) model with roots 0.9 and 0.7 will
% have transfer function
%1 -a1*z^{-1} -a2*z^{-2} =(1 -0.9*z^{-1}) * (1 -0.7*z^{-1})%
%which implies a1 = 1.6, a2 = -0.63.
a1=1.6; a2=-0.63;
b0=5; b1=.3;

% generate the data 
xprev=[0 0];
w=[0 0];
burnin=5e3;
x=zeros(1,N+burnin);
for i=1:N+burnin
    w= [sigma*randn w(1)];
    x(i)=[a1 a2]*xprev' + [b0 b1]*w';
    xprev=[x(i) xprev(1)];
end
x=x(burnin+1:end); % discard burn in samples
 
% compute the likelihood

% we grid the roots of the AR(2) process to ensure stability
rut1=[0.8:0.025:0.99];
%rut1=[0.9:0.9]
rut2=[0.5:0.025: 0.9];


loglik=zeros(length(rut1),length(rut2));
alpha1=loglik;
alpha2=loglik;

for ii=1:length(rut1)
    for  jj=1:length(rut2);   
        a1=rut1(ii)+rut2(jj);
        a2=-1*rut1(ii)*rut2(jj);
        alpha1(ii,jj)=a1;
        alpha2(ii,jj)=a2;
        
        A=[a1 1; a2 0];
        B=[b0; b1];
        C=[1 0];
        I=eye(2);
        xhat=[0 0]';
        
        % solving P=A*P*A' + B*B'*sigma^2 for the true P
        T=[a1^2 a1 a1 1; a1*a2 0 a2 0; a1*a2 a2 0 0; a2^2 0 0 0];
        vecP=inv(eye(4)-T)*[b0^2 b0*b1 b0*b1 b1^2]';
        P=[vecP(1) vecP(2); vecP(3) vecP(4)];
        %P-A*P*A' - B*B'*sigma^2
        
        % uncomment for false P
       P=zeros(2,2);

        loglikTemp=0;
        for i=1:N
            xhatPred=A*xhat;
            Ppred=A*P*A' + B*B'*sigma^2;
            K=Ppred*C'/(C*Ppred*C');
            xhat=xhatPred+K*(x(i)-C*xhatPred);
            P=[I-K*C]*Ppred;
            loglikTemp=loglikTemp-0.5*log(C*Ppred*C')...
                -0.5*(x(i)-C*xhatPred)^2/(C*Ppred*C');
        end
        

        loglik(ii,jj)=loglikTemp;
    end
end

figure
surf(rut1,rut2,loglik'/N)
xlabel('root of 1-\alpha_1 z^{-1}')
ylabel('root of 1-\alpha_2 z^{-1}')
hold on
temp=axis; % temp=[XMIN XMAX YMIN YMAX ZMIN ZMAX]
plot3(.9,.7,temp(5),'*')
title(['Log likelihood (N=' num2str(N) ')'])
