MODULE FRSEL            ! fits fixed and random effects selectivity model of Fryer (1991)

! programmer:               Rob Fryer
! date:                     21 October 1996
! language:                 Fortran 90 (standard)
! auxiliary routines:       NAG F01ABF, F03ABF
! documentation:            available from author


USE FRSEL_SETUP

CONTAINS
	  
  SUBROUTINE FRSEL_FIT(haul, mltype, niter, tol, monit, d_specified, d, log_lik, alpha, se, vcov, error_message)

    IMPLICIT NONE

  ! Arguments

    TYPE (HAUL_INFO), DIMENSION(:), INTENT(INOUT) :: haul     ! haul information

    CHARACTER*1, OPTIONAL, INTENT(IN)  :: mltype              ! type of likelihood: 'm' = maximum likelihood, or the default,
                                                              ! 'r' = residual maximum likelihood
    INTEGER, OPTIONAL, INTENT(IN)      :: niter               ! number of iterations:  default 1000
    REAL (wp), OPTIONAL, INTENT(IN)    :: tol                 ! convergence tolerance: default 10**(-12)
    INTEGER, OPTIONAL, INTENT(IN)      :: monit               ! controls monitoring, printing every monit iterations: default 0
    LOGICAL, OPTIONAL, INTENT(IN)      :: d_specified         ! determines whether d is specified: default false

    REAL (wp), OPTIONAL, INTENT(INOUT) :: d(nnu, nnu)         ! between-haul variance: default identity matrix

    REAL (wp), OPTIONAL, INTENT(OUT)   :: log_lik             ! (residual) log-likelihood
    REAL (wp), OPTIONAL, INTENT(OUT)   :: alpha(nalpha)       ! parameters to be estimated
    REAL (wp), OPTIONAL, INTENT(OUT)   :: se(nalpha)          ! standard errors
    REAL (wp), OPTIONAL, INTENT(OUT)   :: vcov(nalpha,nalpha) ! variance matrix
    INTEGER, OPTIONAL, INTENT(OUT)     :: error_message       


  ! Local variables

    INTEGER     :: i, niter_wk
    REAL (wp)   :: tol_wk, d_wk(nnu,nnu), log_lik_wk, alpha_wk(nalpha), vcov_wk(nalpha,nalpha)
    CHARACTER*1 :: mltype_wk


    IF (PRESENT(error_message)) error_message = 0

    CALL INITIALISE_VARIABLES(haul, mltype, niter, tol, d_specified, d, mltype_wk, niter_wk, tol_wk, d_wk, error_message)
    IF (PRESENT(error_message) .and. error_message /= 0) RETURN
      

    CALL EM_ALGORITHM(haul, mltype_wk, niter_wk, tol_wk, monit, d_wk, log_lik_wk, alpha_wk, vcov_wk, error_message)
    IF (PRESENT(error_message) .and. error_message /= 0 .and. error_message /= 1) RETURN


    IF (PRESENT(log_lik)) log_lik = log_lik_wk           ! assign optional output variables given convergence, or if run out of 
    IF (PRESENT(alpha)) alpha = alpha_wk                 ! iterations
    IF (PRESENT(se)) THEN
      DO i = 1, nalpha
        se(i) = sqrt(vcov_wk(i,i))
      END DO
      END IF
    IF (PRESENT(d)) d = d_wk
    IF (PRESENT(vcov)) vcov = vcov_wk

    RETURN

  CONTAINS

    ! ****************************************************************************************************************************

    SUBROUTINE INITIALISE_VARIABLES(haul, mltype, niter, tol, d_specified, d, mltype_wk, niter_wk, tol_wk, d_wk, error_message)

      IMPLICIT NONE

    ! Arguments

      TYPE (HAUL_INFO), INTENT(INOUT), DIMENSION(:)   :: haul      

      CHARACTER*1, OPTIONAL, INTENT(IN)               :: mltype     
      INTEGER, OPTIONAL, INTENT(IN)                   :: niter
      REAL (wp), OPTIONAL, INTENT(IN)                 :: tol
      LOGICAL, OPTIONAL, INTENT(IN)                   :: d_specified
      REAL (wp), OPTIONAL, INTENT(IN)                 :: d(nnu,nnu)

      CHARACTER*1, INTENT(OUT)                        :: mltype_wk   
      REAL (wp), INTENT(OUT)                          :: tol_wk
      INTEGER, INTENT(OUT)                            :: niter_wk
      REAL (wp), INTENT(OUT)                          :: d_wk(nnu,nnu)

      INTEGER, OPTIONAL, INTENT(INOUT)                :: error_message       


    ! Local variables    
    
      INTEGER   :: i, j, ifail
      REAL (wp) :: wk(nnu), det


      DO i = 1, SIZE(haul)                             ! fill in haul variance matrices
	DO j = 1, nnu
	  haul(i)%r(j,j:nnu) = haul(i)%r(j:nnu,j)
	END DO
      END DO


      IF (PRESENT(mltype)) THEN                        ! initialise monitoring type
        mltype_wk = mltype
      ELSE
        mltype_wk = 'r'
      END IF
	
      IF (mltype_wk .ne. 'm' .and. mltype_wk .ne. 'r') THEN
        CALL ERROR_CHECK(2, error_message)
        RETURN
      END IF

      
      IF (PRESENT(niter) .and. niter > 1) THEN         ! initialise number of iterations                
        niter_wk = niter
      ELSE
        niter_wk = 1000
      END IF


      IF (PRESENT(tol)) THEN                           ! initialise convergence tolerance 
        tol_wk = tol
      ELSE
        tol_wk = 10.0_wp**(-12)
      END IF

     
      IF (PRESENT(d_specified) .and. d_specified) THEN ! initialise between-haul variance

        IF (PRESENT(d)) THEN
          d_wk = d
          DO i = 1, nnu                                ! fill in upper triangle
            d_wk(i,i:nnu) = d_wk(i:nnu,i)
          END DO

          ifail = 1
          CALL F03ABF(d_wk, nnu, nnu, det, wk, ifail)  ! check d_wk is positive-definite
          IF (ifail /= 0) THEN
            CALL ERROR_CHECK(3, error_message)
            RETURN
          END IF 

        ELSE
   
          CALL ERROR_CHECK(4, error_message)
          RETURN

        END IF

      ELSE

        d_wk = 0.0_wp
        DO i = 1, nnu
          d_wk(i,i) = 1.0_wp
        END DO

      END IF

      RETURN

    END SUBROUTINE INITIALISE_VARIABLES


    ! ****************************************************************************************************************************

    SUBROUTINE EM_ALGORITHM(haul, mltype, niter, tol, monit, d, log_lik, alpha, vcov, error_message)

      IMPLICIT NONE

    ! Arguments

      TYPE (HAUL_INFO), DIMENSION(:), INTENT(INOUT) :: haul   
      CHARACTER*1, INTENT(IN)                       :: mltype 
      INTEGER, INTENT(IN)                           :: niter
      REAL (wp), INTENT(IN)                         :: tol        
      INTEGER, OPTIONAL, INTENT(IN)                 :: monit
      REAL (wp), INTENT(INOUT)                      :: d(nnu,nnu)
      REAL (wp), INTENT(OUT)                        :: log_lik
      REAL (wp), INTENT(OUT)                        :: alpha(nalpha)
      REAL (wp), INTENT(OUT)                        :: vcov(nalpha,nalpha)
      INTEGER, OPTIONAL, INTENT(INOUT)              :: error_message

    ! Local variables

      INTEGER   :: iiter
      REAL (wp) :: log_lik_old, log_lik_1


      CALL LOG_LIK_1_CALC(haul, mltype, log_lik_1, error_message)   ! calculate constant term in log-likelihood
      IF (PRESENT(error_message) .and. error_message /= 0) RETURN

 
      DO iiter = 1, niter                                           ! EM algorithm
 
        IF (iiter > 1) log_lik_old = log_lik                        ! store old log-likelihood
        log_lik = log_lik_1                                         ! initialise log-likelihood with constant
 

        CALL W_CALC(haul, d, log_lik, error_message)                ! calculate variance matrices w and increment log-likelihood
        IF (PRESENT(error_message) .and. error_message /= 0) RETURN 


        CALL VCOV_CALC(haul, mltype, vcov, log_lik, error_message)  ! calculate var(alpha) and increment likelihood
        IF (PRESENT(error_message) .and. error_message /= 0) RETURN 


        CALL B_CALC(haul, vcov, d, alpha, log_lik)                  ! calculate alpha, b, and increment log-likelihood


        CALL D_CALC(haul, vcov, mltype, d)                          ! calculate new value of d


        log_lik = - log_lik / 2.0_wp                                ! convert to log-likelihood
        IF (iiter > 1) THEN                                         ! check for convergence
          IF (log_lik - log_lik_old .le. tol) RETURN
        END IF

        IF (PRESENT(monit) .and. monit > 0 .and. MOD(iiter, monit) == 0) CALL MPRINT(log_lik, alpha, d)

      END DO

      CALL ERROR_CHECK(1, error_message)

      RETURN

    END SUBROUTINE EM_ALGORITHM

    ! ****************************************************************************************************************************
 
    SUBROUTINE LOG_LIK_1_CALC(haul, mltype, log_lik_1, error_message)
 
    ! calculate constant term in (residual) log-likelihood

      IMPLICIT NONE

    ! Arguments

      TYPE (HAUL_INFO), DIMENSION(:), INTENT(IN) :: haul
      CHARACTER*1, INTENT(IN)                    :: mltype 
      REAL (wp), INTENT(OUT)                     :: log_lik_1
      INTEGER, OPTIONAL, INTENT(INOUT)           :: error_message

    ! Local variables

      INTEGER    :: i, ifail, nhaul
      REAL (wp)  :: wk(nalpha, nalpha), det, wkspce(nalpha)
      REAL (wp)  :: pi = 3.14159265


      nhaul = SIZE(haul)

      IF (mltype == 'r') THEN

        wk = 0.0_wp                                                 ! calculate sum(transpose(x) *+ x)            
        DO i = 1, nhaul
          wk = wk + MATMUL(TRANSPOSE(haul(i)%x), haul(i)%x)
        END DO
      
        ifail = 1
        CALL F03ABF(wk, nalpha, nalpha, det, wkspce, ifail)         ! calculate det(sum)
        IF (ifail .ne. 0) THEN
          CALL ERROR_CHECK(ifail + 10, error_message)
          RETURN
	END IF
  
        log_lik_1 = REAL(nnu * nhaul - nalpha, wp) * log(2.0_wp * pi) - log(det)

      ELSE
      
        log_lik_1 = REAL(nnu * nhaul, wp) * log(2.0_wp * pi)
     
      END IF

      RETURN

    END SUBROUTINE LOG_LIK_1_CALC

    ! ****************************************************************************************************************************
 
    SUBROUTINE W_CALC(haul, d, log_lik, error_message)
 
    ! calculate w matrices and increment log-likelihood  

      IMPLICIT NONE

    ! Arguments

      TYPE (HAUL_INFO), DIMENSION(:), INTENT(INOUT) :: haul
      REAL (wp), INTENT(IN)                         :: d(nnu,nnu)
      REAL (wp), INTENT(INOUT)                      :: log_lik 
      INTEGER, OPTIONAL, INTENT(INOUT)              :: error_message
 
    ! Local variables

      INTEGER    :: i, j, ifail          
      REAL (wp)  :: v(nnu+1,nnu), wk(nnu), det
 

      DO i = 1, SIZE(haul)
 
	v(1:nnu,:) = haul(i)%r + d                                   ! v = r + d     
      
        ifail = 1
        CALL F01ABF(v, nnu+1, nnu, haul(i)%w, nnu, wk, ifail)        ! w = inverse(v)
        IF (ifail /= 0) THEN
          CALL ERROR_CHECK(20+ifail, error_message)
          RETURN
        END IF

        DO j = 1, nnu
          haul(i)%w(j,j:nnu) = haul(i)%w(j:nnu,j)                    ! fill in upper triangle
        END DO

        ifail = 1                                                    ! det(v)  
        CALL F03ABF(v(1:nnu,:), nnu, nnu, det, wk, ifail)
        IF (ifail /= 0) THEN
          CALL ERROR_CHECK(30+ifail, error_message)
          RETURN
        END IF

        log_lik = log_lik + log(det)                                 ! increment log-likelihood by sum(log(det(v))) 

      END DO

    END SUBROUTINE W_CALC
 
    ! ***************************************************************************************************************************
 
    SUBROUTINE VCOV_CALC(haul, mltype, vcov, log_lik, error_message)

    ! calculate variance(alpha) and increment log-likelihood further

      IMPLICIT NONE

    ! Arguments
  
      TYPE (HAUL_INFO), DIMENSION(:), INTENT(IN) :: haul        
      CHARACTER*1, INTENT(IN)                    :: mltype
      REAL (wp), INTENT(OUT)                     :: vcov(nalpha,nalpha)
      REAL (wp), INTENT(INOUT)                   :: log_lik     
      INTEGER, INTENT(INOUT)                     :: error_message

    ! Local variables

      INTEGER    :: i, ifail
      REAL (wp)  :: det, wk(nalpha+1,nalpha), wkspce(nalpha)
      

      wk = 0.0_wp                                                    ! sum(t(x) *+ w *+ x)
      DO i = 1, SIZE(haul)
	wk(1:nalpha,:) = wk(1:nalpha,:) + MATMUL(TRANSPOSE(haul(i)%x), MATMUL(haul(i)%w, haul(i)%x))
      END DO


      ifail = 1                                                      ! vcov = inverse(sum(t(x) *+ w *+ x))
      CALL F01ABF(wk, nalpha+1, nalpha, vcov, nalpha, wkspce, ifail)
      IF (ifail .ne. 0) THEN
	CALL ERROR_CHECK(40+ifail, error_message)
	RETURN
      END IF

      DO i = 1, nalpha                                               ! fill in upper triangle
	vcov(i,i:nalpha) = vcov(i:nalpha,i)
      END DO


      IF (mltype == 'r') THEN                                        ! increment log-likelihood by log(det(sum(t(x) *+ w *+ x)))

        CALL F03ABF(wk(1:nalpha,:), nalpha, nalpha, det, wkspce, ifail)          
        IF (ifail .ne. 0) THEN
          CALL ERROR_CHECK(50+ifail, error_message)
	  RETURN
        END IF

        log_lik = log_lik + log(det)

      END IF     

    END SUBROUTINE VCOV_CALC

    ! ****************************************************************************************************************************
 
    SUBROUTINE B_CALC(haul, vcov, d, alpha, log_lik)
  
    ! calculate alpha, b, and increment the log-likelihood       

      IMPLICIT NONE

    ! Arguments

      TYPE (HAUL_INFO), DIMENSION(:), INTENT(INOUT) :: haul     
      REAL (wp), INTENT(IN)                         :: vcov(nalpha,nalpha)
      REAL (wp), INTENT(IN)                         :: d(nnu,nnu)
      REAL (wp), INTENT(OUT)                        :: alpha(nalpha)
      REAL (wp), INTENT(INOUT)                      :: log_lik      

    ! Local variables

      INTEGER   :: i
      REAL (wp) :: wk1(nnu), wk2(nnu)


      alpha = 0.0_wp                                                 ! alpha = vcov *+ sum(t(x) *+ w *+ nu)
      DO i = 1, SIZE(haul)
	alpha = alpha + MATMUL(TRANSPOSE(haul(i)%x), MATMUL(haul(i)%w, haul(i)%nu))
      END DO
      alpha = MATMUL(vcov, alpha)


      DO i = 1, SIZE(haul)                                           ! b = d *+ w *+ (nu - x *+ alpha) and increment log-lik
	wk1 = haul(i)%nu - MATMUL(haul(i)%x, alpha)                  ! by sum(t(nu - x *+ alpha) *+ w *+ (nu - x *+ alpha))
	wk2 = MATMUL(haul(i)%w, wk1)
	haul(i)%b = MATMUL(d, wk2)
	log_lik = log_lik + SUM(wk1 * wk2)
      END DO
   
    END SUBROUTINE B_CALC

    ! ****************************************************************************************************************************
 
    SUBROUTINE D_CALC(haul, vcov, mltype, d)
 
    ! update between-haul variance d

      IMPLICIT NONE
      
    ! Arguments

      TYPE (HAUL_INFO), DIMENSION(:), INTENT(IN) :: haul        
      REAL (wp), INTENT(IN)                      :: vcov(nalpha,nalpha)
      CHARACTER*1, INTENT(IN)                    :: mltype
      REAL (wp), INTENT(INOUT)                   :: d(nnu,nnu)

    ! Local variables   
    
      INTEGER                                    :: nhaul, i
      REAL (wp), DIMENSION(SIZE(haul), nnu)      :: bwk
      REAL (wp), DIMENSION(nnu,nnu)              :: dnew, wk
      

    ! ML: calculate dnew = dold - sum(dold *+ w *+ d) / nhaul + sum(b *+ t(b)) / nhaul

      nhaul = SIZE(haul)

      wk = 0.0_wp
      DO i = 1, nhaul
	bwk(i,:) = haul(i)%b
	wk = wk + haul(i)%w
      END DO
	
      wk = - MATMUL(d, MATMUL(wk, d)) + MATMUL(TRANSPOSE(bwk), bwk) 
      dnew = d + wk / nhaul
							 
      
    ! REML: dnew = as above + sum(dold *+ w *+ x *+ vcov *+ t(x) *+ w *+ dold) / nhaul
    
      IF (mltype .eq. 'r') THEN               
	wk = 0.0_wp
	DO i = 1, nhaul
	  wk = wk + MATMUL(haul(i)%w, MATMUL(haul(i)%x, MATMUL(vcov, MATMUL(TRANSPOSE(haul(i)%x), haul(i)%w))))
	END DO
	dnew = dnew + MATMUL(d, MATMUL(wk, d)) / nhaul
      END IF
    
      d = dnew

      RETURN

    END SUBROUTINE D_CALC      

    ! ****************************************************************************************************************************

    SUBROUTINE MPRINT(log_lik, alpha, d)

    ! monitor convergence

      IMPLICIT NONE

    ! Arguments

      REAL(wp), INTENT(IN) :: log_lik
      REAL(wp), INTENT(IN) :: alpha(nalpha)
      REAL(wp), INTENT(IN) :: d(nnu,nnu)

    ! Local variables

      INTEGER i, j 

      WRITE(*,'(//,A,//,A,e15.8)') '*** convergence monitoring ***', 'log-likelihood = ', log_lik
     
      WRITE(*,'(/,A,/)') '    alpha'
      DO i = 1, nalpha
        WRITE(*,'(f10.4)') alpha(i)
      END DO
    
      WRITE(*,'(/,A,/)') '    d'
      DO i = 1, nnu
        WRITE(*,'(8f10.4)') (d(i,j), j = 1, nnu)
      END DO

      RETURN

    END SUBROUTINE MPRINT

    ! ***************************************************************************************************************************

    SUBROUTINE ERROR_CHECK(icheck, error_message)

      IMPLICIT NONE

      INTEGER, INTENT(IN)            :: icheck     
      INTEGER, OPTIONAL, INTENT(OUT) :: error_message

      IF (PRESENT(error_message)) THEN
      
        error_message = icheck
        RETURN

      ELSE

        SELECT CASE (icheck)
        CASE (1) 
          WRITE(*,111) '*** ERROR ran out or iterations:  error_message = ', icheck
        CASE (2)
          WRITE(*,111) '*** ERROR mltype is invalid: error_message = ', icheck
        CASE (3)
          WRITE(*,111) '*** ERROR d is not positive definite: error_message = ', icheck
        CASE (4)
          WRITE(*,111) '*** ERROR d is not specified: error_message = ', icheck
        CASE (11:14)
          WRITE(*,111) '*** ERROR calculating det(sum(t(x) *+ x)):  NAG F03ABF ifail = ', icheck - 10
        CASE (21:22)
          WRITE(*,111) '*** ERROR calculating w matrix for some haul:  NAG F01ABF ifail = ', icheck - 20
        CASE (31:34)
          WRITE(*,111) '*** ERROR calculating det(v) for some haul:  NAG F03ABF ifail = ', icheck - 30
        CASE (41:42)
          WRITE(*,111) '*** ERROR calculating variance(alpha):  NAG F01ABF ifail = ', icheck - 40
        CASE (51:54)
          WRITE(*,111) '*** ERROR calculating det(information matrix):  NAG F03ABF ifail = ', icheck - 50
        END SELECT

        STOP

      END IF

      RETURN
      
      111 FORMAT(//,A,i1,//)

    END SUBROUTINE ERROR_CHECK     

  END SUBROUTINE FRSEL_FIT

END MODULE FRSEL