% Used to make the plots for the arma/RLS example

% Example 4.2

function [] = optarma()

% Define the input to the plant
Ts = 0.01;  % sample period
t = 0:Ts:2; t=t(:);
u = 0.5*sin(2*pi*10*t) + cos(2*pi*5*t);
n = length(t);
y = zeros(n, 1);
k = 1:n; k=k(:);

% Define the plant output
for i=2:n;
  b(i) =  1 + 0.2*cos(2*pi*1*t(i));
  y(i) = 0.9*y(i-1) + b(i)*u(i-1);
end

plot(k,b);

% This is the loop to see how the forgetting factor affects the ID

lambdas = [0.2, 0.8, 0.99];
lstyle = [': '; '-.'; '--'];
hold on
m = length(lambdas);
for i=1:m
  [B,A,bcoef] = arma(y, u, [1,1], lambdas(i));
  plot(k, bcoef, lstyle(i, :));
end

axis([1 200 0.7 1.4])
xlabel('sample k');
legend('b(k)', '\lambda = 0.2', '\lambda = 0.8', '\lambda = 0.99');
hold off;



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% [B,A] = arma(y, u, m, lambda)
% [B,A,b] = arma(y, u, m, lambda)
% [B,A] = arma(y, u, m)
%
% Find the coefficients of the direct form II transposed discrete system
%
%         y(k) = b(1)*u(k-1) + ... + b(nb)*u(k-nb-1)
%              - a(1)*y(k-1) - ... - a(na)*y(k-na-1)
%
% where m=[na, nb].  A and B are normalized so that a(1) == 1.  This uses
% a recursive LS routine with forgetting factor lambda.  If lambda is not
% specified, a value of 0.99 is used.  If b is returned, the time history
% of b(1) is returned

function [B,A,b] = arma(y, u, m, lambda)

% Do some error checking
if (size(y,2) > 1)
  error('ARX only works for single output systems');
end
if (length(y) ~= length(u))
  error('ARX requires that the input and output be the same length');
end
if (length(m) ~= 2)
  error('ARX requires length(m) == 2');
end
if (nargin == 3)
  lambda = 0.99;
elseif (nargin == 4);
else
  error('incorrect number of inputs');
end

if (nargout == 3)
  b = zeros(length(y), 1);
end

na = m(1);
nb = m(2);
ncoef = na+nb;        % number of coefficients to ID
P = 2*eye(ncoef);     % Initialize P
th = zeros(ncoef, 1); % Initialize the coefficients
Pinv = inv(P);

% Initialize the data
ndata = length(y);
nmax = max(na, nb);

% run the recursive least squares routine
for i=(1+nmax):ndata
 
  % specify zeta
  ydata = -y(i-1:-1:i-na);
  udata = u(i-1:-1:i-nb);
  zeta = [ydata(:); udata(:)];

  % update the least squares estimate
  nrm = (lambda + zeta'*P*zeta);
  P = (P - P*zeta*zeta'*P/nrm)/lambda;
  th = th + P*zeta*(y(i) - zeta'*th);

  % return the leading b coef if requested
  if (nargout == 3)
    b(i) = th(na+1);
  end

end

if (na > 1)
  A = [1, th(1:na-1)'];
else
  A = 1;
end

if (nb > 0)
  B = th(na:na+nb-1)';
else
  B = [];
end