function [F1,P] = dlqmarkov(beta,A1,A2,A3,B,Q,R,W,pi,S,ct,Theta,Psi)
%DLQMARKOV Linear quadratic regulator design for discrete-time systems
%          with some states governed by a Markov chain.
%          [F,P] = dlqmarkov(beta,A1,A2,A3,B,Q,R,W,pi,S,ct,Theta,Psi) 
%          calculates the decision functions (u[t]) that maximize:
%
%		J = Sum beta^t{x[t]'Qx[t] + u[t]'Ru[t] + 2x[t]'Wu[t]}
%
%          subject to {x3[t]} given and the constraint equations: 
%		
%	        x1[t+1] = A1*x1[t] + A2*x2[t] + B*u[t];
%               pi(i,j) = probability that x2[t+1] = S(:,j) given that
%                         x2[t] = S(:,i), i,j in {1,2,..,n}, all t;
%               element ct of x1 is equal to 1.
%      
%          The solution is found by imposing the following side constraints
%          (e.g. market clearing conditions) on first-order conditions
%
%               x3[t] = Theta*[x1[t];x2[t]] + Psi*u[t].
%
%          The solution is given by 
%               
%               u[t] = -Fj *x1[t]  if x2[t] = S(:,j) 
%         
%          where F = [F1; ... ;Fn].   The value function, given the state
%          vector is [x1;S(:,j)], is equal to x1'*Pj*x1 and P = [P1;... ;Pn]. 

%	   Ellen McGrattan, 2-4-93

A      = [A1,A2,A3];
[m1,n] = size(A1);
[n,m3] = size(A3);
[m2,n] = size(S);
m12    = m1+m2;
m      = m12+m3;
e      = zeros(m1,1);
e(ct)  = 1;
Wy = W(1:m12,:);
Wz = W(m12+1:m,:);
ay = [A1,A2]-B/R*Wy';
az = A3-B/R*Wz';
q  = Q - W/R*W';
qy = q(1:m12,1:m12);
qz = q(1:m12,m12+1:m);
s  = inv(eye(m3)+Psi/R*Wz');
t  = s*(Theta-Psi/R*Wy');
s  = s*Psi;

%  Compute matrices needed for Hamiltonian

ahat = ay + az*t;
bhat = B + az*s;
qhat = qy + qz*t;

P      = -kron(ones(n,1),eye(m1));
p      = kron(pi,eye(m1))*P;
F0     = [ ];
for j = 1:n;
  pj  = p((j-1)*m1+1:j*m1,:);
  vj  = [eye(m1);S(:,j)*e'];
  %F0  = [F0,(R+beta*B'*pj*bhat)\(beta*B'*pj*ahat*vj)];
  F0  = [F0,(R+beta*B'*pj*(B+A3*Psi)+Wz'*Psi)\ ...
       (beta*B'*pj*[A1,A2]+Wy'+beta*B'*pj*A3*Theta+Wz'*Theta)*vj];
end;
dd=1;
it=1;
maxit=5000;
tol=1e-5;       % reset tolerance if necessary 
while (dd> tol & it<= maxit);
  for j=1:n;
    pj  = p((j-1)*m1+1:j*m1,:);
    fj  = F0(:,(j-1)*m1+1:j*m1);
    vj  = [eye(m1);S(:,j)*e'];
    %vj  = [vj;t*vj-s*fj];
    vj  = [vj;Theta*vj-Psi*fj];
    %Pj  = vj'*Q*vj+fj'*R*fj+beta*(A*vj-B*fj)'*pj*(A*vj-B*fj);
    Pj  = vj'*Q*vj-vj'*W*fj-fj'*W'*vj+fj'*R*fj+ ...
          beta*(A*vj-B*fj)'*pj*(A*vj-B*fj);
    P((j-1)*m1+1:j*m1,:) = Pj;
  end;
  p     = kron(pi,eye(m1))*P;
  F1    = [ ];
  for j = 1:n;
    pj  = p((j-1)*m1+1:j*m1,:);
    vj  = [eye(m1);S(:,j)*e'];
    %F1  = [F1,(R+beta*B'*pj*bhat)\(beta*B'*pj*ahat*vj)];
    F1  = [F1,(R+beta*B'*pj*(B+A3*Psi)+Wz'*Psi)\ ...
       (beta*B'*pj*[A1,A2]+Wy'+beta*B'*pj*A3*Theta+Wz'*Theta)*vj];
  end;
  dd = norm(F1-F0);
  it = it+1;
  F0 = F1;
end;
if it>=maxit; 
  disp(sprintf('WARNING: Iteration limit of %g reached in DLQMARKOV',maxit))
end;
%for j = 1:n;
%  pj  = p((j-1)*m1+1:j*m1,:);
%  fj  = F1(:,(j-1)*m1+1:j*m1);
%  vj  = [eye(m1);S(:,j)*e'];
%  fj  = (R+Wz'*Psi)\(R*fj+(Wy'+Wz'*Theta)*vj);
%  F1(:,(j-1)*m1+1:j*m1) = fj;
%end;
