! You need to typecast only the integers when calling the Fortran subroutines, 
! as R seems to initialise all numbers to float (well, "numeric"), and those seem
! to be handled well enough.
!
! The temperature matrix, the observations and the current inverse correlation matrix
! are saved between different function calls, they are globally defined variables
! in the module TUPDATER_DATA
!
! The CISCM needs to be written after each updater step of the parameters
!
! 2014-05-29
! Moved the proxy parameters CPROX into the main module to plan for the future.
! 2011-02-19
! Removed unnecessary cholesky decomposition at each T_t update step. It is now
! calculated when updating the Post.Cov.Now Array after parameter updates
!
MODULE TUPDATER_DATA
  
  ! One problem with variables defined in modules:
  ! It slows execution down, probably a lot, because the code needs to
  ! constantly check for changes and possible reallocations. This is a bit
  ! strange to me. This is avoided when calling from inside a module.
  
  SAVE
  
  INTEGER, PARAMETER :: dp = KIND(1.0D0)

  TYPE :: CSCMatrix
    INTEGER                       :: NNZ, Nrow
    REAL( kind = dp ), ALLOCATABLE :: a(:)
    INTEGER, ALLOCATABLE          :: ia(:), ja(:)
  END TYPE CSCMatrix

  INTEGER :: IX=123                             ! random seed for ziggurat
  INTEGER :: NYEARS, NLOCS, NPT                 ! #years, #locations, #proxies
  REAL( kind = dp ), ALLOCATABLE :: T_MATRIX(:,:), DATA_ALL(:,:), Ciscm(:,:), Cscm(:,:)
  ! T_MATRIX	: NLOCS x NLOCS 		current posterior of temperature
  ! DATA_ALL	: NLOCS*(NPT+1) x NYEARS 	observations / proxies
  ! Ciscm	: NLOCS x NLOCS			current inverse covariance matrix
  REAL( kind = dp ), ALLOCATABLE :: UPatterns(:,:), CProx(:,:,:)
  INTEGER, ALLOCATABLE :: PatByYear(:)
  INTEGER :: NumPat, CovMatType
  TYPE( CSCMatrix), ALLOCATABLE   :: PhiList(:)


CONTAINS

SUBROUTINE getzig( DUMMY, NRAN)

  USE Ziggurat

  IMPLICIT INTEGER(I - K)
  INTEGER :: NRAN
  REAL( kind = dp ) :: DUMMY( NRAN)

  DO I = 1, NRAN
     DUMMY( I) = rnor( )
  END DO

END SUBROUTINE getzig

!------------------------------------------------------
! subroutines to draw new temperature estimates
!------------------------------------------------------
subroutine T0Updater( T0p, Cpars)

  !USE TUPDATER_DATA

  IMPLICIT INTEGER (I-K)

  REAL( kind = dp ), intent(in) :: T0p(2), Cpars(*)

  REAL( kind = dp ), dimension(Nlocs) :: Ones
  REAL( kind = dp ), dimension(NLocs, NLocs) :: postIcovM, postCovMat, Eye
  integer :: INFO 
  
  postCovMat = 0.0
  do i = 1, NLocs
    postCovMat(i,i) = 1
  enddo
  postIcovM = postCovMat
  Eye = postCovMat

  ! Find the (conditional) posterior covariance matrix 
  !  (In the notation of the handwritten notes, this is Psi):
  CALL DGEMM( 'N', 'N', NLocs, NLocs, NLocs, Cpars(1)*Cpars(1)/Cpars(3), Ciscm, NLocs, &
    Eye, NLocs, 1/T0p(2), postIcovM, NLocs)

  ! now invert postCovMat to get Psi
  CALL DPOSV('U', NLocs, NLocs, postIcovM, NLocs, postCovMat, NLocs, INFO)
!  PRINT *, "The matrix inversion Error Code", INFO

  ! Calculate the mean vector:
  ! construct a vector for the residuals
  Ones = 1
  
  T_MATRIX(1:NLOCS,0) = Cpars(1)*T_MATRIX(1:NLOCS,1) - Cpars(1)*(1-Cpars(1))*Cpars(2)*Ones
  CALL DGEMV( 'N', NLocs, NLocs, 1/Cpars(3), Ciscm, NLocs, T_MATRIX(1:NLOCS,0), 1, T0p(1)/T0p(2), Ones, 1)
  T_MATRIX(1:NLOCS,0) = Ones
  CALL DGEMV( 'N', NLOCS, NLOCS, 1.0_dp, postCovMat, NLOCS, Ones, 1, 0.0_dp, T_MATRIX(1:NLOCS,0), 1)
  ! store the mean in T0
  ! in order to sample:
  ! draw NLocs-dimensional normally dist iid random numbers
  ! calculate Cholesky factor of postCovMat = R' %*% R
  ! then
  ! T0.new = postMean + R' %*% Y
  CALL DPOTRF('U', NLocs, postCovMat, NLocs, INFO)   ! Cholesky
  DO i = 2, NLocs
    DO j = 1, i - 1
      postCovMat(i, j) = 0.0_dp
    ENDDO
  ENDDO
  ! overwrite "Ones" by the random numbers, don't need it in this routine anymore
  CALL getzig( Ones, NLOCS)
  CALL DGEMV( 'N', NLOCS, NLOCS, 1.0_dp, postCovMat, NLOCS, Ones, 1, 1.0_dp, T_MATRIX(1:NLOCS, 0), 1)

END SUBROUTINE T0Updater

subroutine TkUpdater_sparse( TIDX, Cpars, HMat, PostCovNow, CholPostCovNow)

  !USE TUPDATER_DATA

  IMPLICIT INTEGER (I-K)

  integer, intent(in) :: TIDX ! The idx of the current year / season / whatever
  REAL( kind = dp ), intent(in) :: Cpars(*)
  ! The H matrix, a stack of diagonal matrices with the factor beta_1
!  REAL( kind = dp ), intent(in) :: Hk(*)
!  integer, intent(in) :: Hkia(*), Hkja(*)
  TYPE( CSCMatrix)             :: HMat
  REAL( kind = dp ), intent(in) :: PostCovNow(NLOCS, NLOCS)
  REAL( kind = dp ), intent(in) :: CholPostCovNow(NLOCS, NLOCS)

  integer :: BIDX, EIDX ! the indices for the proxy loops
  REAL( kind = dp ), dimension(Nlocs) :: Ones
  REAL( kind = dp ), dimension(NLocs, NLocs) :: Eye
  REAL( kind = dp ), allocatable :: BVec(:), Wgood(:)
  integer :: NObsTotal
  
  Ones = 1.0_dp
  Eye= 0.0_dp
  do i = 1, NLocs
    Eye(i,i) = 1.0_dp
  enddo
  NObsTotal = NLOCS * (NPT + 1)
  ! Need to build:
  ! BVec
  ALLOCATE( BVec(NObsTotal) )
  ALLOCATE( Wgood(NObsTotal) )
  ! populate the matrices / vectors, remove the missing obs for this year later (mask with WHERE)
  BVec(1:NLOCS) = 0.0_dp
  DO I = 1, NPT
    BIDX = NLOCS * I + 1
    EIDX = (I+1) * NLOCS
    BVec(BIDX:EIDX) = CProx( I, 3, :)
    IF ( CProx( I, 1, 1) == 1) THEN
      IF ( TIDX > 1) THEN
        BVec(BIDX:EIDX) = BVec(BIDX:EIDX) + &
          UPatterns(BIDX:EIDX,PatByYear(TIDX - 1)) * CProx( I, 5,:) * DATA_ALL(BIDX:EIDX, TIDX - 1)
      END IF
    END IF
  ENDDO
  T_MATRIX(1:NLOCS, TIDX) = Cpars(1)*(T_MATRIX(1:NLOCS, TIDX + 1) + T_MATRIX(1:NLOCS, TIDX -1) ) +&
    (1-Cpars(1))*(1-Cpars(1))*Cpars(2)*Ones
  
  ! this is not strictly necessary and should be caught be the multiplication
  ! with Hk anyway.
  WHERE( DATA_ALL(1:NObsTotal, TIDX) == -99 )
    Wgood = 0
  ELSEWHERE
    Wgood = DATA_ALL(1:NObsTotal,TIDX) - BVec(1:NObsTotal)
  END WHERE
  
  CALL amux( NLOCS, Wgood(1:NObsTotal), Ones(1:NLOCS), HMat%a, HMat%ja, HMat%ia)
!  PRINT *, "A MU X"
!  WRITE (*, '((30F5.1))') ONES(1:NLOCS)
  CALL DGEMV( 'N', NLOCS, NLOCS, 1.0_dp/Cpars(3), Ciscm, NLOCS, T_MATRIX(1:NLOCS, TIDX), 1, 1.0_dp, Ones, 1)
!  PRINT *, "TMAT"
!  WRITE (*, '((30F5.1))') T_MATRIX(1:NLOCS, TIDX)

  CALL DGEMV( 'N', NLOCS, NLOCS, 1.0_dp, PostCovNow, NLOCS, Ones, 1, 0.00_dp, T_MATRIX(1:NLOCS, TIDX), 1)
  CALL getzig( Ones, NLOCS)
  CALL DGEMV( 'N', NLOCS, NLOCS, 1.0_dp, CholPostCovNow, NLOCS, Ones, 1, 1.0_dp, T_MATRIX(1:NLOCS, TIDX), 1)
  
  DEALLOCATE( BVec)
  DEALLOCATE( Wgood)

END SUBROUTINE TkUpdater_sparse

subroutine TlastUpdater( Cpars, UPatNow, PostCovNow, CholPostCovNow)

  !USE TUPDATER_DATA

  IMPLICIT INTEGER (I-K)

  integer :: BIDX, EIDX ! the indices for the proxy loops
  REAL( kind = dp ) :: Cpars(*)
  REAL( kind = dp ), intent(in) :: PostCovNow(NLOCS, NLOCS)
  REAL( kind = dp ), intent(in) :: CholPostCovNow(NLOCS, NLOCS)
  REAL( kind = dp ) :: UPatNow(*)
  integer :: NObsI, NObsP(NPT)
  REAL( kind = dp ), dimension(Nlocs) :: Ones
  REAL( kind = dp ), dimension(NLocs, NLocs) :: Eye
  ! The H matrix, a stack of diagonal matrices with the factor beta_1
  REAL( kind = dp ), allocatable :: Hk(:,:), Hdummy(:,:)
  REAL( kind = dp ), allocatable :: ErrIcovM(:,:)
  REAL( kind = dp ), allocatable :: BVec(:), Wgood(:)
  
  Ones = 1.0_dp
  Eye= 0.0_dp
  do i = 1, NLocs
    Eye(i,i) = 1.0_dp
  enddo
  
  ! Need to build:
  ! Hk, ErrIcovM, BVec
  ! Find number of prox and inst obs, allocate memory
  NObsI = NINT( SUM( UPatNow(1:NLOCS) ))
  DO i = 1, NPT
    NObsP(i) = NINT( SUM( UPatNow(NLOCS * i + 1: (NLOCS+1) * i) ))
  ENDDO
  !NObsTotal = SUM( NObsP) + NObsI
  ! I will first do this without removing the "missing obs" lines
  NObsTotal = NLOCS * (NPT + 1)
  ALLOCATE( Hk( NObsTotal, NLOCS) )
  ALLOCATE( Hdummy( NLOCS, NObsTotal) )
  ALLOCATE( ErrIcovM( NObsTotal, NObsTotal) )
  ALLOCATE( BVec(NObsTotal) )
  ALLOCATE( Wgood(NObsTotal) )
  ! find locations w/ inst + proxy observations this year and populate the matrices
  ! using the UPatNow and the SPREAD command
  ErrIcovM = 0.0_dp
  Hk(1:NLOCS,1:NLOCS) = SPREAD(UPatNow(1:NLOCS), 1, NLOCS)
  Hk(1:NLOCS,1:NLOCS) = MERGE(Hk(1:NLOCS,1:NLOCS), ErrIcovM(1:NLOCS,1:NLOCS), Eye == 1)
  ErrIcovM(1:NLOCS,1:NLOCS) = 1.0_dp / Cpars(6) * Eye
  BVec(1:NLOCS) = 0.0_dp
  DO I = 1, NPT
    BIDX = NLOCS * I + 1
    EIDX = (I+1) * NLOCS
    Hk(BIDX:EIDX,1:NLOCS) = SPREAD(UPatNow(BIDX:EIDX)* CProx( I, 4, :), 1, NLOCS) 
    Hk(BIDX:EIDX,1:NLOCS) = MERGE(Hk(BIDX:EIDX,1:NLOCS), ErrIcovM(BIDX:EIDX,BIDX:EIDX), Eye == 1)
    ErrIcovM(BIDX:EIDX,BIDX:EIDX) = Eye * SPREAD(UPatNow(BIDX:EIDX) / CProx( I, 2, :), 1, NLOCS)
    BVec(BIDX:EIDX) = CProx( I, 3, :) * UPatNow(BIDX:EIDX)
  ENDDO
  T_MATRIX(1:NLOCS, NYEARS) = Cpars(1)*(T_MATRIX(1:NLOCS, NYEARS -1) ) + (1-Cpars(1))*Cpars(2)*Ones
  Wgood(1:NObsTotal) = DATA_ALL(1:NObsTotal,NYEARS) - BVec(1:NObsTotal)
  CALL DGEMM( 'T', 'N', NLOCS, NObsTotal, NObsTotal, 1.0_dp, Hk, NObsTotal, &
    ErrIcovM, NObsTotal, 0.0_dp, Hdummy, NLOCS)
  CALL DGEMV( 'N', NLOCS, NLOCS, 1.0_dp/Cpars(3), Ciscm, NLOCS, T_MATRIX(1:NLOCS, NYEARS), 1, 0.0_dp, Ones, 1)
  CALL DGEMV( 'N', NLOCS, NObsTotal, 1.0_dp, Hdummy, NLOCS, Wgood, 1, 1.0_dp, Ones, 1)
  CALL DGEMV( 'N', NLOCS, NLOCS, 1.0_dp, PostCovNow, NLOCS, Ones, 1, 0.00_dp, T_MATRIX(1:NLOCS, NYEARS), 1)
  CALL getzig( Ones, NLOCS)
  CALL DGEMV( 'N', NLOCS, NLOCS, 1.0_dp, CholPostCovNow, NLOCS, Ones, 1, 1.0_dp, T_MATRIX(1:NLOCS, NYEARS), 1)
  
  DEALLOCATE( Hk)
  DEALLOCATE( Hdummy)
  DEALLOCATE( ErrIcovM)
  DEALLOCATE( BVec)
  DEALLOCATE( Wgood)

END SUBROUTINE TlastUpdater

! -----------------------------------------------------------------------------
! The full time updater step / driver for the updates
! -----------------------------------------------------------------------------
SUBROUTINE T_Updater( PriorT0, Cpars)

  !USE TUPDATER_DATA

  IMPLICIT INTEGER (I-K)

  REAL( kind = dp ) :: PriorT0(2)
  REAL( kind = dp ) :: Cpars(*)
  REAL( kind = dp ), dimension(1:NLOCS*(1+NPT) ) :: UPatNow
  REAL( kind = dp ) :: CCov(NLOCS, NLOCS)
  REAL( kind = dp ) :: UCCov(NLOCS, NLOCS)
  integer :: PatIDX, NObsTotal, TotalObs

  REAL( kind = dp ) :: NObsI( NumPat), NObsP( NPT, NumPat)
  REAL( kind = dp ), dimension(NLocs, NLocs) :: Eye, Dummy
  ! The Hk matrix space will have to contain all the products Hk^T %*% Xi^-1
  ! actually this is Hk^T !
  ! First it will be populated with H for each pattern, 
  ! then it will be converted to CSR
  ! then it will be multiplied
  REAL( kind = dp ), allocatable :: Hdummy(:,:)
!  REAL( kind = dp ), allocatable :: Hk(:)
!  integer, dimension(1:(NLOCS + 1)*NumPat):: Hkia
!  integer, allocatable          :: Hkja(:)
  ! begin and end indices for addressing the different submatrices of Hk and Xi
!  integer :: BIDX, EIDX, HAdd( 0:NumPat)
  integer :: BIDX, EIDX
  TYPE(CSCMatrix)  :: HMat( NumPat)
  ! integer error locator
  integer :: IERR
  REAL( kind = dp ), allocatable :: ErrIcovMDiag(:)

  ! create diagonal matrices to use e.g. as MASK
  Eye= 0.0_dp
  do i = 1, NLocs
    Eye(i,i) = 1.0_dp
  enddo
  
  NObsTotal = NLOCS * (NPT + 1)
  ! iterate over all different pattern types:
  ! Need to build the Hk and Xi^-1 matrices, convert Hk to spares representation
  ! and (finally, cf. next step) pass to TkUpdater
  ! There is no need for Xi^-1 in CSR as we can use AMUDIA or DIAMUA to multiply
  ! a diagonal matrix with a CSR one
  !
  ! Create the Xi^-1 (ErrIcovM) Matrix in diagonal representation
  ! needs additionally information on 
  !   number of diagonals (ndiag) = 1
  !   dimension of matrix (n) = NLOCS * (NPT + 1) = NObsTotal
  ! if real DIA format is needed. However, the AMUDIA subroutine can do without
  ! this information.
  ALLOCATE( ErrIcovMDiag(NObsTotal) )
  ErrIcovMDiag(1:NLOCS) = 1.0_dp / Cpars(6)
  DO i = 1, NPT
    BIDX = NLOCS * I + 1
    EIDX = (I+1) * NLOCS
    ErrIcovMDiag(BIDX:EIDX) = 1.0_dp / CProx( I, 2, :)
  ENDDO

  ! Now: 
  ! - Create Hk for all patterns
  ! - convert to CSR
  ! - multiply with Xi^-1
  TotalObs = NINT( SUM(UPatterns(1:NLOCS*(1+NPT),1:NumPat)))
  ALLOCATE( Hdummy(NLOCS, NObsTotal))
!  ALLOCATE( Hk( TotalObs) )
!  ALLOCATE( Hkja( TotalObs))
  ! IA is #of rows*#patterns + #of patterns, at each #rows+1 the array size of 
  ! Hk and Hkja is stored

!  HAdd(0) = 1
  DO J = 1, NumPat
    UPatNow = UPatterns(1:NObsTotal, J)
    ! Need to build Hk
    ! Find number of prox and inst obs
    ! Store this in the voctors NObsI and NObsP
    ! Store also in the address vector Hadd for the Hk matrix addresses
    HMat( J)%NNZ  = NINT( SUM (UPatNow) )
    HMat( J)%Nrow = NLOCS
    ALLOCATE( HMat( J)%a( HMat( J)%NNZ) )
    ALLOCATE( HMat( J)%ja( HMat( J)%NNZ) )
    ALLOCATE( HMat( J)%ia( HMat( J)%Nrow + 1) )
!    HAdd(J) = HAdd(J-1) + INT(SUM (UPatNow) )
    NObsI( J) = SUM( UPatNow(1:NLOCS) )
    ! find locations w/ inst + proxy observations this year and populate the matrices
    ! spread creates a matrix from a vector by repeating the vector for each col (or row)
    Hdummy(1:NLOCS, 1: NLOCS) = SPREAD(UPatNow(1:NLOCS), 1, NLOCS)
    ! the next line selects the "allowed" elements on the main diagonal
    Hdummy(1:NLOCS, 1: NLOCS) = MERGE(Hdummy(1:NLOCS, 1: + NLOCS), Eye, Eye == 1)
    DO I = 1, NPT
      BIDX = NLOCS * I + 1
      EIDX = (I+1) * NLOCS
      NObsP( I, J) = SUM( UPatNow( BIDX: EIDX) )
      Hdummy(1:NLOCS, BIDX:EIDX) = SPREAD(UPatNow( BIDX: EIDX)*CProx( I, 4, :), 1, NLOCS)
      Hdummy(1:NLOCS, BIDX:EIDX) = MERGE(Hdummy(1:NLOCS, BIDX:EIDX), Eye, Eye == 1)
    ENDDO
    ! convert to sparse representation
    ! input is the Hdummy array, output is a portion of the Hk and Hkia, Hkja matrices
    CALL dnscsr(NLOCS, NObsTotal, HMat( J)%NNZ, Hdummy(1:NLOCS, 1:NObsTotal), NLOCS, &
      HMat( J)%a, HMat( J)%ja, HMat( J)%ia, IErr)
    ! multiply with ErrIcovMDia using AMUDIA
    CALL amudia(NLOCS, 1, HMat( J)%a, HMat( J)%ja, HMat( J)%ia, ErrIcovMDiag(1:NObsTotal), &
      HMat( J)%a, HMat( J)%ja, HMat( J)%ia)
  ENDDO
  ! The matrix Hk should now contain all matrices HXI_k = (Hk^T %*% Xi^-1) 
  
  PatIDX = PatByYear(NYEARS)
  !  CCov = Cscm - Cscm %*% PhiList( PatIdx) %*% Cscm
  !PRINT *, "PHILIST"
  !PRINT *, PhiList( PatIDX)%NNZ
  !PRINT *, "PHILIST Nrow"
  !PRINT *, PhiList( PatIDX)%Nrow
  !PRINT *, PhiList( PatIDX)%a(1: PhiList( PatIDX)%NNZ) 
  !PRINT *, PhiList( PatIDX)%ja(1: (PhiList( PatIDX)%NNZ+1) )
  !PRINT *, PhiList( PatIDX)%ia(1: (PhiList( PatIDX)%Nrow +1)) 
  CCov = Cscm*Cpars(3)/(1+Cpars(1)**2)
  CALL amuxmat( NLOCS , NLOCS, NLOCS, CCov, Dummy, &
    PhiList( PatIdx)%a, PhiList( PatIdx)%ja, PhiList( PatIdx)%ia)
  CALL DGEMM( 'N', 'N', NLocs, NLocs, NLocs, -Cpars(3)/(1+Cpars(1)**2), &
    Cscm, NLocs, Dummy, NLocs, 1.0_dp, CCov, NLocs)
!  CCov = CCovArray(1:NLOCS, 1:NLOCS, PatIDX)
  UCCov = CCov
  ! While the code says "upper chol of the cov", we need the "lower triangle".
  ! This is a remnant of interfacing with R.
  CALL DPOTRF( 'L', NLOCS, UCCov, NLocs, Info) 
  DO I = 1, NLOCS-1
    DO J = I+1, NLOCS
      UCCov(I,J) = 0.0
    ENDDO
  ENDDO
!  UCCov = UCCovArray(1:NLOCS, 1:NLOCS, PatIDX)
!  PRINT *, PatIDX
!  PRINT *, PhiList( PatIdx)%NNZ
!  PRINT *, "UCCovArray(1:NLOCS, 1:NLOCS, PatIDX)"
!  WRITE (*, '(30(30F6.1))') SUM( (UCCov - UCCovArray(1:NLOCS, 1:NLOCS, PatIDX))*(UCCov - UCCovArray(1:NLOCS, 1:NLOCS, PatIDX)) )
  CALL TlastUpdater(Cpars, UPatterns(1:NLOCS*(1+NPT),PatIDX), CCov, UCCov)
  ! Loop pver all years
  DO I = NYEARS-1, 1, -1
    PatIDX = PatByYear(I)
!    PRINT *, PatIDX
!    PRINT *, PhiList( PatIdx)%NNZ
  CCov = Cscm*Cpars(3)/(1+Cpars(1)**2)
  CALL amuxmat( NLOCS , NLOCS, NLOCS, CCov, Dummy, &
    PhiList( PatIdx)%a, PhiList( PatIdx)%ja, PhiList( PatIdx)%ia)
  CALL DGEMM( 'N', 'N', NLocs, NLocs, NLocs, -Cpars(3)/(1+Cpars(1)**2), &
    Cscm, NLocs, Dummy, NLocs, 1.0_dp, CCov, NLocs)
  !CCov = CCovArray(1:NLOCS, 1:NLOCS, PatIDX)
    UCCov = CCov
    CALL DPOTRF( 'L', NLOCS, UCCov, NLocs, Info) 
    DO K = 1, NLOCS-1
      DO J = K+1, NLOCS
        UCCov(K,J) = 0.0
      ENDDO
    ENDDO
!  UCCov = UCCovArray(1:NLOCS, 1:NLOCS, PatIDX)
!  PRINT *, "CCovArray(1:NLOCS, 1:NLOCS, PatIDX)"
!  WRITE (*, '(30(30F6.1))') SUM( (CCov - CCovArray(1:NLOCS, 1:NLOCS, PatIDX))*(CCov - CCovArray(1:NLOCS, 1:NLOCS, PatIDX)) )
    CALL TkUpdater_sparse(I, Cpars, HMat( PatIDX), CCov, UCCov)
  ENDDO
  CALL T0Updater(PriorT0, Cpars)

  DEALLOCATE( ErrIcovMDiag)
  DEALLOCATE( Hdummy)


END SUBROUTINE T_Updater

END MODULE TUPDATER_DATA

!--------------------------------------------------------------------------------
! INTERFACE to the TUPDATER_DATA MODULE
!--------------------------------------------------------------------------------

SUBROUTINE TUpdater( PriorT0, Cpars)
  ! call the Temp updater from the module
  USE TUPDATER_DATA
  
  REAL( kind = dp ) :: PriorT0(2)
  REAL( kind = dp ) :: Cpars(*)
  
  CALL T_Updater( PriorT0, Cpars)

END SUBROUTINE TUpdater

SUBROUTINE init_updater( INITIAL_T, CProxCurr, ProxData, UPat, PbyY, NPat, currinvcovmat, NT, NL, NP, ZIGSEED, CMT)
  ! initialise the data
  USE Ziggurat
  USE TUPDATER_DATA
  
  IMPLICIT NONE
  
  INTEGER :: ZIGSEED, NT, NL, NP, NPat
  REAL( kind = dp ) :: INITIAL_T( NL, NT + 1), currinvcovmat(NL, NL), ProxData(NL*(NP+1),NT)
  REAL( kind = dp ) :: UPat(1:NL*(NP+1), 1:NPat)
  INTEGER :: PbyY(1:NT), I, Info
  REAL( kind = dp ) :: CProxCurr(1:NP, 1:5, 1:NL)
  INTEGER :: CMT
  
  NYEARS = NT
  NLOCS = NL
  NPT = NP
  NumPat = NPat

  IF ( ZIGSEED == 0 ) ZIGSEED = 123
  CALL zigset( ZIGSEED)
  
  IF ( .NOT. ALLOCATED( T_MATRIX) )  ALLOCATE ( T_MATRIX(1 : NLOCS, 0 : NYEARS))
  IF ( .NOT. ALLOCATED( DATA_ALL) )  ALLOCATE ( DATA_ALL(1 : NLOCS*(NPT + 1), 1 : NYEARS))
  IF ( .NOT. ALLOCATED( Ciscm) )     ALLOCATE ( Ciscm(1 : NLOCS, 1 : NLOCS))
  IF ( .NOT. ALLOCATED( Cscm) )      ALLOCATE ( Cscm(1 : NLOCS, 1 : NLOCS))
  IF ( .NOT. ALLOCATED( UPatterns) ) ALLOCATE ( UPatterns (1 : NLOCS*(NPT+1), 1:NumPat) )
  IF ( .NOT. ALLOCATED( PatByYear) ) ALLOCATE ( PatByYear( NYEARS))
  IF ( .NOT. ALLOCATED( CProx) )     ALLOCATE ( CProx(1:NP, 1:5, 1:NLOCS) )
  IF ( .NOT. ALLOCATED( PhiList ) )  ALLOCATE ( PhiList( 1:NumPat) )

  T_MATRIX = INITIAL_T
  DATA_ALL = ProxData
  Ciscm    = currinvcovmat
  Cscm = 0_dp
  DO I = 1, NLOCS
    Cscm(I,I) = 1E0
  END DO
  CALL DPOSV('U', NLOCS, NLOCS, currinvcovmat, NLOCS, Cscm, NLOCS, Info)
  UPatterns = UPat
  PatByYear = PbyY
  CProx = CProxCurr
  CovMatType = CMT

END SUBROUTINE init_updater


SUBROUTINE updatePhiList( PhiMat_a, PhiMat_ja, PhiMat_ia, PhiMat_NNZ, PhiMat_Nrow, thisPat)
  ! Update the list of the PhiMatrices. It can be necessary to reallocate things 
  ! first, but this needs to be done elsewhere.
  ! Note: MOVE_ALLOC is not a solution, since this interfaces with R.
  USE TUPDATER_DATA

  IMPLICIT NONE

  INTEGER, INTENT( IN)      :: PhiMat_ia(*), PhiMat_ja(*), PhiMat_NNZ, PhiMat_Nrow
  REAL( kind = dp ), INTENT( IN) :: PhiMat_a(*)
  INTEGER, INTENT( IN)      :: thisPat

  IF (ALLOCATED( PhiList( thisPat)%a))  DEALLOCATE ( PhiList( thisPat)%a)
  IF (ALLOCATED( PhiList( thisPat)%ja)) DEALLOCATE ( PhiList( thisPat)%ja)
  IF (ALLOCATED( PhiList( thisPat)%ia)) DEALLOCATE ( PhiList( thisPat)%ia)
  
  ALLOCATE( PhiList( thisPat)%a( PhiMat_NNZ))
  ALLOCATE( PhiList( thisPat)%ja( PhiMat_NNZ) )
  ALLOCATE( PhiList( thisPat)%ia( PhiMat_Nrow + 1) )

  ! the spam representation counts differently, everything is long by one.
  ! Needs to be done correctly when calling this function!!
  PhiList( thisPat)%NNZ = PhiMat_NNZ
  PhiList( thisPat)%Nrow = PhiMat_Nrow
  PhiList( thisPat)%a = PhiMat_a(1: PhiMat_NNZ )
  PhiList( thisPat)%ia = PhiMat_ia(1: (PhiMat_Nrow +1) )
  PhiList( thisPat)%ja = PhiMat_ja(1: (PhiMat_NNZ +1 ))

  ! looks sensible 2015-10-15
  !PRINT *, "PATTERN"
  !PRINT *, thisPat
  !PRINT *, PhiList( thisPat)%NNZ
  !PRINT *, PhiList( thisPat)%Nrow
  !PRINT *, PhiList( thisPat)%a(1: PhiMat_NNZ) 
  !PRINT *, PhiList( thisPat)%ja(1: (PhiMat_NNZ+1) )
  !PRINT *, PhiList( thisPat)%ia(1: (PhiMat_Nrow +1)) 

END SUBROUTINE updatePhiList

SUBROUTINE newPhiList( NPat)

  USE TUPDATER_DATA

  IMPLICIT NONE

  INTEGER, INTENT( IN)      :: NPat

  IF (ALLOCATED( PhiList ) )  DEALLOCATE ( PhiList)
  ALLOCATE ( PhiList( 1:NPat) )

END SUBROUTINE newPhiList

SUBROUTINE updateUPat( UPat, PbyY, NPat)
  ! populate the new pattern data sets
  USE TUPDATER_DATA
  
  IMPLICIT NONE
  
  INTEGER :: NPat
  REAL( kind = dp ) :: UPat(1:NLOCS*(NPT+1), 1:NPat)
  INTEGER :: PbyY(1:NYEARS)
  
  NumPat = NPat
  
  DEALLOCATE( UPatterns)
  IF ( .NOT. ALLOCATED( UPatterns) ) ALLOCATE ( UPatterns (1 : NLOCS*(NPT+1), 1:NumPat) )

  UPatterns = UPat
  PatByYear = PbyY

END SUBROUTINE updateUPat

SUBROUTINE updateSglProx( ProxData, PNum, PIdx)
  ! change the data of a single proxy
  USE Ziggurat
  USE TUPDATER_DATA
  
  IMPLICIT NONE
  
  INTEGER :: PNum, PIdx
  REAL( kind = dp ) :: ProxData(1: NYEARS)
  
  DATA_ALL( PNum*NLOCS + PIdx, 1: NYEARS) = ProxData

  ! PatByYear = PbyY also needs to be updated!

END SUBROUTINE updateSglProx

SUBROUTINE set_ciscm( currinvcovmat)
  ! update the inverse covariance matrix
  USE TUPDATER_DATA

  REAL( kind = dp ) :: currinvcovmat(NLOCS,NLOCS)
  Ciscm    = currinvcovmat

END SUBROUTINE set_ciscm

SUBROUTINE get_ciscm( currinvcovmat)
  ! return the inverse covariance matrix
  USE TUPDATER_DATA

  REAL( kind = dp ) :: currinvcovmat(NLOCS,NLOCS)
  currinvcovmat = Ciscm

END SUBROUTINE get_ciscm

SUBROUTINE get_cscm( currcovmat)
  ! return the inverse covariance matrix
  USE TUPDATER_DATA

  REAL( kind = dp ) :: currcovmat(NLOCS,NLOCS)
  currcovmat = Cscm

END SUBROUTINE get_cscm

SUBROUTINE cleanup_updater()
  ! deallocate everything
  USE TUPDATER_DATA
  
  IMPLICIT NONE
 
  DEALLOCATE ( T_MATRIX)
  DEALLOCATE ( DATA_ALL)
  DEALLOCATE ( Ciscm)
  DEALLOCATE ( Cscm)
  DEALLOCATE ( UPatterns )
  DEALLOCATE ( PatByYear )
  DEALLOCATE ( PhiList )
  DEALLOCATE ( CProx)

END SUBROUTINE cleanup_updater

SUBROUTINE GetTMatrix( TMat)
  !
  ! returns the current temperature matrix
  USE TUPDATER_DATA

  IMPLICIT NONE
  REAL( kind = dp ), DIMENSION(1:NLOCS, 0:NYEARS) :: TMat

  TMat = T_MATRIX

END SUBROUTINE GetTMatrix


! -----------------------------------------------------------------------------
! 
! The parameters Theta=[alpha, mu, sigma, phi, sigma_I, sigma_P, beta_{0,1...}

SUBROUTINE Alpha_Updater( alphaprior, Cpars)
  !inputs are: [PRIORS.alpha, CURRENT_PARS ]
  !
  !#UPDATES the AR(1) coefficient alpha in the main BARCAST code
  !
  ! TRUNCATED NORMAL
  
  USE Ziggurat
  USE TUPDATER_DATA

  IMPLICIT NONE
  REAL( kind = dp ) :: TMatLessMean(1:NLOCS, 0:NYEARS)
  REAL( kind = dp ) :: AlphaMean, AlphaVar
  REAL( kind = dp ) :: Dummy( 1:NLOCS, 0:NYEARS)
  REAL( kind = dp ) :: alpha(1), alphaprior(*), Cpars(*)
  
  !calculate the inverse posterior variance.
  !first lay down the matrix of the deviatons of the temperature values from
  !the dependant means. 
  TMatLessMean(1:NLOCS, 0:NYEARS) = T_MATRIX(1:NLOCS, 0:NYEARS) - Cpars(2)
  
  CALL DGEMM('N', 'N', NLOCS, NYEARS, NLOCS, 1.0/Cpars(3), Ciscm, NLOCS, TMatLessMean(1:NLOCS,0:NYEARS-1), NLOCS, &
    0.0_dp, Dummy(1:NLOCS,0:NYEARS-1), NLOCS)
 
  AlphaVar = 1.0_dp/(SUM( TMatLessMean(1:NLOCS,0:NYEARS-1) * Dummy(1:NLOCS,0:NYEARS-1) ) + 1.0/alphaprior(2))

  AlphaMean = ( SUM( TMatLessMean(1:NLOCS,1:NYEARS) * Dummy(1:NLOCS,0:NYEARS-1) ) + alphaprior(1)/alphaprior(2))
  AlphaMean = AlphaVar * AlphaMean
  
  CALL getzig( alpha, 1)
  alpha = AlphaMean + alpha * AlphaVar**.5
  Cpars(1) = alpha(1)
  IF (Cpars(1) < 0) Cpars(1) = 0

END SUBROUTINE Alpha_Updater
  
SUBROUTINE Mu_Updater( muprior, Cpars)
  
  USE Ziggurat
  USE TUPDATER_DATA

  IMPLICIT NONE
  REAL( kind = dp ) :: TMatDiff(1:NLOCS, 1:NYEARS)
  REAL( kind = dp ) :: MuMean, MuVar
  REAL( kind = dp ) :: Ones_L(1:NLOCS), Ones_Y(1:NYEARS)
  REAL( kind = dp ) :: muprior(*), Cpars(*), DummyVec(1:NLOCS)

  ! function New_mu=Mu_Updater_vNM(mu_p, T_Mat, C_pars, C_iscm);
  !
  !inputs are: [PRIORS.mu, CURRENT_PARS] 
  !
  !#UPDATES the constant mean coefficient of the AR(1) process in the
  !#main BARCAST code
  !
  ! NORMAL

  Ones_Y(1:NYEARS) = 1

  MuVar = 1.0_dp/(1.0_dp/muprior(2) + 1.0_dp/Cpars(3) * SUM(Ciscm) * NYEARS*(1-Cpars(1))**2 )
  
  !calculate the posterior mean:
  ! first the difference, T_mat less shifted (in time) and scaled (by
  ! a*N_T*(1-alpha) T_mat, then summed across time:
  TMatDiff = T_MATRIX(1:NLOCS, 1:NYEARS) - Cpars(1)*T_MATRIX(1:NLOCS, 0:NYEARS-1)
  CALL DGEMV('N', NLOCS, NYEARS, 1.0_dp, TMatDiff, NLOCS, Ones_Y, 1, 0.0_dp, Ones_L, 1)
  CALL DGEMV('N', NLOCS, NLOCS, 1.0_dp/Cpars(3), Ciscm, NLOCS, Ones_L, 1, 0.0_dp, DummyVec, 1)

  MuMean = MuVar * ( muprior(1)/muprior(2) + (1-Cpars(1)) * SUM(DummyVec)  )
  
  CALL getzig( Cpars(2), 1)
  Cpars(2) = MuMean + Cpars(2) * MuVar**.5
END SUBROUTINE Mu_Updater

SUBROUTINE Sigma2_Updater( sigma2prior, Cpars)

  USE Ziggurat
  USE random
  USE TUPDATER_DATA

  IMPLICIT NONE
  REAL              :: alpha
  REAL( kind = dp ) :: sigma2prior(*), Cpars(*), beta
  REAL( kind = dp ) :: TDiffMat(1:NLOCS, 1:NYEARS), Dummy(1:NLOCS, 1:NYEARS)

  alpha = NYEARS * NLOCS /2.0 + sigma2prior(1)

  TDiffMat = T_MATRIX(1:NLOCS,1:NYEARS) - Cpars(1) * T_MATRIX(1:NLOCS,0:NYEARS - 1) - &
    (1-Cpars(1) ) * Cpars(2)
  CALL DGEMM( 'N', 'N', NLOCS, NYEARS, NLOCS, 1.0_dp, Ciscm, NLOCS, TDiffMat, NLOCS, 0.0_dp, Dummy, NLOCS)
  TDiffMat = TDiffMat * Dummy
  beta = sigma2prior(2) + 0.5*SUM(TDiffMat)
  !
  ! concerning Hobbits (the Inv-Gamma distribution)
  ! for a random varaible X the following holds:
  ! 1/X ~ Gamma (a, 1/b) <-> X ~ InvGamma(a, b)
  ! X ~ InvGamma (a, b) <-> k*X ~ InvGamma(a, k*b)
  ! So:
  ! Y ~ Gamma(a, 1)
  ! X = 1/Y ~ InvGamma(a, 1)
  ! sigma2 = beta*X ~ InvGamma(a, beta)
  ! sigma2 ~ beta/Y
  Cpars(3) = MIN( beta / random_gamma( alpha , .TRUE.), 5.0E0)

END SUBROUTINE Sigma2_Updater

SUBROUTINE Tau2I_Updater(tau2Iprior, Cpars, MInst)
  
  USE random
  USE TUPDATER_DATA

  IMPLICIT NONE

  REAL( kind = dp )  ::  tau2Iprior(*), Cpars(*), MInst
  REAL( kind = dp )  ::  TMatDiff(1:NLOCS, 1:NYEARS)
  REAL               ::  TauAlpha
  REAL( kind = dp )  ::  TauBeta
  !
  !UPDATES the variance of the instrumental error in the main BARCAST code
  !
  ! INVERSE GAMMA
  !
  
  !need to find the sum of the squared residuals between the Inst
  !observations and the corresponding Temp values:
  !extract the Inst part of the T,D matrices:
  !also get rid of the time=0 value from the temp mat:
  !we don't actually need the HH_Select matrix: just take the difference
  !between the Temperature matrix and ther Inst Obs matrix:: wherever there
  !is a NaN in the Data matrix, there will be a NaN in the difference. Then
  !find all non NaN entries, square them, add them up.
  WHERE( DATA_ALL(1:NLOCS, 1:NYEARS) == -99 )
    TMatDiff = 0;
  ELSEWHERE  
    TMatDiff = T_MATRIX(1:NLOCS, 1:NYEARS) - DATA_ALL(1:NLOCS, 1:NYEARS)
  END WHERE
  
  !find the sum of the squares of the residuals::
  
  !we can now calculate the postertior beta parameter:
  TauBeta = SUM(TMatDiff**2)/2.0_dp + tau2Iprior(2)
  !can then calculate the first parameter, alpha, for the posterior inv-gamma
  !dist.:
  TauAlpha = MInst / 2.0_dp + tau2Iprior(1)
  !print(paste("alpha post", alpha.post, "; beta post", beta.post))
  
  !make the draw:
  Cpars(6) = MIN( TauBeta/ random_gamma( TauAlpha , .TRUE.), 1.0_dp)

!  print *, TauBeta
!  print *, TauAlpha
  
END SUBROUTINE Tau2I_Updater


SUBROUTINE Phi_Updater( phiprior, Cpars, inv_scm, DistMat, MHParsLogphi)

  USE random
  USE TUPDATER_DATA

  IMPLICIT INTEGER(I-K)

  REAL( kind = dp )  ::  phiprior(*), Cpars(*), MHParsLogphi(*)
  REAL( kind = dp ), DIMENSION(1:NLOCS,1:NYEARS)  ::  TMatDiff, W_now, W_prop
  REAL( kind = dp ), DIMENSION(1:NLOCS,1:NLOCS)   ::  scm_prop, scm_now, AT_prop, AT_now, &
      DistMat, Eye, inv_scm
  REAL( kind = dp ), DIMENSION(1:NLOCS**2)        ::  scm_vec
  REAL( kind = dp )  ::  w_squared_now, w_squared_prop, J_par, lgphi_prop
  REAL( kind = dp )  ::  acceptance, Met_ratio, log_met_ratio, tester, VAE_prop
  REAL( kind = dp ), ALLOCATABLE :: lgphiVAE(:,:)
  INTEGER  :: INFO, N_Its
  INTEGER, DIMENSION(1:NLOCS, 1:NLOCS)  ::  IEye

  !UPDATES the range parameter in the spatial covaraince matrix for the main
  !BARCAST code
  Eye = 0.0_dp
  DO I = 1, NLOCS
    Eye(I,I) = 1.0_dp
  ENDDO
  IEye = 0
  DO I = 1, NLOCS
    IEye(I,I) = 1
  ENDDO

  J_par = MHparsLogphi(1)**.5;
  N_Its = floor(MHparsLogphi(2) );

  !find the matrix of adjusted temeprature differences that appear in the
  !quadratic forms. 
  TMatDiff = T_MATRIX(1:NLOCS,1:NYEARS)- Cpars(1)*T_MATRIX(1:NLOCS,0:NYEARS-1) - (1-Cpars(1))*Cpars(2)
  !label the spatial correlation matrix and inverse:
  IF ( CovMatType .EQ. 10) THEN
    CALL MaternCov2d( 2.5_dp, Cpars(4), DistMat, scm_Now, NLOCS)
  ELSE
    scm_Now = exp(-Cpars(4)*DistMat)
  ENDIF
  AT_now = scm_Now
  ! Cholesky decomposition
  CALL DPOTRF( 'U', NLOCS, AT_now, NLOCS, INFO)
  
  IF ( INFO .ne. 0) THEN
    print *, "DPOTRF ERROR!"
    print *, INFO
  ENDIF
  ! Don't need to clear the lower triangular part of AT. The DTRTRS function
  ! references the upper ('U') part only anyways
  !DO I = 2, NLOCS
  !  DO J = 1, I-1
  !    AT_now(I,J) = 0.0
  !  ENDDO
  !ENDDO
  w_now = TMatDiff
  CALL DTRTRS('U', 'T', 'N', NLOCS, NYEARS, AT_now, NLOCS, w_now, NLOCS, INFO )
  IF ( INFO .ne. 0) THEN
    print *, "DTRTRS ERROR!"
    print *, INFO
  ENDIF
  w_squared_now = (sum(w_now**2))
  !we will also make a "proposed" versions, and replace the "now" versions as
  !necessary. This is an attempt to avoid unnecessary matrix inversions. 

  !Create an empty matrix, N_Its+1 by 2. Left Column: log_phi values, starting
  !with the log of the inputed phi value and ending with the value after N_Its of the
  !Metropolis algorithm. Right column: value of the argument of the
  !exponenent given the value of Beta_1. 

  ALLOCATE(lgphiVAE(1:N_Its+1,2))
  lgphiVAE(1,1) = log(Cpars(4))

  !Find the argument of the exponent for the current logphi value, once more
  !using the matrix then array multiplication trick:
  lgphiVAE(1,2) = -(lgphiVAE(1,1)-phiprior(1))**2/(2*phiprior(2))  - w_squared_now / (2*Cpars(3))
  
  acceptance = 0;
  DO K = 1, N_Its
    lgphi_prop = rnor()
    lgphi_prop = lgphi_prop * J_par + lgphiVAE(K,1) 
    !calculate the spatial correlation matrix using this phi:
    IF ( CovMatType .EQ. 10) THEN
      CALL MaternCov2d( 2.5_dp, exp(lgphi_prop), DistMat, scm_prop, NLOCS)
    ELSE
      scm_prop = exp(-exp(lgphi_prop)*DistMat)
    ENDIF

    AT_prop = scm_prop
    CALL DPOTRF( 'U', NLOCS, AT_prop, NLOCS, INFO)
    IF ( INFO .ne. 0) THEN
      print *, "DPOTRF ERROR!"
      print *, INFO
    ENDIF

    w_prop = TMatDiff
    CALL DTRTRS('U', 'T', 'N', NLOCS, NYEARS, AT_prop, NLOCS, w_prop, NLOCS, INFO )
    IF ( INFO .ne. 0) print *, "ERROR!"
    w_squared_prop = sum(w_prop**2)

    !calculate the value of the arg of the exponent for this proposed
    !log_phi:
    VAE_prop = -(lgphi_prop - phiprior(1))**2/(2*phiprior(2))  - w_squared_prop / (2*Cpars(3))

    !calculate the log metropolis ratio, using logs to avoid NaNs and other
    !foolishness. 

    log_Met_ratio = VAE_prop - lgphiVAE(K,2) + (NYEARS)* &
      (log(PRODUCT(AT_now, MASK = IEye == 1)) - log(PRODUCT(AT_prop, MASK = IEye == 1)) )
                        ! 2*NYEARS/2 because of AT just being the "root" of R

    Met_ratio = exp(log_Met_ratio);
    !decide on the next value of logphi and arg of exp, and update the values of (i)scm_Now:
    IF (Met_ratio > 1) THEN
      !if ratio greater than 1, accept:
      lgphiVAE(K+1, 1) = lgphi_prop
      lgphiVAE(K+1, 2) = VAE_prop
      scm_Now = scm_prop
      AT_now = AT_prop
      w_now = w_prop
      w_squared_now = w_squared_prop
      acceptance = acceptance + 1
    ELSE
      !else accept with probability equal to the ratio
      tester = uni()
      IF ( Met_ratio > tester) THEN
        lgphiVAE(K+1, 1) = lgphi_prop
        lgphiVAE(K+1, 2) = VAE_prop
        scm_Now = scm_prop
        AT_now = AT_prop
        w_now = w_prop
        w_squared_now = w_squared_prop
        acceptance = acceptance + 1
      ELSE
        lgphiVAE(K + 1,1:2) = lgphiVAE(K, 1:2)
            !no need to update the (i)scm_Now.
      ENDIF
    ENDIF
!    write (*, '(A)',advance='no') "."
  ENDDO
  Cpars(4) = EXP(lgphiVAE(N_Its+1 ,1) )
  Cpars(5) = 0
  Cscm = scm_Now
  Ciscm = Eye
  CALL DPOSV('U', NLOCS, NLOCS, scm_now, NLOCS, Ciscm, NLOCS, Info)
  inv_scm = Ciscm
  write (*, '(A, F7.2, E9.2 )'), "acceptance / last step:", acceptance/N_Its, EXP(lgphiVAE(N_Its+1 ,1) )
  DEALLOCATE(lgphiVAE)

END SUBROUTINE Phi_Updater

