function [theta,se,logl,res] = arma(X,AR,MA,C)
% estimate ARMA(p,q) model using conditional MLE
% [theta,se,logl,res] = armamle(X,AR,MA,C)
% Input:
%   X: observations
%   AR: AR specification. Example: if AR=[1;3]; then AR(1) and AR(3) are
%   included.
%   MA: MA specification.
%   C: Constant term. If there is a constant term, set C=1, otherwise C=0.
% Output:
%   theta: estimated parameters
%   se: standard error
%   logl: log likelihood 
%   res: residuals

if nargin<4
    C = 1;
end
p = length(AR);
pmax = max(AR);
q = length(MA);
qmax = max(MA);
m = p + q + C + 1; % number of parameters.
theta0 = [zeros(m-1,1);std(X)];
[theta,logl,exitflag,output,lambda,grad,hessian] = fmincon(@lharma,theta0,[],[],[],[],[],[],@stationary,[],X,AR,MA,C);
V = inv(hessian);
se = sqrt(diag(V));

% estimate residual
T = length(X);

if p>0
    alpha = theta(1:p);
    A = zeros(max(AR),1);
    A(AR) = alpha;
end
if q>0
    beta = theta(p+1:p+q);
    B = zeros(max(MA),1);
    B(MA) = beta;
end
if C
    c = theta(1+p+q);
else
    c = 0;
end
e = zeros(T,1);
mpq = max([qmax,pmax]);
if p > 0 && q > 0 
    for ii = mpq+1:T
        e(ii) = X(ii) - c - A'*X(ii-1:-1:ii-max(AR)) ...
        - B'*e(ii-1:-1:ii-max(MA));
    end
end
if p > 0 && q == 0
    for ii = mpq+1:T
        e(ii) = X(ii) - c - A'*X(ii-1:-1:ii-max(AR)); 
    end
end
if p==0 && q > 0
    for ii = mpq+1:T
        e(ii) = X(ii) - c - B'*e(ii-1:-1:ii-max(MA));
    end
end
    
res = e;
logl = -logl;

function y=lharma(theta0,X,AR,MA,C)
%likelihood function to estimate MA(q) model

p = length(AR);
q = length(MA);
pmax = max(AR);
qmax = max(MA);
m = p + q + C + 1; % number of parameters.

T = length(X);

A = zeros(pmax,1);
B = zeros(qmax,1);
if p>0
    alpha = theta0(1:p);
    A(AR) = alpha;
end
if q>0
    beta = theta0(p+1:p+q);
    B(MA) = beta;
end
if C
    c = theta0(1+p+q);
else
    c = 0;
end
sigma = theta0(m);

y = 0;
e = zeros(T,1);
mpq = max(pmax,qmax);
if p>0 && q>0
    for ii = mpq+1:T
        e(ii) = X(ii) - c - A'*X(ii-1:-1:ii-pmax) ...
            - B'*e(ii-1:-1:ii-qmax);
        y = y + e(ii)^2;
    end
end
if p>0 && q==0
    for ii = mpq+1:T
        e(ii) = X(ii) - c - A'*X(ii-1:-1:ii-pmax);
        y = y + e(ii)^2;
    end
end
if p==0 && q>0
    for ii = mpq+1:T
        e(ii) = X(ii) - c - B'*e(ii-1:-1:ii-qmax);
        y = y + e(ii)^2;
    end
end

y = y/(2*sigma^2);
y = y + log(sigma^2)*(T-mpq)/2 + log(2*pi)*(T-mpq)/2;

if isnan(y) || isinf(y) || ~isreal(y)
    y = 1e10;
end

function [C, Ceq] = stationary(theta0,X,AR,MA,C)
p = length(AR);
pmax = max(AR);
alpha = theta0(1:p);
A = zeros(pmax,1);
A(AR) = alpha;
z = roots([-A(pmax:-1:1);1]);
if isempty(z)
    C = 0;
else
    C = 1-min(abs(z));
end
Ceq = [];



