program fullproblem
c -----------------------------------------------------------------------------
c Purpose: This routine is to assemble the linear system
c          (6) in the paper entitled 
c          `A note on integrals for birth-death processes'
c           by Valery Stefanov and Song Wang. The routine
c           calls the linear sparse matrix solver -- Harwell
c           MA28. This set of subroutines are freely available
c           for non-commercial research at
c           http://www.netlib.org/harwell/
c           You need to download all the subroutines of MA28 listed
c           in the makefile
c
c Algorithm description: We number the variables in the following order
c           x_{ij} = Cov(N_i, N_j), j = 1,2, ..., N-1, i = 1,2,..., N-1,
c           y_{ij} = Cov(N_i, S_j), j = 1,2, ..., N, i = 1,2,..., N-1,
c           z_{ij} = Cov(S_i, S_j), j = 1,2, ..., N, i = 1,2,..., N.
c           So, there are (N-1)^2 + N(N-1) + N^2 variables altogether.
c    
c
c Authors:  Drs S. Wang and V. Stefanov
c           Department of Mathematics & Statistics
c           The Univeristy of Western Australia
c           Perth, Australia
c
c Version:  1.0
c
c Input variables:  nsize (size of the problem, i.e. variable N in the paper)
c           lambda, mu and kk (i.e.,initial number of infectives, k)
c
c Output variables: zsum (standard diviation of the cost)
c           ey (expected cost)
c
c Other variables: 
c           nz, irn, icn, ikeep, iw, iflag, mtype, a, u, w, rhs
c           are the same as those described in harwell MA28.
c           nsizex, nsizey and nsizez denote the numbers of x_{ij}, 
c           y_{ij} and z_{ij} respectively (see the above).
c
c Functions called:
c           lami : defines the function lambda_i. In this work it is
c                  the SIS model -- lambda_i = lambda*i*(N-i)/N
c           hj   : H_j(i) in the paper
c           ai   : ai = \sum_{j=1}^k H_j (i)/mu_j
c           bi   : bi = mu* \sum_{j=1}^k H_j(i)/mu_j
c
c Note:     In the SIS model, mu_i = mu*i. So, we do not define it
c           as a separate function because of its simplicity. Users
c           may wish to rewrite it as a callable function if other
c           models are used.
c
c ------------------------------------------------------------------------------ 
      integer nmax, licn, lirn
      parameter(nmax = 200000, licn = 4*nmax, lirn = 4*nmax)
      integer nz, irn(lirn), icn(licn), ikeep(nmax,5), iw(nmax,8), 
     *        iflag,mtype, nsize, n, i, j, k, nsizex, nsizey, nsizez,
     *        kx, ky, kyji, kz, nrow, kk,  kji
      real*8 a(licn), u, w(nmax), rhs(nmax), lambda, mu, ai, bi,
     *       tmp1, tmp2, zsum, lami, ey, hj
c      
      write(6,10)
10    format(1x,'Input N, lambda, mu and k : ', $)
      read(5,*) nsize, lambda, mu, kk
      print *, 'M = ', nsize, ', lambda = ', lambda, 
     &       ', mu = ', mu, ', k = ', kk
c --------------------------------------------------------
c   set up parameters for the MA28 subroutines
      u = 1.0
      mtype = 1
c -------------------------------------------------------
c
      nsizex = (nsize-1)*(nsize-1)
      nsizey = (nsize-1)*nsize
      nsizez = nsize*nsize
      n = nsizex + nsizey + nsizez
      print *, 'N = ', n
      if(n .gt. nmax) then
        print *, 'Order of matrix = ', n, ' > nmax = ', nmax
        print *, 'Increase nmax and run again'
        stop
      end if
c
c ------------- Assembly of the system matrix and RHS----------
c
      nz = 0
      nrow = 0
c -----------------------------------------------------------------
c assemble the first equation system
c
c  x_{i,j} - lambda_j y_{i,j} - lambda_i y_{j,i}
c         + lambda_i*lambda_j/(i*j) z_{i,j} = 0 if i =|= j 
c         or ai*lambda_i if i == j
c
      do i = 1, nsize -1
         do j = 1, nsize - 1
            kx = (i-1)*(nsize-1) + j
            ky = (i-1)*nsize + j + nsizex
            kyji = (j-1)*nsize + i + nsizex
            kz = (i-1)*nsize + j + nsizex + nsizey
            nrow = nrow + 1
            tmp1 = lami(i,nsize,lambda)
            tmp2 = lami(j,nsize,lambda)
c assemble  x_{i,j}
            nz = nz + 1
            a(nz) = 1.0d0
            irn(nz) = nrow
            icn(nz) = kx
c assemble  - lambda_j y_{i,j} - lambda_i y_{j,i}
            if(i .eq. j) then
               nz = nz + 1
               a(nz) = -2.0d0*tmp1
               irn(nz) = nrow
               icn(nz) = ky
            else
               if(ky .gt. kyji) then
                  nz = nz + 1
                  a(nz) = -tmp1
                  irn(nz) = nrow
                  icn(nz) = kyji
                  nz = nz + 1
                  a(nz) = -tmp2
                  irn(nz) = nrow
                  icn(nz) = ky
               else
                  nz = nz + 1
                  a(nz) = -tmp2
                  irn(nz) = nrow
                  icn(nz) = ky
                  nz = nz + 1
                  a(nz) = -tmp1
                  irn(nz) = nrow
                  icn(nz) = kyji
               end if
            end if
c assemble   lambda_i*lambda_j/(i*j)  z_{ij}
             nz = nz + 1
             a(nz) = tmp1*tmp2
             irn(nz) = nrow
             icn(nz) = kz
c assemble the RHS
            if(j .eq. i) then
               rhs(nrow) = ai(i,kk,nsize,lambda,mu)*tmp1
            else
               rhs(nrow) = 0.0d0
            end if
         end do
      end do 
c -----------------------------------------------------------------
c assemble the second equation system
c
c  x_{i-1,j-1} - mu_j y_{i-1,j} - mu_i y_{j-1,i}
c         + mu_i*mu_j/(i*j) z_{i,j} = 0 if i =! j
c         or bi if i == j
c
      do i = 1, nsize 
         do j = 1, nsize
            kx = (i-2)*(nsize-1) + j - 1
            ky = (i-2)*nsize + j + nsizex
            kyji = (j-2)*nsize + i + nsizex
            kz = (i-1)*nsize + j + nsizex + nsizey
            nrow = nrow + 1
            tmp1 = mu*dble(i)
            tmp2 = mu*dble(j)
c assemble  x_{i-1,j-1}
            if(i.gt.1 .and. j.gt.1) then
              nz = nz + 1
              a(nz) = 1.0d0
              irn(nz) = nrow
              icn(nz) = kx
            end if
c assemble  - mu_j y_{i-1,j} - mu_i y_{j-1,i}
            if(i .eq. 1 .and. j .eq.1 ) goto 500
            if(i.eq.1) then
              nz = nz + 1
              a(nz) = -tmp1
              irn(nz) = nrow
              icn(nz) = kyji   
              goto 500
            end if
            if(j.eq.1) then
              nz = nz + 1
              a(nz) = -tmp2
              irn(nz) = nrow
              icn(nz) = ky   
              goto 500
            end if
            if(i .eq. j ) then
               nz = nz + 1
               a(nz) = -2.0d0*tmp1
               irn(nz) = nrow
               icn(nz) = ky
            else
              if(ky .gt. kyji) then
                 nz = nz + 1
                 a(nz) = -tmp1
                 irn(nz) = nrow
                 icn(nz) = kyji
                 nz = nz + 1
                 a(nz) = -tmp2
                 irn(nz) = nrow
                 icn(nz) = ky
              else
                 nz = nz + 1
                 a(nz) = -tmp2
                 irn(nz) = nrow
                 icn(nz) = ky
                 nz = nz + 1
                 a(nz) = -tmp1
                 irn(nz) = nrow
                 icn(nz) = kyji
              end if
            end if
c---- assemble   mu_i*mu_j/(i*j)  z_{ij}
500         continue
            nz = nz + 1
            a(nz) = tmp1*tmp2
            irn(nz) = nrow
            icn(nz) = kz
c assemble the RHS
            if(j .eq. i) then
              rhs(nrow) = bi(i,kk,nsize,lambda,mu)
            else
               rhs(nrow) = 0.0d0
            end if
         end do
      end do
c -----------------------------------------------------------------
c assemble the third equation system
c
c  x_{i,j-1} - mu_j y_{i,j} - lambda_i y_{j-1,i}
c         + lambda_i*mu_j/(i*j) z_{i,j} = 0 
c
      do i = 1, nsize - 1
         do j = 1, nsize
            kx = (i-1)*(nsize - 1) + j - 1
            ky = (i-1)*nsize + j + nsizex
            kyji = (j-2)*nsize + i + nsizex
            kz = (i-1)*nsize + j + nsizex + nsizey
            nrow = nrow + 1
            tmp1 = lami(i,nsize,lambda)
            tmp2 = mu*dble(j)
c assemble  x_{i,j-1}
            if(j.gt.1) then
              nz = nz + 1
              a(nz) = 1.0d0
              irn(nz) = nrow
              icn(nz) = kx
            end if
c assemble  - mu_j y_{i,j} - lambda y_{j-1,i}
            if(j.eq.1) then
               nz = nz + 1
               a(nz) = -tmp2
               irn(nz) = nrow
               icn(nz) = ky
               goto 600
            end if
            if(ky .gt. kyji) then
               nz = nz + 1
               a(nz) = -tmp1
               irn(nz) = nrow
               icn(nz) = kyji
               nz = nz + 1
               a(nz) = -tmp2
               irn(nz) = nrow
               icn(nz) = ky
            else
               nz = nz + 1
               a(nz) = -tmp2
               irn(nz) = nrow
               icn(nz) = ky
               nz = nz + 1
               a(nz) = -tmp1
               irn(nz) = nrow
               icn(nz) = kyji
            end if
c assemble   mu_i*mu_j/(i*j)  z_{ij}
600         continue
             nz = nz + 1
             a(nz) = tmp1*tmp2
             irn(nz) = nrow
             icn(nz) = kz
c assemble the RHS
             rhs(nrow) = 0.0d0
         end do
      end do
c 
c Check for errors
c
      if( nrow .ne. n) then
        print *, 'Matrix is not square!'
        print *, 'No of rows = ', nrow, ',  No of columns = ', n
        stop
      end if
c
c ----------- solve the linear system ------------------------------
c
      print *, 'LU decomposition'
      call  ma28ad(n, nz, a, licn, irn, lirn, icn, u, ikeep, iw, w,
     * iflag)
      print *, 'Backward substituting'
      call ma28cd(n, a, licn, icn, ikeep, rhs, w, mtype)
c
c evaluate \sum_{i,j = 1}^nsize} i*j*z_{i,j}
c
      zsum = 0.0d0
      do i = 1, nsize
         do j = 1, nsize
            k = (i-1)*nsize + j + nsizex + nsizey
            kji = (j-1)*nsize + i + nsizex + nsizey
            zsum = zsum + rhs(k)*(dble(i)*dble(j))
         end do
      end do
      zsum = dsqrt(zsum)
c
c evaluate the expectation 
c   E(y_f) = \sum_{j=1}^kk (1/mu_j \sum_{i=j}^nsize i*H_j(i))
c
      ey = 0.0d0
      do j = 1, kk
         tmp1 = 0.0d0
         do i = j, nsize
            tmp1 = tmp1 + dble(i)*hj(i, j, nsize,lambda,mu)
         end do
         ey = ey + tmp1/(mu*dble(j))
      end do
c
      print *, 'Standard diviation = ', zsum, ', E(y_f) = ', ey
c      print *, 'sum of z_{ij} = ', sqrt(zsum), ',
c     &          E(y_f) = ', ey
      end
c -----  function ai ---------------------
      real*8 function ai(i,k,nsize,lambda,mu)
      integer i, k, nsize,j
      real*8 lambda,mu, hj
c
      ai = 0.0d0
      do j = 1, k
         ai = ai + hj(i,j,nsize,lambda,mu)/(mu*dble(j))
      end do
      return
      end
c --- function bi ----------------
      real*8 function bi(i,k,nsize,lambda,mu)
      integer i, k, nsize,j
      real*8 lambda,mu, hj, tmp, lami
c
c      tmp = 0.0d0
c      do j = 1, k
c         tmp = tmp + hj(i-1,j,nsize,lambda,mu)/(mu*dble(j))
c      end do
c      if(i .le. k) then
c        bi =  lami(i-1,nsize,lambda)*tmp + 1.0d0 
c      else
c        bi = lami(i-1,nsize,lambda)*tmp
c      end if
c -------- new form of bi  -----------------
      bi = 0.0d0
      do j = 1, k
         bi = bi + hj(i,j,nsize,lambda,mu)/(mu*dble(j))
      end do
      bi = bi*mu*dble(i)
      return
      end
c --- function Hj ----------------
      real*8 function hj(i,j,nsize,lambda,mu)
      integer i,j,nsize, k
      real*8 lambda, mu, tmp1, tmp2, lami
c
      if( j .gt. i) then
        hj = 0.0d0
        return
      end if
      if(i .eq. j) then
        hj = 1.0d0
        return
      end if
      if(i .gt. j) then
        tmp1 = 1.0d0
        tmp2 = 1.0d0
        do k = j, i-1
           tmp1 = tmp1*lami(k,nsize,lambda)
           tmp2 = tmp2*mu*dble(k+1)
        end do
        if(tmp2 .le. 1.0d-10) then
          print *, 'Error in the evaluation of H_j(i)'
          stop
        end if
        hj = tmp1/tmp2
      end if
      return
      end
c --------- lambda_i -------------
      real*8 function lami(i,nsize,lambda)
      integer i, nsize
      real*8 lambda
c
      lami = lambda*dble(i)*dble(nsize-i)/dble(nsize)
      return
      end