!**********************************************************************************
!
!	SUBROUTINE: Test
!
!	PURPOSE: Takes as input state variable and returns as output all variables
!                 for the current period.
!			 
!	AUTHOR:	 Ben Malin 
!
!	DATE:    11/2007
!**********************************************************************************

subroutine Test(state,control)
  
  ! inputs:
  !   state -- vector of dimension 2*N, contains state variables k and ln(a)
  !
  ! output:
  !   control -- vector of dimension 3*N, contains k', c, and y 
  !
  ! remarks:
  !
  
  use Params

  implicit none
  
  !Variable Declarations
  real(prec),dimension(2*N),intent(in) :: state
  real(prec),dimension(3*N),intent(out) :: control
  integer :: i,j	
  real(prec),dimension(5,2*N) :: Ttemp
 
  !Function used to calculate optimal policy functions (either consumption or capital) given state 
  real(prec),external :: ChebValue

  !Read in state variables
  do i = 1,N
     SimCapital(i) = state(i)
     shock(i) = state(N+i)
  end do
  call ChebPoly(d, state, Ttemp)

  !Compute consumption, output and tomorrow's cap stock for each country
  do i = 1,N
     SimCapTom(i) = ChebValue(Ttemp,CoeffsCap1(:,:,i), &
          & CoeffsCap2(:,:,:,:,i),CoeffsCap3(:,:,i),YCap0(i))
     SimConsumption(i) = ChebValue(Ttemp,CoeffsCons1(:,:,i), &
          & CoeffsCons2(:,:,:,:,i),CoeffsCons3(:,:,i),YCons0(i))
     SimOutput(i) = F(SimCapital(i),shock(i))
  enddo
  
  do i = 1,N
     control(i) = SimCapTom(i)
     control(N+i) = SimConsumption(i)
     control(2*N+i) = SimOutput(i)
  end do

end subroutine Test

!**********************************************************************************
!
!	SUBROUTINE:  ChebPoly
!
!	PURPOSE: Converts a vector of state variables from the [kmin,kmax] x N grid (or 
!                 [kmin,kmax]/[zmin,zmax] x d grid) to the [-1,1] x N (or d) grid 
!                and returns the Chebyshev polynomials for these state variables.
!                This subroutine works in conjunction with ChebValue.
!**********************************************************************************

subroutine ChebPoly(dim1,state,Ttemp)
  
  ! inputs:
  !    dim1 -- dimension of "state" vector
  !    state -- vector of capital stocks and, possibly, technology shocks
  !
  ! output:
  !    Ttemp -- Chebyshev polynomials evaluated at points in state vector
  ! 
  ! remarks:
  !

  use Params
  
  implicit none
  
  !Variable Declaration
  integer, intent(in) :: dim1
  real(prec), dimension(dim1), intent(in) :: state
  real(prec), dimension(2**(q-d)+1,dim1), intent(out) :: Ttemp
  real(prec), dimension(dim1) :: newstate
  integer :: i,j
  
  do i = 1,N
     !Convert capital stock [-1,1] grid
     newstate(i) = 2.0 * (state(i) - kmin)/(kmax-kmin) - 1.0	!kmax, kmin are global variables
     
!     if (newstate(i) < -1) then
!        print*, 'error: capital below kmin in function ChebPoly',i
!        pause
!     elseif (newstate(i) > 1) then
!        print*, 'error: capital above kmax in function ChebPoly',i
!        pause
!     end if
     
     !Construct Chebyshev Polynomials for capital stocks
     Ttemp(1,i) = 1
     Ttemp(2,i) = newstate(i)
     do j = 3,5			!Since q = d+2, 2**(q-d)+1 = 5.
        Ttemp(j,i) = 2.0 * newstate(i) * Ttemp(j-1,i) - Ttemp(j-2,i)
     end do
     
     !Convert technological shocks to [-1,1] grid
     if (dim1 == d) then
        newstate(i+N) = 2.0 * (state(i+N) - zmin)/(zmax-zmin) - 1.0	!zmax, zmin are global variables
        
!        if (newstate(i+N) < -1) then
!           print*, 'error: tech shock below zmin in function ChebPoly'
!           pause
!        elseif (newstate(i+N) > 1) then
!           print*, 'error: tech shock above zmax in function ChebPoly'
!           pause
!        end if
        
        !Construct Chebyshev Polynomials for tech shocks
        Ttemp(1,i+N) = 1
        Ttemp(2,i+N) = newstate(i+N)
        do j = 3,5	!Since q = d+2, 2**(q-d)+1 = 5.
           Ttemp(j,i+N) = 2.0 * newstate(i+N) * Ttemp(j-1,i+N) - Ttemp(j-2,i+N)
        end do
     endif
  end do
  
end subroutine ChebPoly

!**********************************************************************************
!
!	FUNCTION:  ChebValue
!
!	PURPOSE: Given the state variables, returns the value of the 
!                 policy function.
!**********************************************************************************

function ChebValue(Ttemp,tempCoeffs1,tempCoeffs2,tempCoeffs3,ChebValueOrigin)
  
  ! inputs:
  !    Ttemp -- Chebyshev polynomials of state variables
  !    tempCoeffs1 -- coefficients of policy function
  !    tempCoeffs2 -- coefficients of policy function
  !    tempCoeffs3 -- coefficients of policy function
  !    ChebValueOrigin -- coefficient of policy function
  !
  ! output:
  ! 
  ! remarks:
  !    The relevant information concerning the state is included in Ttemp.
  !    The formula for constructing the value is given in Malin/Krueger/Kubler (2007)
  !     equation (11) on pg. 13.
  !

  use Params
  
  implicit none
  
  !Variable Declaration
  real(prec),dimension(2**(q-d)+1,d),intent(in) :: Ttemp
  real(prec),dimension(d,2**(q-d-1)+1),intent(in) :: tempCoeffs1		
  real(prec),dimension(d-1,d,2**(q-d-1)+1,2**(q-d-1)+1),intent(in) :: tempCoeffs2		
  real(prec),dimension(d,2**(q-d)+1),intent(in) :: tempCoeffs3		
  real(prec),intent(in) :: ChebValueOrigin

  real(prec) :: ChebValue
  real(prec),dimension(d) :: ChebValue1, ChebValue3
  real(prec),dimension(d-1,d) :: ChebValue2
  integer :: i,j,i1,j1
  
  !Construct 4th order polynomials in each of d dimensions
  ChebValue3 = 0.0
  do i = 1,d
     do j = 1,5		!Note that 5 = 2**(q-d) + 1, since (q-d) = 2
        ChebValue3(i) = tempCoeffs3(i,j)*Ttemp(j,i) + ChebValue3(i) 
     end do
  end do
  
  !Construct 2nd order polynomials in each of d dimensions
  ChebValue1 = 0.0
  do i = 1,d
     do j = 1,3		!Note that 3 = 2**(q-d-1) + 1, since (q-d) = 2
        ChebValue1(i) = tempCoeffs1(i,j)*Ttemp(j,i) + ChebValue1(i) 
     end do
  end do
  
  !Construct 2-dimensional tensor products
  ChebValue2 = 0.0
  do i = 1,d-1
     do i1 = i+1,d
        do j = 1,3		!Note that 3 = 2**(q-d) + 1, since (q-d) = 2
           do j1 = 1,3		!Note that 3 = 2**(q-d) + 1, since (q-d) = 2
              ChebValue2(i,i1) = tempCoeffs2(i,i1,j,j1)*Ttemp(j,i)*Ttemp(j1,i1) + ChebValue2(i,i1)
           enddo
        enddo
     enddo
  enddo
  
  !Constuct ChebValue ... Using formula on pg. 13 of Malin/Krueger/Kubler
  ChebValue = 0.0
  do i = 1,d
     ChebValue = ChebValue3(i) + ChebValue
  enddo
  
  do i = 1,d
     ChebValue = (-1.0)*coc1*ChebValue1(i) + ChebValue
  enddo
  
  do i = 1,d-1
     do i1 = i+1,d
        ChebValue = ChebValue2(i,i1) + ChebValue
     enddo
  enddo
  
  if (d >= 3) then
     ChebValue = coc2*ChebValueOrigin + ChebValue
  endif

end function ChebValue
