/*
 * Copyright (C) 1996-2011 Roger Farmer, Daniel Waggoner and Tao Zha
 *
 * This free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * It is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * If you did not received a copy of the GNU General Public License
 * with this software, see <http://www.gnu.org/licenses/>.
 */

#include "dw_matrix.h"
#include "dw_matrix_array.h"
#include "dw_matrix_rand.h"
#include "bmatrix.h"
#include "dw_error.h"

#include <string.h>
#include <stdlib.h>

/*
 Computes MSV solution of 

   A(s(t) y(t) = B(s(t) y(t-1) + Psi(s(t)) epsilon(t) + Pi eta(t)

 using Newton's Method.  Assumes that Pi' = [zeros(s,n-s)  I(s)] and
 that the first n-s rows of A[i] are linearly independent.  is invertible.  P is the transition matrix and P(i,j) is the
 probability that s(t+1)=i given that s(t)=j.  Note that the columns of P
 must sum to one.  x is the initial value.  max_iteration is the maximum 
 number of iterations of Newton's method before failing and tolerance is 
 the convergence criterion.

 The solution is of the form

    y(t) = V{s(t)}*F1{s(t)}*y(t-1) + V{s(t)}*G1{s(t)}*epsilon(t)

  eta(t) = F2{s(t)}*x(t-1) + G2{s(t)}*epsilont(t)

 A positive value of err is the number of iterations needed to obtain 
 convergence and indicates success.  A negitive value of err is the number of 
 iterations before the method terminated without convergence and indicates
 failure.
*/
int fwz_msv_msre(TMatrix *F1, TMatrix *F2, TMatrix *G1, TMatrix *G2, TMatrix *V, TMatrix P, TMatrix *A, TMatrix *B, TMatrix *Psi, int s, TVector z, int max_iteration, PRECISION tolerance)
{
  int i, j, k, count, h, n, r, ns, verbose_errors, terminal_errors;
  TVector g;
  TMatrix *ZI, *IZ, **BU, **BU_row, **BU_col, *C1, *C2, *U, W0, W1, W2, W3, W4, W5, W6, D, Q, R, G;
  PRECISION norm;

  n=RowM(A[0]);
  r=Psi[0] ? ColM(Psi[0]) : 0;
  h=DimV(z)/(s*(n-s));
  ns=n-s;

  W0=CreateMatrix(s,ns);
  W1=CreateMatrix(ns,ns);
  W2=CreateMatrix(s*ns,s*ns);
  W3=CreateMatrix(s,s);
  W4=CreateMatrix(s*ns,s*ns);
  W5=CreateMatrix(s,n);
  W6=CreateMatrix(n,ns);

  g=CreateVector(h*s*ns);
  D=CreateMatrix(h*s*ns,h*s*ns);
  IZ=dw_CreateArray_matrix(h); 
  ZI=dw_CreateArray_matrix(h); 
  C1=dw_CreateArray_matrix(h);
  C2=dw_CreateArray_matrix(h);
  BU=dw_CreateRectangularArray_matrix(h,h);
  BU_row=dw_CreateRectangularArray_matrix(h,h);
  BU_col=dw_CreateRectangularArray_matrix(h,h);
  U=dw_CreateArray_matrix(h);

  Q=CreateMatrix(n,n);
  R=CreateMatrix(n,n);

  verbose_errors=dw_SetVerboseErrors(NO_ERR);
  terminal_errors=dw_SetTerminalErrors(NO_ERR);
  for (j=h-1; j >= 0; j--)
    {
      if (!(U[j]=Inverse_LU((TMatrix)NULL,A[j])))
	{
	  Transpose(R,A[j]);
	  QR(Q,R,R);
	  Transpose(R,R);
	  C1[j]=SubMatrix((TMatrix)NULL,R,ns,0,s,ns);
	  SubMatrix(W1,R,0,0,ns,ns);
	  if (!Inverse_LU(W1,W1))
	    {
	      dw_SetVerboseErrors(verbose_errors);
	      dw_SetTerminalErrors(terminal_errors);
	      printf("fwz_msv_msre():  First %d rows of A[%d] not linearly independent.\n",ns,j);
	      dw_exit(0);
	    }
	  ProductMM(C1[j],C1[j],W1);
	  C2[j]=SubMatrix((TMatrix)NULL,R,ns,ns,s,s);
	  InitializeMatrix(U[j]=CreateMatrix(n,n),0.0);
	  for (k=ns; k < n; k++) ElementM(U[j],k,k)=1.0;
	  InsertSubMatrix(U[j],W1,0,0,0,0,ns,ns);
	  ProductMM(U[j],Q,U[j]);
	}
      else
	C2[j]=IdentityMatrix((TMatrix)NULL,s);

      InitializeMatrix(IZ[j]=CreateMatrix(n,ns),0.0);
      for (k=n-s-1; k >= 0; k--) ElementM(IZ[j],k,k)=-1.0;

      InitializeMatrix(ZI[j]=CreateMatrix(s,n),0.0);
      for (k=s-1; k >= 0; k--) ElementM(ZI[j],k,ns+k)=1.0;

      for (i=h-1; i >= 0; i--)
	{
	  BU[i][j]=ProductMM((TMatrix)NULL,B[i],U[j]);
	  ProductMS(BU[i][j],BU[i][j],-ElementM(P,i,j));
	  BU_row[i][j]=SubMatrix((TMatrix)NULL,BU[i][j],0,0,ns,n);
	  BU_col[i][j]=SubMatrix((TMatrix)NULL,BU[i][j],0,ns,n,s);
	}
    }
  dw_SetVerboseErrors(verbose_errors);
  dw_SetTerminalErrors(terminal_errors);

  for (count=max_iteration; count > 0; count--)
    {
      // divides and reshapes z into h matrices
      for (k=h-1; k >= 0; k--)
	{
	  if (MajorForm(W0))
	    memcpy(pElementM(W0),pElementV(z)+k*s*(ns),s*(ns)*sizeof(PRECISION));
	  else
	    bTranspose(pElementM(W0),pElementV(z)+k*s*(ns),s,ns,1);

	  InsertSubMatrix(IZ[k],W0,ns,0,0,0,s,ns);
	  if (C1[k])
	    {
	      ProductMM(W0,C2[k],W0);
	      SubtractMM(W0,W0,C1[k]);
	    }
	  InsertSubMatrix(ZI[k],W0,0,0,0,0,s,ns);
	}

      for (j=h-1; j >= 0; j--)
	{
	  for (i=h-1; i >= 0; i--)
	    {
	      ProductMM(W1,BU_row[i][j],IZ[j]);
	      Transpose(W1,W1);
	      KroneckerProduct(W2,W1,C2[i]);
	      if (i == j)
		{
		  ProductMM(W3,ZI[0],BU_col[0][j]);
		  for (k=h-1; k > 0; k--) UpdateProductMM(1.0,W3,1.0,ZI[k],BU_col[k][j]);
		  for (k=ns-1; k >= 0; k--)
		    InsertSubMatrix(W4,W3,k*s,k*s,0,0,s,s);
		  AddMM(W2,W2,W4);
		}
	      InsertSubMatrix(D,W2,j*s*(n-s),i*s*(n-s),0,0,s*(n-s),s*(n-s));
	    }
	}

      for (j=h-1; j >= 0; j--)
	{
	  InitializeMatrix(W5,0.0);
	  for (i=h-1; i >= 0; i--) UpdateProductMM(1.0,W5,-ElementM(P,i,j),ZI[i],B[i]);
	  ProductMM(W6,U[j],IZ[j]);
	  ProductMM(W0,W5,W6);

	  // vectorize
	  if (MajorForm(W0))
	    memcpy(pElementV(g)+j*s*(n-s),pElementM(W0),s*(n-s)*sizeof(PRECISION));
	  else
	    bTranspose(pElementV(g)+j*s*(n-s),pElementM(W0),s,n-s,0);
	}
	 
      norm=Norm(g);  
      InverseProductMV(g,D,g);
      SubtractVV(z,z,g);    

      if (norm < tolerance) break;
    }

  // Convert solution
  InitializeMatrix(R,0.0);
  for (k=n-1; k >= 0; k--) ElementM(R,k,k)=1.0;
  G=CreateMatrix(n,r);
  for (k=h-1; k >= 0; k--)
    {
      if (MajorForm(W0))
	memcpy(pElementM(W0),pElementV(z)+k*s*(ns),s*(ns)*sizeof(PRECISION));
      else
	bTranspose(pElementM(W0),pElementV(z)+k*s*(ns),s,ns,1);
      InsertSubMatrix(IZ[k],W0,ns,0,0,0,s,ns);

      if (C1[k])
	{
	  ProductMM(W0,C2[k],W0);
	  SubtractMM(W0,W0,C1[k]);
	}
      InsertSubMatrix(R,W0,ns,0,0,0,s,ns);

      ProductMM(V[k],U[k],IZ[k]);
      MinusM(V[k],V[k]);

      ProductMM(Q,R,B[k]);
      SubMatrix(F1[k],Q,0,0,ns,n);
      SubMatrix(F2[k],Q,ns,0,s,n);
      ProductMM(G,R,Psi[k]);
      SubMatrix(G1[k],G,0,0,ns,r);
      SubMatrix(G2[k],G,ns,0,s,r);
    }

  // Cleanup
  FreeMatrix(G);
  FreeMatrix(Q);
  FreeMatrix(R);
  dw_FreeArray(U);
  dw_FreeArray(BU_col);
  dw_FreeArray(BU_row);
  dw_FreeArray(BU);
  dw_FreeArray(C2);
  dw_FreeArray(C1);
  dw_FreeArray(ZI);
  dw_FreeArray(IZ);
  FreeMatrix(D);
  FreeVector(g);
  FreeMatrix(W6);
  FreeMatrix(W5);
  FreeMatrix(W4);
  FreeMatrix(W3);
  FreeMatrix(W2);
  FreeMatrix(W1);
  FreeMatrix(W0);
 
  return (count) ? max_iteration-count+1 : -max_iteration;
}

/*
 Verifies that 

     x(t) = V(s(t))(F1(s(t))x(t-1) + G1(s(t))epsilon(t)
   eta(t) = F2(s(t))x(t-1) + G2(s(t))epsilon(t)

 is a solution of 
 
    A(s(t)) x(t) = B(s(t)) x(t-1) + Psi(s(t) epsilon(t) + Pi eta(t)

 where Pi' = [zeros(s,n-s)  eye(s)] by verifying that

  [A(j)*V(j)  Pi] [ F1(j)   
                    F2(j) ] = B(j)

  [A(j)*V(j)  Pi] [ G1(j)   
                    G2(j) ] = Psi(j)

  (P(i,1)*F2(1) + ... + P(i,h)*F2(h))*V(i) = 0
*/
PRECISION verify_solution(TMatrix P, TMatrix *A, TMatrix *B, TMatrix *Psi, int s, TMatrix *F1, TMatrix *F2, TMatrix *G1, TMatrix *G2, TMatrix *V)
{
  int i, j, n, r, h;
  PRECISION diff=0.0, tmp;
  TMatrix W, X, Y, Z, F, G;

  n=RowM(A[0]);
  r=Psi[0] ? ColM(Psi[0]) : 0;
  h=RowM(P);

  InitializeMatrix(X=CreateMatrix(n,n),0.0);
  for (i=n-s; i < n; i++) ElementM(X,i,i)=1.0;

  Y=CreateMatrix(n,n-s);
  F=CreateMatrix(n,n);
  G=CreateMatrix(n,r);
  Z=CreateMatrix(s,n);
  W=CreateMatrix(s,n-s);

  for (j=0; j < h; j++)
    {
      ProductMM(Y,A[j],V[j]);
      InitializeMatrix(X,0.0);
      for (i=n-s; i < n; i++) ElementM(X,i,i)=1.0;
      InsertSubMatrix(X,Y,0,0,0,0,n,n-s);
      Inverse_LU(X,X);
    
      InsertSubMatrix(F,F1[j],0,0,0,0,n-s,n);
      InsertSubMatrix(F,F2[j],n-s,0,0,0,s,n);
      UpdateProductMM(1.0,F,-1.0,X,B[j]);
      tmp=MatrixNormEuclidean(F);
      if (tmp > diff) diff=tmp;

      InsertSubMatrix(G,G1[j],0,0,0,0,n-s,r);
      InsertSubMatrix(G,G2[j],n-s,0,0,0,s,r);
      UpdateProductMM(1.0,G,-1.0,X,Psi[j]);
      tmp=MatrixNormEuclidean(G);
      if (tmp > diff) diff=tmp;
    }

  for (i=0; i < h; i++)
    {
      InitializeMatrix(Z,0.0);
      for (j=0; j < h; j++) UpdateM(1.0,Z,ElementM(P,j,i),F2[j]);
      ProductMM(W,Z,V[i]);
      tmp=MatrixNormEuclidean(W);
      if (tmp > diff) diff=tmp;
    }

  FreeMatrix(W);
  FreeMatrix(Z);
  FreeMatrix(G);
  FreeMatrix(F);
  FreeMatrix(Y);
  FreeMatrix(X);

  return diff;
}

/*******************************************************************************/
/******************************** Test MSV Code ********************************/
/*******************************************************************************/
//#define FWZ_Example
void TestMSVCode(void)
{
  TMatrix P, *A, *B, *Psi, *F1, *F2, *G1, *G2, *V;
  TVector x;
  int err, i, j;
  PRECISION error, scale;

#ifdef FWZ_Example
  int r=3, s=2, n=7, h=2;
  TVector tau, kappa, beta, gamma1, gamma2, rhoD, rhoS, rhoR, sigmaD, sigmaS, sigmaR;
#else
  int r=1, s=2, n=4, h=3;
#endif

  P=CreateMatrix(h,h);
  if (h == 1)
    ElementM(P,0,0)=1.0;
  else if (h == 2)
    {
      ElementM(P,0,0)=0.9872;  ElementM(P,0,1)=0.0;
      ElementM(P,1,0)=0.0128;  ElementM(P,1,1)=1.0;
    }
  else
    {
      dw_UniformMatrix(P);
      for (j=h-1; j >= 0; j--)
	{
	  scale=ElementM(P,0,j);
	  for (i=h-1; i > 0; i--) scale+=ElementM(P,i,j);
	  for (i=h-1; i >= 0; i--) ElementM(P,i,j)/=scale;
	}
    }

  A=dw_CreateArray_matrix(h);
  B=dw_CreateArray_matrix(h);
  Psi=dw_CreateArray_matrix(h);
  for (i=0; i < h; i++)
    {
      A[i]=CreateMatrix(n,n);
      B[i]=CreateMatrix(n,n);
      Psi[i]=CreateMatrix(n,r);
    }

  x=CreateVector(s*(n-s)*h);

  F1=dw_CreateArray_matrix(h);
  F2=dw_CreateArray_matrix(h);
  G1=dw_CreateArray_matrix(h);
  G2=dw_CreateArray_matrix(h);
  V=dw_CreateArray_matrix(h);
  for (i=0; i < h; i++)
    {
      F1[i]=CreateMatrix(n-s,n);
      F2[i]=CreateMatrix(s,n);
      G1[i]=CreateMatrix(n-s,r);
      G2[i]=CreateMatrix(s,r);
      V[i]=CreateMatrix(n,n-s);
    }

  while (1)
    {
#ifdef FWZ_Example
      tau=CreateVector(h);    ElementV(tau,0)=0.69;    ElementV(tau,1)=0.54;
      kappa=CreateVector(h);  ElementV(kappa,0)=0.77;  ElementV(kappa,1)=0.58;
      beta=CreateVector(h);   ElementV(beta,0)=0.997;  ElementV(beta,1)=0.993;
      gamma1=CreateVector(h); ElementV(gamma1,0)=0.77; ElementV(gamma1,1)=2.19;
      gamma2=CreateVector(h); ElementV(gamma2,0)=0.17; ElementV(gamma2,1)=0.30;
      rhoD=CreateVector(h);   ElementV(rhoD,0)=0.68;   ElementV(rhoD,1)=0.83;
      rhoS=CreateVector(h);   ElementV(rhoS,0)=0.82;   ElementV(rhoS,1)=0.85;
      rhoR=CreateVector(h);   ElementV(rhoR,0)=0.60;   ElementV(rhoR,1)=0.84;
      sigmaD=CreateVector(h); ElementV(sigmaD,0)=0.27; ElementV(sigmaD,1)=0.18;
      sigmaS=CreateVector(h); ElementV(sigmaS,0)=0.87; ElementV(sigmaS,1)=0.37;
      sigmaR=CreateVector(h); ElementV(sigmaR,0)=0.23; ElementV(sigmaR,1)=0.18;

      for (i=0; i < h; i++)
	{
	  InitializeMatrix(A[i],0.0);

	  ElementM(A[i],0,0)=1.0; ElementM(A[i],0,1)=-1.0; ElementM(A[i],0,3)=-ElementV(tau,i); ElementM(A[i],0,4)=ElementV(tau,i); ElementM(A[i],0,5)=-1.0; 
	  ElementM(A[i],1,0)=-ElementV(kappa,i); ElementM(A[i],1,2)=1.0; ElementM(A[i],1,3)=-ElementV(beta,i); ElementM(A[i],1,6)=-1.0;
	  ElementM(A[i],2,0)=-(1-ElementV(rhoR,i))*ElementV(gamma2,i); ElementM(A[i],2,2)=-(1-ElementV(rhoR,i))*ElementV(gamma1,i); ElementM(A[i],2,4)=1.0;
	  ElementM(A[i],3,5)=1.0;
	  ElementM(A[i],4,6)=1.0;
	  ElementM(A[i],5,0)=1.0;
	  ElementM(A[i],6,2)=1.0;
       
	  // Test A singular
	  if (i == 0) ElementM(A[i],6,2)=0.0;
       
	  InitializeMatrix(B[i],0.0);
	  ElementM(B[i],2,4)=ElementV(rhoR,i);
	  ElementM(B[i],3,5)=ElementV(rhoD,i);
	  ElementM(B[i],4,6)=ElementV(rhoS,i);
	  ElementM(B[i],5,1)=1.0;
	  ElementM(B[i],6,3)=1.0;
  
	  InitializeMatrix(Psi[i],0.0);
	  ElementM(Psi[i],2,2)=ElementV(sigmaR,i);
	  ElementM(Psi[i],3,0)=ElementV(sigmaD,i);
	  ElementM(Psi[i],4,1)=ElementV(sigmaS,i);
	}
#else
      for (i=0; i < h; i++)
	{
	  dw_NormalMatrix(A[i]);
	  dw_NormalMatrix(B[i]); ProductMS(B[i],B[i],0.001);
	  dw_NormalMatrix(Psi[i]);
	}
#endif

      dw_NormalVector(x);
      //InitializeVector(x,0.0);

      err=fwz_msv_msre(F1,F2,G1,G2,V,P,A,B,Psi,s,x,10000,1.0e-8);

      for (i=0; i < h; i++)
	{
	  printf("F1[%d]\n",i);
	  dw_PrintMatrix(stdout,F1[i],"%lf ");

	  printf("F2[%d]\n",i);
	  dw_PrintMatrix(stdout,F2[i],"%lf ");

	  printf("G1[%d]\n",i);
	  dw_PrintMatrix(stdout,G1[i],"%lf ");

	  printf("G2[%d]\n",i);
	  dw_PrintMatrix(stdout,G2[i],"%lf ");

	  printf("V[%d]\n",i);
	  dw_PrintMatrix(stdout,V[i],"%lf ");
	}

      error=verify_solution(P,A,B,Psi,s,F1,F2,G1,G2,V);
      printf("return code (should be positive): %d\n",err);
      printf("error (should be approximately zero): %le\n",error);
      //getchar();
    }
}

int main(void)
{
  TestMSVCode();
  return 0;
}

/*******************************************************************************/
/*******************************************************************************/
/*******************************************************************************/
