!************************************************************************
!
!    Subroutine:  Policy
!
!    Purpose:    Solves for equilibrium using policy function iteration
!
!    Author:   Ben Malin
!
!    Date:   6/04
!    Revised: 9/06/07
!************************************************************************

subroutine Policy

 ! inputs:
 !
 ! output:
 ! 
 ! remarks:
  
  use Params		
  use Regression	!Holds subroutine ChebyshevRegression
  use EulerEquation	!Holds subroutine EulerEq 
  
  implicit none
  
  !Declaration of Variables
  real(prec) :: reldifcons, reldifcap, reldiflab
  real(prec) :: reldifcons0, reldifcons2, reldifcons3, reldifcap0, reldifcap2, reldifcap3
  real(prec) :: reldiflab0, reldiflab2, reldiflab3
  integer :: j,i,i1,h,h1	!Counters used in do loops
  
  !Holds previous iteration of policy function
  real(prec),dimension(d,2**(q-d),N) :: OldYCons3	!Value of policy functions (consumption) at points3 
  real(prec),dimension(d,2**(q-d),N) :: OldYCap3	!Value of policy functions (capital) 
  real(prec),dimension(d,2**(q-d),N) :: OldYLab3	!Value of policy functions (labor) 
  real(prec),dimension(d-1,d,2**(q-d-1),2**(q-d-1),N) :: OldYCons2 !Value of policy function (cons) at points2 
  real(prec),dimension(d-1,d,2**(q-d-1),2**(q-d-1),N) :: OldYCap2  !Value of policy function (cap) at points2 
  real(prec),dimension(d-1,d,2**(q-d-1),2**(q-d-1),N) :: OldYLab2  !Value of policy function (lab) at points2 
  
  real(prec),dimension(2**(q-d)+1,d) :: Ttemp
  
  !Variables used in DNEQNF function - in solving Euler Equations to determine policy function
  integer,parameter :: itmax = 200
  integer,parameter :: count = Nplus
  real(prec) :: Relerr,fnorm
  real(prec),dimension(count) :: guess, final
  
  ! Initialize inputs to nonlinear equation solver, DNEQNF
  Relerr = 0.0001
  fnorm = 0
  guess = 0.1
  final = 0.1
  
  !Start timer
  call CPU_TIME (time_begin)
  
  !Policy Function Iteration
  print*, 'Here we go'

  do iter = 1, maxiter
     print*, iter
!if (iter > 50) then 
!   pause
!endif
    
     !Stores policy function from previous iteration
     !Uses Acceleration methods
     if (iter < 5) then 
        !At Origin
        OldYCap0 = yCap0
        OldYCons0 = yCons0  
        OldYLab0 = yLab0
         
        !At points2
        OldYCap2 = yCap2
        OldYCons2 = yCons2
        OldYLab2 = yLab2
        
        !At points3
        OldYCap3 = yCap3
        OldYCons3 = yCons3
        OldYLab3 = yLab3

     elseif (iter > 4 .and. reldifcons < 5*tol .and. reldifcap < 5*tol) then
        !At Origin
        OldYCap0 = yCap0
        OldYCons0 = yCons0  
        OldYLab0 = yLab0
        
        !At points2
        OldYCap2 = yCap2
        OldYCons2 = yCons2
        OldYLab2 = yLab2
        
        !At points3
        OldYCap3 = yCap3
        OldYCons3 = yCons3
        OldYLab3 = yLab3

     else
        !At Origin
        OldYCap0 = acc*yCap0 + (1-acc)*OldYCap0
        OldYCons0 = acc*yCons0 + (1-acc)*OldYCons0  
        OldYLab0 = acc*yLab0 + (1-acc)*OldYLab0
        
        !At points2
        OldYCap2 = acc*yCap2 + (1-acc)*OldYCap2
        OldYCons2 = acc*yCons2 + (1-acc)*OldYCons2
        OldYLab2 = acc*yLab2 + (1-acc)*OldYLab2
        
        !At points3
        OldYCap3 = acc*yCap3 + (1-acc)*OldYCap3
        OldYCons3 = acc*yCons3 + (1-acc)*OldYCons3
        OldYLab3 = acc*yLab3 + (1-acc)*OldYLab3
     endif
     
     !Calculates Coefficients of Smolyak Algorithm
     do j = 1,N
        call interpolate(OldYCap0(j),OldYCap2(:,:,:,:,j),OldYCap3(:,:,j), & 
             & CoeffsCap1(:,:,j),CoeffsCap2(:,:,:,:,j),CoeffsCap3(:,:,j))
        call interpolate(OldYCons0(j),OldYCons2(:,:,:,:,j),OldYCons3(:,:,j), &
             & CoeffsCons1(:,:,j),CoeffsCons2(:,:,:,:,j),CoeffsCons3(:,:,j))
        call interpolate(OldYLab0(j),OldYLab2(:,:,:,:,j),OldYLab3(:,:,j), &
             & CoeffsLab1(:,:,j),CoeffsLab2(:,:,:,:,j),CoeffsLab3(:,:,j))
     enddo
     
     !Use Nonlinear Equation Solver to find policy function for each of smolyak points and each shock
     !!	!At the origin
     do j = 1,N
        stind(j) = 3		!Note: 3 = ((2**(q-d)+1)+1)/2
        stind(N+j) = 3
        guess(1+j) = OldYCap0(j)
     end do
        guess(1) = OldYCons0(1)
     
     iter1 = 0		!Used to keep track of number of iterations done by DNEQNF
     
     call DNEQNF(EulerEq, Relerr, count, itmax, guess, final, fnorm)
     
     yCons0(1) = final(1)
     do j = 2,N
        yCons0(j) = EECons(j)
     end do
     do j = 1,N
        yCap0(j) = final(1+j)
        yLab0(j) = EELab(j)
     enddo
     
     !Error check to make sure solution is sensible
     do j = 1,N
        if (yCap0(j) < kmin(j)) then
           print*, 'solution below kmin in Policy for country', j
!           pause
        elseif (yCap0(j) > kmax(j)) then
           print*, 'solution above kmax in Policy for country', j
!           pause
        endif
        if (yCons0(j) < 0) then
           print*, 'negative consumption in Policy for country', j
           pause
        endif
        if (yLab0(j) < 0) then
           print*, 'negative labor in Policy for country', j
           pause
        elseif (yLab0(j) > Le) then
           print*, 'labor supply above Le in Policy for country', j
           pause
        endif
     enddo
     
     !!	!At points2
     do i = 1,d-1
        do i1 = i+1,d
           do h = 1, 2		!That is, h = 1, 2**(q-d-1)  where q-d = 2
              do h1 = 1, 2	!Same comment applies
                 do j = 1,N
                    !Set current index of capital stock for each country
                    if (j == i .and. h == 1) then
                       stind(j) = 1
                    elseif (j == i .and. h == 2) then
                       stind(j) = 5
                    elseif (j == i1 .and. h1 == 1) then
                       stind(j) = 1
                    elseif (j == i1 .and. h1 == 2) then
                       stind(j) = 5
                    else 
                       stind(j) = 3
                    endif
                    
                    !Set current index of technology shock for each country
                    if (N+j == i .and. h == 1) then
                       stind(N+j) = 1
                    elseif (N+j == i .and. h == 2) then
                       stind(N+j) = 5
                    elseif (N+j == i1 .and. h1 == 1) then
                       stind(N+j) = 1
                    elseif (N+j == i1 .and. h1 == 2) then
                       stind(N+j) = 5
                    else
                       stind(N+j) = 3
                    endif
                    
                    !Set guess for policy functions for each country
                    guess(1+j) = OldYCap2(i,i1,h,h1,j)
                 end do
                 guess(1) = OldYCons2(i,i1,h,h1,1)

                 iter1 = 0	!Used to keep track of number of iterations done by DNEQNF

                 call DNEQNF(EulerEq, Relerr, count, itmax, guess, final, fnorm)
                
                 yCons2(i,i1,h,h1,1) = final(1)
                 do j = 2,N
                    yCons2(i,i1,h,h1,j) = EECons(j)
                 end do
                 do j = 1,N
                    yCap2(i,i1,h,h1,j) = final(1+j)
                    yLab2(i,i1,h,h1,j) = EELab(j)
                 enddo

                 !Error check to make sure solution is sensible
                 do j = 1,N
                    if (yCap2(i,i1,h,h1,j) < kmin(j)) then
                       print*, 'solution below kmin in Policy for country', j, stind
!                       pause
                    elseif (yCap2(i,i1,h,h1,j) > kmax(j)) then
                       print*, 'solution above kmax in Policy for country', j, stind
!                       pause
                    endif
                    if (yCons2(i,i1,h,h1,j) < 0) then
                       print*, 'negative consumption in Policy for country', j
                       pause
                    endif
                    if (yLab2(i,i1,h,h1,j) < 0) then
                       print*, 'negative labor in Policy for country', j
                       pause
                    elseif (yLab2(i,i1,h,h1,j) > Le) then
                       print*, 'labor supply above Le in Policy for country', j
                       pause
                    endif
                 enddo
              enddo
           enddo
        enddo
     enddo
!print*, yCap2(1,2,1,1,:)
!print*, yLab2(1,2,1,1,:)     
!print*, yCons2(1,2,1,1,:)     
!!!!!At points3
     do i = 1,d
        do h = 1, 4		!That is, h = 1, 2**(q-d)  where q-d = 2
           do j = 1,N
              !Set current index of capital stock for each country
              if (j==i .and. h < 3) then
                 stind(j) = h
              elseif (j==i .and. h >= 3) then
                 stind(j) = h+1
              else
                 stind(j) = 3
              endif
              
              !Set current index of technology stock for each country
              if (N+j == i .and. h < 3) then
                 stind(N+j) = h
              elseif (N+j == i .and. h >= 3) then
                 stind(N+j) = h+1
              else
                 stind(N+j) = 3
              endif
              
              !Set guess for policy functions for each country
              guess(1+j) = OldYCap3(i,h,j)
           enddo
           guess(1) = OldYCons3(i,h,1)
 
           iter1 = 0	!Used to keep track of number of iterations done by DNEQNF
           
           call DNEQNF(EulerEq, Relerr, count, itmax, guess, final, fnorm)
            
           yCons3(i,h,1) = final(1)
           do j = 2,N
              yCons3(i,h,j) = EECons(j)
           end do
           do j = 1,N
              yCap3(i,h,j) = final(1+j)
              yLab3(i,h,j) = EELab(j)
           enddo

           !Error check to make sure solution is sensible
           do j = 1,N
              if (yCap3(i,h,j) < kmin(j)) then
                 print*, 'solution below kmin in Policy for country', j
!                 pause
              elseif (yCap3(i,h,j) > kmax(j)) then
                 print*, 'solution above kmax in Policy for country', j
!                 pause
              endif
              if (yCons3(i,h,j) < 0) then
                 print*, 'negative consumption in Policy for country', j
                 pause
              endif
              if (yLab3(i,h,j) < 0) then
                 print*, 'negative labor in Policy for country', j
                 pause
              elseif (yLab3(i,h,j) > Le) then
                 print*, 'labor supply above Le in Policy for country', j
                 pause
              endif
           enddo
        enddo
     enddo
     
     reldifcons0 = maxval(abs((yCons0-OldYCons0)/OldYCons0))
     reldifcons2 = maxval(abs((yCons2-OldYCons2)/OldYCons2))
     reldifcons3 = maxval(abs((yCons3-OldYCons3)/OldYCons3))
     
     reldifcap0 = maxval(abs((yCap0-OldYCap0)/OldYCap0))
     reldifcap2 = maxval(abs((yCap2-OldYCap2)/OldYCap2))
     reldifcap3 = maxval(abs((yCap3-OldYCap3)/OldYCap3))
     
     reldiflab0 = maxval(abs((yLab0-OldYLab0)/OldYLab0))
     reldiflab2 = maxval(abs((yLab2-OldYLab2)/OldYLab2))
     reldiflab3 = maxval(abs((yLab3-OldYLab3)/OldYLab3))

     reldifcons = max(reldifcons0,reldifcons2,reldifcons3)
     reldifcap = max(reldifcap0,reldifcap2,reldifcap3)
     reldiflab = max(reldiflab0,reldiflab2,reldiflab3)
     print*, reldifcons, reldifcap, reldiflab

     if (reldifcons < tol .and. reldifcap < tol .and. reldiflab < tol) then
        exit 
     endif

  end do
  
  print*,'Convergence Achieved after Iteration',iter
  
  !Calculate Running Time
  call CPU_TIME (time_end)
  policy_runtime = policy_runtime + time_end - time_begin
  
  !Calculates Coefficients of Smolyak Algorithm
  do j = 1,N
     call interpolate(yCap0(j),yCap2(:,:,:,:,j),yCap3(:,:,j), & 
          & CoeffsCap1(:,:,j),CoeffsCap2(:,:,:,:,j),CoeffsCap3(:,:,j))
     call interpolate(yCons0(j),yCons2(:,:,:,:,j),yCons3(:,:,j), & 
          & CoeffsCons1(:,:,j),CoeffsCons2(:,:,:,:,j),CoeffsCons3(:,:,j))
     call interpolate(yLab0(j),yLab2(:,:,:,:,j),yLab3(:,:,j), & 
          & CoeffsLab1(:,:,j),CoeffsLab2(:,:,:,:,j),CoeffsLab3(:,:,j))
  enddo
  
end subroutine Policy
