program lsq
use printmat2
implicit none
!
!This program calculates least square solutions for a model with 2 effects.
!All the storage is dense, and solutions are obtained iteratively by
!Gauss-Seidel.
! This program is easily upgradable to any number of effects
!
real, allocatable:: xx(:,:),xy(:),sol(:) !storage for the equations
integer, allocatable:: indata(:) !storage for one line of effects
!integer,parameter:: neff=2,nlev(2)=(/2,3/) !number of effects and levels
real :: y ! observation value
integer :: neq,io,i,j ! number of equations and io-status
integer :: neff,addeff
real :: var_e,var_a
integer,allocatable :: nlev(:),address(:)
real,allocatable :: a(:,:)
character (40) :: parfile,datafile,pedfile,random

call read_parameter

neq=sum(nlev)
print*,'# equations =',neq
allocate (xx(neq,neq), xy(neq), sol(neq),indata(neff),address(neff))
xx=0; xy=0; sol=0
!
open(1,file=datafile)
!
do
   read(1,*,iostat=io)indata,y
   if (io.ne.0) exit
   call find_addresses
   do i=1,neff
      do j=1,neff
         xx(address(i),address(j))=xx(address(i),address(j))+1.0/var_e
      enddo
      xy(address(i))=xy(address(i))+y/var_e
   enddo
enddo

do i=1,neff
   if(i == addeff) then
      if(random == 'add') then
         call ainv(i)
         print*,'Additive relationship matrix'
         call printnice(a,'(20f10.4)')
      elseif(random == 'diag') then
         call diag(i)
      else
         print*,'* choose add or diag'
         stop
      endif
   endif
enddo
!
print*,'left hand side'
do i=1,neq
   print '(100f10.4)',xx(i,:)
enddo
!
print '( '' right hand side:'' ,100f10.4)',xy

call solve_dense_gs(neq,xx,xy,sol) !solution by Gauss-Seidel
open(99,file='solutions')
do i=1,neq
   write(99,*) sol(i)
enddo
print '( '' solution:'' ,100f10.4)',sol

contains

subroutine find_addresses
integer :: i
do i=1,neff
   address(i)=sum(nlev(1:i-1))+indata(i)
enddo
end subroutine

function address1(e,l)
! returns address for given level l of effect e and trait t
integer :: e,l,address1
address1= sum(nlev(1:e-1))+(l-1)+1
end function

subroutine ainv(eff)
integer :: eff,i,k,l,m,n,io,p(3),nped=0
real ::w(3),res(4),mendel

w=(/1., -.5, -.5/)
res=(/2., 4/3., 1., 0./)

do
   read(2,*,iostat=io) p
   if(io /= 0) exit
   nped=nped+1
enddo
print*,'# records in pedigree =',nped
allocate(a(nped,nped))
a=0
rewind 2
do
   read(2,*,iostat=io) p
   if (io /= 0) exit
   i=1
   if(p(2) == 0) i=i+1
   if(p(3) == 0) i=i+1
   mendel=res(i)
   do k=1,3
      do l=1,3
         if (p(k) /= 0 .and. p(l) /=0) then
            m=address1(eff,p(k))
            n=address1(eff,p(l))
            a(p(k),p(l))=a(p(k),p(l))+w(k)*w(l)*mendel
            xx(m,n)=xx(m,n)+w(k)*w(l)*mendel/var_a
         endif
      enddo
   enddo
enddo
end subroutine

subroutine diag(eff)
integer :: eff,i,m

do i=1,nlev(eff)
   m=address1(eff,i)
   xx(m,m)=xx(m,m)+1.0/var_a
enddo

end subroutine

subroutine read_parameter
integer :: i

print*,'Parameter file?'
read(*,'(a)') parfile
open(99,file=parfile)

read(99,'(a)') datafile
print*,'Data file: ',datafile

read(99,'(a)') pedfile
print*,'Pedigree file: ',pedfile

open(2,file=pedfile)

read(99,*) neff
print*,'# effects =',neff

allocate(nlev(neff))
read(99,*) nlev

do i=1,neff
   print*,i,'-th effect # levels =',nlev(i)
enddo

read(99,*) addeff
print*,'Additive effect =',addeff

read(99,'(a)') random
print*,'Random effect = ',random

read(99,*) var_a
print*,'Additive variance =',var_a

read(99,*) var_e
print*,'Residual variance =',var_e

end subroutine

end program lsq

subroutine solve_dense_gs(n,lhs,rhs,sol)
! finds sol in the system of linear equations: lhs*sol=rhs
! the solution is iterative by Gauss-Seidel
integer :: n
real :: lhs(n,n),rhs(n),sol(n),eps
integer :: round
!
round=0
do
   eps=0; round=round+1
   do i=1,n
      solnew=sol(i)+(rhs(i)-sum(lhs(i,:)*sol))/lhs(i,i)
      eps=eps+ (sol(i)-solnew)**2
      sol(i)=solnew
   end do
   print*,round,' rounds conv =',eps
   if (eps.lt. 1e-10) exit
end do
print*,'solutions computed in ',round,' rounds of iteration'
end subroutine
program BLUP1
implicit none
!
! As lsqmt but with support for random effects

!
integer,parameter::effcross=0,& !effects can be cross-classified
effcov=1 !or covariables
real, allocatable:: xx(:,:),xy(:),sol(:) !storage for the equations
integer, allocatable:: indata(:) !storage for one line of effects

!integer,parameter:: neff=2,& !number of effects
! nlev(2)=(/2,4/),& !number of levels
! ntrait=2,& !number of traits
integer,parameter:: ntrait=1,& !number of traits
miss=0 !value of missing trait/effect
integer :: neff,addeff
integer, allocatable :: nlev(:),effecttype(:),nestedcov(:),&
randomtype(:),randomnumb(:)
real :: r(ntrait,ntrait),& !residual (co)variance matrix
rinv(ntrait,ntrait) ! and its inverse
!integer :: effecttype(neff)=(/effcross, effcross/)
!integer :: nestedcov(neff) =(/0,0/)
real, allocatable :: weight_cov(:,:),g(:,:,:)

! Types of random effects
integer, parameter :: g_fixed=1,& ! fixed effect
g_diag=2, & ! diagonal
g_A=3, & ! additive animal
g_As=4 ! additive sire

!integer :: randomtype(neff),& ! status of each effect, as above
! randomnumb(neff) ! number of consecutive correlated effects
!real :: g(neff,20,20)=0 ! The random (co)variance matrix for each trait
! maximum size is 20 of trait x corr. effect combinations
real :: var_a,var_e
real :: y(ntrait) ! observation value
integer :: neq,io,i,j,k,l ! number of equations and io-status
integer,allocatable:: address(:,:) ! start and address of each effect

character (40) :: parfile,datafile,pedfile,random

call read_parameter

allocate(effecttype(neff),nestedcov(neff),weight_cov(neff,ntrait),&
randomtype(neff),randomnumb(neff),g(neff,ntrait,ntrait))
effecttype=0
nestedcov=0
!
neq=ntrait*sum(nlev)
print*,'# equations =',neq
allocate (xx(neq,neq), xy(neq), sol(neq),indata(neff*ntrait),&
address(neff,ntrait))
xx=0; xy=0; sol=0
!
! Assign random effects and G and R matrices

randomtype='g_fixed' !types of random effect
randomnumb=0 !number of correlated effects per effect
if(random == 'diag') then
randomtype(addeff)=g_diag
elseif(random == 'add') then
randomtype(addeff)=g_A
else
stop 'error'
endif
randomnumb(addeff)=1

!g(2,1:2,1:2)=reshape( (/2, -1, -1, 1/), (/2,2/))
!r=reshape( (/1, 0, 0, 1/), (/2,2/))
g(addeff,1,1)=var_a
r(1,1)=var_e

call setup_g ! invert G matrices

open(2,file=pedfile) !pedigree file
open(1,file=datafile) !data file
!

do
   read(1,*,iostat=io)indata,y
   if (io.ne.0) exit
   call find_addresses
   call find_rinv
   do i=1,neff
      do j=1,neff
         do k=1,ntrait
            do l=1,ntrait
               xx(address(i,k),address(j,l))=xx(address(i,k),address(j,l))+&
               weight_cov(i,k)*weight_cov(j,l)*rinv(k,l)
            enddo
         enddo
      enddo
      do k=1,ntrait
         do l=1,ntrait
            xy(address(i,k))=xy(address(i,k))+rinv(k,l)*y(l)*weight_cov(i,k)
         enddo
      enddo
   enddo
enddo
!
! Random effects' contributions
do i=1,neff
   select case (randomtype(i))
   case (g_fixed)
      continue ! fixed effect, do nothing
   case (g_diag)
      call add_g_diag(i)
   case (g_A, g_As)
      call add_g_add(randomtype(i),i)
   case default
      print*,'unimplemented random type',randomtype(i)
   endselect
enddo
!
print*,'left hand side'
do i=1,neq
   print '(100f5.1)',xx(i,:)
enddo
!
print '( '' right hand side:'' ,100f6.1)',xy
!
call solve_dense_gs(neq,xx,xy,sol) !solution by Gauss-Seidel
open(99,file='solutions')
do i=1,neq
   write(99,*) sol(i)
enddo
print '( '' solution:'' ,100f7.3)',sol

contains

subroutine find_addresses
integer :: i,j,baseaddr
do i=1,neff
   do j=1,ntrait
      if (datum(i,j) == miss) then !missing effect
         address(i,j)=0 !dummy address
         weight_cov(i,j)=0.0
         cycle
      endif
      baseaddr=sum(nlev(1:i-1))*ntrait+j !base address (start)
      select case (effecttype(i))
      case (effcross)
! address(i,j)=baseaddr+(datum(i,j)-1)*ntrait
         address(i,j)=address1(i,int(datum(i,j)),j)
         weight_cov(i,j)=1.0
      case (effcov)
         weight_cov(i,j)=datum(i,j)
         if (nestedcov(i) == 0) then
            address(i,j)=baseaddr
         elseif (nestedcov(i)>0 .and. nestedcov(i).lt.neff) then
            address(i,j)=baseaddr+(datum(nestedcov(i),j)-1)*ntrait
         else
            print*,'wrong description of nested covariable'
            stop
         endif
      case default
         print*,'unimplemented effect ',i
         stop
      end select
   enddo
enddo
end subroutine

function datum(ef,tr)
real:: datum
integer :: ef,tr
! calculates the value effect ef and trait tr
datum=indata(ef +(tr-1)*neff)
end function

subroutine find_rinv
! calculates inv(Q R Q), where Q is an identity matrix zeroed for
! elements corresponding to y(i)=miss
integer :: i,irank
! real:: w(10)
rinv=r
do i=1,ntrait
   if (y(i) == miss) then
      rinv(i,:)=0; rinv(:,i)=0
   endif
enddo
call ginv(rinv,ntrait,1e-5,irank)
end subroutine

subroutine setup_g
! inverts g matrices
! real :: w(20)
integer :: rank,j
do i=1,neff
   if (randomnumb(i).ne.0) then
      j=randomnumb(i)*ntrait
      call printmat('original G',g(i,:,:),j,j)
      call ginv(g(i,:,:),j,1e-5,rank)
      call printmat('inverted G',g(i,:,:),j,j)
   endif
enddo
end subroutine

function address1(e,l,t)
! returns address for given level l of effect e and trait t
integer :: e,l,t, address1
address1= sum(nlev(1:e-1))*ntrait+(l-1)*ntrait+t
end function

subroutine add_g_diag(eff)
! adds diagonal (IID) contributions to MME
integer :: eff, i,j,k,l,m,t1,t2
do i=1,nlev(eff)
   do j=0,randomnumb(eff)-1
      do k=0,randomnumb(eff)-1
         do t1=1,ntrait
            do t2=1,ntrait
               m=address1(eff+j,i,t1); l=address1(eff+k,i,t2)
               xx(m,l)=xx(m,l)+g(eff,t1+j*ntrait,t2+k*ntrait)
            enddo
         enddo
      enddo
   enddo
enddo
end subroutine

subroutine add_g_add(type,eff)
! generates contributions for additive sire or animal effect
integer :: type,eff,i,j,t1,t2,k,l,m,n,io,p(3)
real ::w(3),res(4),mendel
!
select case (type)
case (g_A)
   w=(/1., -.5, -.5/)
   res=(/2., 4/3., 1., 0./)
case (g_As)
   w=(/1., -.5, -.25/)
   res=(/16/11., 4/3., 16/15., 1./)
end select

do
   read(2,*,iostat=io) p
   if (io /= 0) exit
   i=1
   if(p(2) == 0) i=i+1
   if(p(3) == 0) i=i+1
   mendel=res(i)
   do i=0,randomnumb(eff) - 1
      do j=0,randomnumb(eff) - 1
         do t1=1,ntrait
            do t2=1,ntrait
               do k=1,3
                  do l=1,3
                     if (p(k) /= 0 .and. p(l) /=0) then
                        m=address1(eff+i,p(k),t1)
                        n=address1(eff+j,p(l),t2)
                        xx(m,n)=xx(m,n)+g(eff,t1+i*ntrait,t2+j*ntrait)*&
                        w(k)*w(l)*mendel
                     endif
                  enddo
               enddo
            enddo
         enddo
      enddo
   enddo
enddo
end subroutine

subroutine read_parameter
integer :: i

print*,'Parameter file?'
read(*,'(a)') parfile
open(99,file=parfile)

read(99,'(a)') datafile
print*,'Data file: ',datafile

read(99,'(a)') pedfile
print*,'Pedigree file: ',pedfile

open(2,file=pedfile)

read(99,*) neff
print*,'# effects =',neff

allocate(nlev(neff))
read(99,*) nlev

do i=1,neff
   print*,i,'-th effect # levels =',nlev(i)
enddo

read(99,*) addeff
print*,'Additive effect =',addeff

read(99,'(a)') random
print*,'Random effect = ',random

read(99,*) var_a
print*,'Additive variance =',var_a

read(99,*) var_e
print*,'Residual variance =',var_e

end subroutine

end program blup1

subroutine solve_dense_gs(n,lhs,rhs,sol)
! finds sol in the system of linear equations: lhs*sol=rhs
! the solution is iterative by Gauss-Seidel
integer :: n
real :: lhs(n,n),rhs(n),sol(n),eps
integer :: round
!
round=0
do
   eps=0; round=round+1
   do i=1,n
      if (lhs(i,i).eq.0) cycle
      solnew=sol(i)+(rhs(i)-sum(lhs(i,:)*sol))/lhs(i,i)
      eps=eps+ (sol(i)-solnew)**2
      sol(i)=solnew
   end do
   print*,round,' rounds conv =',eps
   if (eps.lt. 1e-10) exit
end do
print*,'solutions computed in ',round,' rounds of iteration'
end subroutine


subroutine printmat(text,matrix,n,m)
! prints text and matrix
character text*(*)
integer :: n,m,i
real :: matrix(n,n)
print*,text
do i=1,m
   print '(11F7.2)',matrix(i,1:m)
enddo
end