!**********************************************************************************
!
!	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 4*N, contains k', c, l, and y 
  !
  ! remarks:
  !
  
  use Params

  implicit none
  
  !Variable Declarations
  real(prec),dimension(2*N),intent(in) :: state
  real(prec),dimension(4*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

  !Subroutines called by nonlinear equation solver
  external :: AConsLabOne
  external :: AConsLabOther

  !Variables used in nonlinear equation solver to solve for country 1's labor supply
  integer, parameter :: itmaxA = 300
  integer, parameter :: countA = 1
  real(prec) :: RelerrA, fnormA
  real(prec), dimension(countA) :: guessA, finalA, fvecA
  
  !Variables used in nonlinear equation solver to solve for cons and labor of countries 2-N
  integer, parameter :: itmaxB = 300
  integer, parameter :: countB = 2
  real(prec) :: RelerrB, fnormB
  real(prec), dimension(countB) :: guessB, finalB, fvecB

  ! Variables required by HYBRD
  real(prec), parameter :: epsfcn = 1E-16
  integer, dimension(countA) :: diagA 
  integer, dimension(countB) :: diagB
  integer, parameter :: mode = 1
  real(prec), parameter :: factor = 100.0
  integer, parameter :: nprint = 0
  integer :: info, nfev, iflag
  real(prec), dimension(countA,countA) :: fjacA
  real(prec), dimension(countB,countB) :: fjacB
  integer :: lrA, lrB
  real(prec), dimension(:), allocatable :: rA, rB
  real(prec), dimension(countA) :: qtfA, wa1A, wa2A, wa3A, wa4A
  real(prec), dimension(countB) :: qtfB, wa1B, wa2B, wa3B, wa4B

  !Initialize inputs to nle solver, DNEQNF
  RelerrA = 0.0000001
  fnormA = 0
  guessA = 1.0
  finalA = 1.0
  lrA = countA*(countA+1)/2  
  allocate(rA(lrA))

  !Initialize inputs to nle solver, DNEQNF
  RelerrB = 0.0000001
  fnormB = 0
  guessB(1) = A_tfp
  guessB(2) = 1.0
  !finalB = 1.0
  finalB(1) = A_tfp
  finalB(2) = 1.0
  lrB = countB*(countB+1)/2  
  allocate(rB(lrB))

  iflag = 0

  !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 labor supply and consumption of country 1
  SimConsumption(1) = ChebValue(Ttemp,CoeffsCons1(:,:,1), &
       & CoeffsCons2(:,:,:,:,1),CoeffsCons3(:,:,1),YCons0(1))
  !  guessA = lss
  !  call DNEQNF(AConsLabOne, RelerrA, countA, itmaxA, guessA, finalA, fnormA)
  finalA = lss
  call HYBRD(AConsLabOne,countA,finalA,fvecA,RelerrA,itmaxA,countA,countA, &
       & epsfcn,diagA,mode,factor,nprint,info,nfev,fjacA,countA,rA,lrA,qtfA, &
       & wa1A,wa2A,wa3A,wa4A)
  SimLabor(1) = finalA(1)
  SimOutput(1) = F(SimCapital(1),SimLabor(1),shock(1))
  
  do i = 1,N
     SimCapTom(i) = ChebValue(Ttemp,CoeffsCap1(:,:,i), &
          & CoeffsCap2(:,:,:,:,i),CoeffsCap3(:,:,i),YCap0(i))
  enddo
  !Compute consumption, labor supply and output of countries 2 - N.
  do FOCi = 2,N
!     guessB(1) = log(SimConsumption(1))
!     guessB(2) = SimLabor(1)
!     call DNEQNF(AConsLabOther, RelerrB, countB, itmaxB, guessB, finalB, fnormB)
     finalB(1) = log(SimConsumption(1))
     finalB(2) = SimLabor(1)
     call HYBRD(AConsLabOther,countB,finalB,fvecB,RelerrB,itmaxB,countB,countB, &
          & epsfcn,diagB,mode,factor,nprint,info,nfev,fjacB,countB,rB,lrB,qtfB, &
          & wa1B,wa2B,wa3B,wa4B)
     SimConsumption(FOCi) = exp(finalB(1))
     SimLabor(FOCi) = finalB(2)
     SimOutput(FOCi) = F(SimCapital(FOCi),SimLabor(FOCi),shock(FOCi))
  end do
    
  do i = 1,N
     control(i) = SimCapTom(i)
     control(N+i) = SimConsumption(i)
     control(2*N+i) = SimLabor(i)
     control(3*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

!*****************************************************************************
!
!	SUBROUTINE:  AConsLabOne
!
!	PURPOSE: Specifies equation whose solution is country-1 labor supply.
!*****************************************************************************

subroutine AConsLabOne(countA,ls,mrsmpl,iflag)
  
  ! inputs:
  !    ls -- Holds guess of nonlinear equation solver
  !    countA -- size of pol vector (1)
  !
  ! output:
  !    mrsmpl -- value of equation evaluated at ls
  ! 
  ! remarks:
  !    ls(1) will be the labor supply of country 1.
  !
  
  use Params
  
  implicit none
  
  !Variable Declaration
  integer,intent(in) :: countA, iflag	
  real(prec),dimension(countA),intent(in) :: ls	
  real(prec),dimension(countA),intent(out) :: mrsmpl
  
  real(prec) :: epsilon = 10.0**(-6)
  real(prec), dimension(countA) :: lstemp
  real(prec) :: BigNumb = 100
  
  !Make sure nle solver's guess for labor supply is in-bounds
  if (ls(1) < epsilon) then
     !print*, 'error: labor supply below epsilon in AConsLabOne'
     !pause
     lstemp(1) = epsilon    !Sets labor supply at the boundary
  elseif (ls(1) > Le) then
     !print*, 'error: labor supply above endowment in AConsLabOne '
     !pause
     lstemp(1) = Le - epsilon    !Sets labor supply near boundary
  else
     lstemp(1) = ls(1)
  endif
  

  mrsmpl(1) = UlUc(SimConsumption(1),lstemp(1)) + &
       & FL(SimCapital(1),lstemp(1),shock(1))

  !Penalty if guess for labor supply is out-of-bounds
  if (ls(1) < epsilon) then
     !print*, mrsmpl(1)
     mrsmpl(1) = mrsmpl(1) - (ls(1) - epsilon) * BigNumb
     !print*, mrsmpl(1), ls(1)
  elseif (ls(1) > Le) then
     !print*, mrsmpl(1)
     mrsmpl(1) = mrsmpl(1) - (ls(1) - Le + epsilon) * BigNumb
     !print*, mrsmpl(1), ls(1)
  endif
  
end subroutine AConsLabOne

!*****************************************************************************
!
!	SUBROUTINE:  AConsLabOther
!
!	PURPOSE: Specifies system of 2 equations whose solution is 
!               country-i consumption and labor supply.
!*****************************************************************************

subroutine AConsLabOther(countB,polcl,StatFOC,iflag)
  
  ! inputs:
  !    polcl -- Holds guess of nonlinear equation solver
  !    countB -- size of polcl vector (2)
  !
  ! output:
  !    StatFOC -- value of equations evaluated at polcl
  ! 
  ! remarks:
  !    polcl(1) will be consumption of country i
  !    polcl(2) will be the labor supply of country i
  !
  
  use Params
  
  implicit none
  
  !Variable Declaration
  integer,intent(in) :: countB, iflag    !countB == 2
  real(prec),dimension(countB),intent(in) :: polcl	
  real(prec),dimension(countB),intent(out) :: StatFOC
  
  real(prec) :: epsilon = 10.0**(-6)
  real(prec), dimension(countB) :: poltemp
  real(prec) :: BigNumb = 100

  !Make sure nle solver's guesses for labor supply and cons are in-bounds
  poltemp(1) = exp(polcl(1))
  if (polcl(2) < epsilon) then
     !print*, 'error: labor supply below epsilon in AConsLabOther'
     !pause
     poltemp(2) = epsilon    !Sets labor supply at the boundary
  elseif (polcl(2) > Le) then
     !print*, 'error: labor supply above endowment in AConsLabOther'
     !pause
     poltemp(2) = Le - epsilon   !Sets labor supply near boundary
  else
     poltemp(2) = polcl(2)
  endif
  
  !FOC for consumption smoothing
  StatFOC(1) = Pareto(1) * Uc(SimConsumption(1),SimLabor(1)) & 
       & - Pareto(FOCi) * Uc(poltemp(1),poltemp(2))
  
  !FOC for labor supply
  StatFOC(2) = UlUc(poltemp(1),poltemp(2)) & 
       & + FL(SimCapital(FOCi),poltemp(2),shock(FOCi))
  
  !Penalty if guess for labor supply is out-of-bounds
  if (polcl(2) < epsilon) then
     !print*, StatFOC(2)
     StatFOC(2) = StatFOC(2) - (polcl(2) - epsilon) * BigNumb
     !print*, StatFOC(2), polcl(2)
  elseif (polcl(2) > Le) then
     !print*, StatFOC(2)
     StatFOC(2) = StatFOC(2) - (polcl(2) - Le + epsilon) * BigNumb
     !print*, StatFOC(2), polcl(2)
  endif
  
end subroutine AConsLabOther
