!  bimrsim.f90
!  
!  Copyright 2025 Igor J. Chybicki <igorchy@ukw.edu.pl>
!  
!  This program is free software; you can redistribute it and/or modify
!  it under the terms of the GNU General Public License as published by
!  the Free Software Foundation; either version 2 of the License, or
!  (at your option) any later version.
!  
!  This program is distributed in the hope that it will be useful,
!  but WITHOUT ANY WARRANTY; without even the implied warranty of
!  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
!  GNU General Public License for more details.
!  
!  You should have received a copy of the GNU General Public License
!  along with this program; if not, write to the Free Software
!  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
!  MA 02110-1301, USA.
!  
!  INFO:
!  The program generates random genotype data under a simulation
!  model that resembles the inference model implemented in bimr2.
!  
!  As the program has no interface, parameter settings can be only changed
!  before compiling. See the subsection 'simulation settings' in the main
!  program below.
!  
!  The program was tested with the GNU Fortran (GFortran) compiler only.
!  The minimum command to compile the code is:
!  gfortran -o bimrsim bimrsim.f90
!  As a result, an executable binary file (named bimrsim) is created.

program bimrsim
   
   implicit none
   
   ! ----------------------
   ! global declarations
   !
   ! useful constants
   character(1), parameter :: tab=char(9)
   double precision, parameter :: mypi = dacos(dble(-1.0)), my2pi = 2.D0 * mypi
   
   ! useful variables
   integer, dimension(8) :: values
   character(256) :: title
   integer :: nthreads
   
   ! RNG-related variables
   integer*8 :: seed,jsr
   
   ! data dimensions
   integer :: npop, nind, nindpop, nloc, nall, nenv
   
   ! simulation replicability
   integer :: repl, nrepl
   
   ! simulation hyper-parameters
   double precision :: m_f, y_f, min_f, max_f, m_r, y_r, pr0, m00
   
   ! simulation parameters
   double precision, allocatable :: qlm(:,:), pklm(:,:,:), fk(:), beta(:,:), ri(:), psi(:,:), b_env(:)
   logical :: constant_divergence=.false.
   logical :: binary_transposed=.false.
   
   ! simulated data
   integer, allocatable :: ipop(:), igen(:,:,:), ianc(:,:)
   double precision, allocatable :: penv(:,:,:)
   
   ! ----------------------
   ! program starts here
   ! ----------------------
   
   ! rng seeded with the time (in milliseconds) since midnight
   call date_and_time(values=values)
   seed=((values(5)*60+values(6))*60+values(7))*1000+values(8)
   jsr=seed
   call init_rnd
   
   ! ----------------------
   ! simulation settings
   
   nrepl = 10            ! number of replications
   
   ! data dimensions
   npop = 15             ! number of populations
   nindpop = 20          ! per-pop number of individuals
   nind=npop*nindpop     ! total number of individuals
   nloc = 1000           ! number of markers
   nall = 2              ! number of alleles per locus
   nenv = 2              ! number of environmental factors
   
   nthreads=npop
   
   if (nall==2) binary_transposed=.true.
   
   ! hyper-parameters
   m00  = 0.7D0          ! expected non-migration rate
   
   allocate(b_env(nenv)) ! effect slopes for environmental factors
   b_env(1)=-0.9D0
   b_env(2)=1.1D0
   
   m_f = 0.10D0          ! mean divergence rate
   y_f = 0.D0            ! dispersion of divergence rates among pops (Beta parameter alpha=m_k*(1-y_f)/(y_f))
   min_f=0.01D0          ! minimum divergence value allowed
   max_f=0.9D0           ! maximum divergence value allowed
   
   m_r = 0.1D0           ! mean inbreeding
   y_r = 0.1D0           ! dispersion of individual inbreeding
   pr0 = 0.D0            ! zero-inflated beta distribution parameter (simple beta distrib. if pr0=0) (NOT USED)
   
   if (y_f==0.D0) constant_divergence=.true.
   ! ----------------------
   
   do repl=1,nrepl
     write (title,"('repl',i0)") repl
     call simulate_data
     write (*,*) 'Simulated data saved to "', trim(title),'"'
   enddo
   
   ! ----------------------
   ! program ends here
   ! ----------------------
   
 contains
 
 !----------------------------------------------------------------------------------------------------------
 ! sorts data per pop (required for proper BayesAss data reading)
 subroutine sort_ind_per_pop
   integer :: i,j,tmp_pop,tmp_gen(nloc,2),tmp_anc(2)
   double precision :: tmp_r
   
   do i=1,nind-1
     do j=i+1,nind
       if (ipop(i)>ipop(j)) then
         tmp_pop=ipop(i)
         tmp_gen=igen(i,:,:)
         ipop(i)=ipop(j)
         igen(i,:,:)=igen(j,:,:)
         ipop(j)=tmp_pop
         igen(j,:,:)=tmp_gen
         tmp_r=ri(i)
         ri(i)=ri(j)
         ri(j)=tmp_r
         tmp_anc=ianc(i,:)
         ianc(i,:)=ianc(j,:)
         ianc(j,:)=tmp_anc
       endif
     enddo
   enddo
 end subroutine sort_ind_per_pop
 !----------------------------------------------------------------------------------------------------------
 
 !----------------------------------------------------------------------------------------------------------
 ! simulates genetic data based on the zygotic/gametic migration model
 subroutine simulate_data
   integer :: k,l,m,n
   double precision :: sa
   double precision, allocatable :: ak(:)
   character(len=256) :: file_name
   
   ! memory allocations
   allocate(qlm(nloc,nall))
   allocate(pklm(npop,nloc,nall))
   allocate(fk(npop))
   allocate(beta(npop,npop))
   allocate(ipop(nind))
   allocate(igen(nind,nloc,2))
   allocate(ianc(nind,2))
   allocate(ak(max(npop,nall,2)))
   allocate(ri(nind))
   allocate(penv(npop,npop,nenv))
   allocate(psi(npop,npop))
   
   ! environmental factors
   ! distance-like
   penv=0.D0
   do k=1,npop
     do l=k+1,npop
       penv(k,l,1)=normaldev(0.D0,1.D0)
       penv(l,k,1)=penv(k,l,1)
     enddo
   enddo
   
   ! pop-specific
   do k=1,npop
     do l=k+1,npop
       penv(k,l,2)=normaldev(0.D0,1.D0)
       penv(l,k,2)=-penv(k,l,2)
     enddo
   enddo
   
   call normalize_factors
   
   ! psi hyper-parameters
   psi=0.D0
   do k=1,npop
     do l=1,npop
       if (k==l) cycle
       do m=1,nenv
         psi(k,l)=psi(k,l)+b_env(m)*penv(k,l,m)
       enddo
       ! assuming no deviation from the regression line
       psi(k,l)=dexp(psi(k,l))
     enddo
   enddo
   
   ! uniform prior global allele frequencies (for the F-model)
   qlm=1.D0/dble(nall)
   
   ! divergence rates (for the F-model)
   do k=1,npop
     if (constant_divergence) then
       fk(k)=m_f
     else
       do
         sa=0.D0 
         fk(k)=gammadev(m_f * (1.D0-y_f)/y_f)
         sa=fk(k)+gammadev((1.D0-m_f) * (1.D0-y_f)/y_f)
         fk(k)=fk(k)/sa
         if (fk(k)>min_f .and. fk(k)<max_f) exit
       enddo
     endif
     
     ! random gametic migration rates taken as draws from Dirichlet distrib.
     sa=0.D0
     ak=0.D0
     do l=1,npop
       if (l==k) cycle
       ak(l)=gammadev(psi(k,l))
       sa=sa+ak(l)
     enddo
     beta(k,:)=(1.D0-m00)*ak(1:npop)/sa
     beta(k,k)=m00
     
     ! random allele frequencies in pops taken as draws from Dirichlet distrib.
     do l=1,nloc
       sa=0.D0
       do m=1,nall
         ak(m)=gammadev(qlm(l,m) * (1.D0-fk(k))/fk(k))
         sa=sa+ak(m)
       enddo
       pklm(k,l,:)=ak(1:nall)/sa
     enddo
   enddo
   
   ! draws a random sample of individual genotypes
   do k=1,nind
     ! sampling location
     ipop(k)=(k-1)/nindpop+1
     
     m=draw_pop(beta(ipop(k),:))
     n=draw_pop(beta(ipop(k),:))
     ianc(k,1)=m
     ianc(k,2)=n
     
     ! draws inbreeding if both gametes from the same pop
     ! otherwise, inbreeding is zero
     if (m==n .and. m_r>0.D0) then
       if (rnd()<(1.D0-pr0)) then ! if zero-inflated Beta distribution is needed, pr0 should be greater than zero
         if (y_r>0.D0) then
           ak(1)=gammadev(m_r * (1.D0-y_r)/y_r)
           ak(2)=gammadev((1.D0-m_r) * (1.D0-y_r)/y_r)
           ri(k)=ak(1)/(ak(1)+ak(2))
         else
           ri(k)=m_r
         endif
       else
         ri(k)=0.D0
       endif
     else
       ri(k)=0.D0  
     endif
     
     ! generates multilocus genotype based on seed and pollen origin
     do l=1,nloc
       igen(k,l,1)=draw_allele(m,l)
       if (rnd()<ri(k)) then
         igen(k,l,2)=igen(k,l,1)
       else
         igen(k,l,2)=draw_allele(n,l)
       endif
     enddo
   enddo
   
   ! for example, BayesAss (or BA3) expects sorted data according to sampling individuals, locations, and loci
   call sort_ind_per_pop
   
   ! creates BayesAss (BA3) data file
   file_name=trim(title)//'.ba3'
   open (unit=789,file=file_name,action="write",status="replace")
   do l=1,nloc
     do m=1,nind
       write (789,"('ind',i0,' pop',i0,' loc',i0,2(' ',i0))") m,ipop(m),l,igen(m,l,1),igen(m,l,2)
     enddo
   enddo
   close (789)
   
   ! creates genenetic data file
   file_name=trim(title)//'.data'
   open (unit=789,file=file_name,action="write",status="replace")
   if (binary_transposed) then
     write (789,"(i0,*(a,i0))") 1,(tab,m,m=2,nind)
     write (789,"(i0,*(a,i0))") ipop(1),(tab,ipop(m),m=2,nind)
     do l=1,nloc
       write (789,"(i0,*(a,i0))") igen(1,l,1)+igen(1,l,2)-2,(tab,igen(m,l,1)+igen(m,l,2)-2,m=2,nind)
     enddo
   else
     do m=1,nind
       write (789,"(i0,*(a,i0))") m,tab,ipop(m),(tab,igen(m,l,1),tab,igen(m,l,2),l=1,nloc)
     enddo
   endif
   close (789)
   
   !creates environmental data file
   file_name=trim(title)//'.factors'
   open (unit=789,file=file_name,action="write",status="replace")
   do m=1,nenv
     write (789,"('[G]=',i0)") m
     do k=1,npop
       write (789,"(e16.8,*(a,e16.8))") penv(k,1,m),(tab,penv(k,l,m),l=2,npop)
     enddo
     write (789,"()")
   enddo
   close (789)
   
   ! creates file with expected migration rates
   file_name=trim(title)//'.params'
   open (unit=789,file=file_name,action="write",status="replace")
   write (789,"(a)") 'Expected migration rates:'
   write (789,"('Sink',a,'Source',a,'mij',a,'psi')",advance="no") tab,tab,tab
   if (nenv>0) then
     do k=1,nenv-1
       write (789,"(a,'env.',i0)",advance="no") tab,k
     enddo
     write (789,"(a,'env.',i0)",advance="yes") tab,nenv
   else
     write (789,"()",advance="yes")
   endif
   
   do k=1,npop
     do l=1,npop
       if (k==l) cycle
       write (789,"(2('Pop_',i0,a),f7.4,a,e16.8)",advance="no") k,tab,l,tab,beta(k,l),tab,psi(k,l)
       if (nenv>0) then
         do m=1,nenv-1
           write (789,"(a,e16.8)",advance="no") tab,penv(k,l,m)
         enddo  
         write (789,"(a,e16.8)",advance="yes") tab,penv(k,l,nenv)
       else
         write (789,"()",advance="yes")
       endif
     enddo
   enddo
   
   write (789,"(a)") ''
   write (789,"(a)") 'Expected divergence rates:'
   write (789,"('Pop',a,'fst')") tab
   do k=1,npop
     write (789,"('Pop_',i0,a,f7.4)") k,tab,fk(k)
   enddo
   write (789,"(a)") ''
   write (789,"(a)") 'Expected inbreeding:'
   write (789,"('Ind',a,'fis')") tab
   do k=1,nind
     write (789,"('Ind_',i0,a,f7.4)") k,tab,ri(k)
   enddo
   
   write (789,"(a)") ''
   write (789,"(a)") 'Hyper-parameters:'
   write (789,"(2a,f7.4)") 'm00',tab,m00
   write (789,"(2a,f7.4)") 'm_fst',tab,m_f
   write (789,"(2a,f7.4)") 'm_fis',tab,m_r
   write (789,"(2a,f7.4)") 'y_fis',tab,y_r
   
   close(789)
   
   !creates settings file for bimr2
   file_name=trim(title)//'.bimr2'
   open (unit=789,file=file_name,action="write",status="replace")
   write (789,"(a,i0)") '#seed = ',0
   write (789,"(a,i0)") '#individuals = ',nind
   write (789,"(a,i0)") '#markers = ',nloc
   write (789,"(a,i0)") '#samples = ',20000
   write (789,"(a,i0)") '#burnin = ',10000
   write (789,"(a,i0)") '#keepevery = ',10
   write (789,"(a,i0)") '#threads = ',nthreads
   write (789,"(2a)") '#gfile = ',trim(title)//'.data'
   if (binary_transposed) write (789,"(a)") '+biallelic_transposed'
   write (789,"(a)") '+rjmcmc'
   write (789,"(a,i0)") '+envir.factors = ',nenv
   write (789,"(a)") '+interactions'
   write (789,"(a)") '+inbreeding'
   write (789,"(2a)") '+efile = ',trim(title)//'.factors'
   write (789,"(a)") '#end'
   close(789)
   
   call BIMr1_data_file
   
   ! frees memory
   deallocate(qlm)
   deallocate(pklm)
   deallocate(fk)
   deallocate(beta)
   deallocate(ipop)
   deallocate(igen)
   deallocate(ianc)
   deallocate(penv)
   deallocate(psi)
   deallocate(ak)
   deallocate(ri)
   
   ! run bimr2
   file_name='./bimr2 '//trim(title)//'.bimr2'
   call execute_command_line (trim(file_name),wait=.true.)
   
 end subroutine simulate_data
 !----------------------------------------------------------------------------------------------------------
 
 !----------------------------------------------------------------------------------------------------------
 ! draws a random allele based on allele freqs in the k-th pop and l-th loc
 function draw_allele(pop,loc)
   integer :: draw_allele,m,pop,loc
   double precision :: r,sp
   
   draw_allele=1
   r=rnd()
   sp=0.D0
   do m=1,nall
     sp=sp+pklm(pop,loc,m)
     if (r<sp) then
       draw_allele=m
       return
     endif
   enddo
   
 end function draw_allele
 !----------------------------------------------------------------------------------------------------------
 
 !----------------------------------------------------------------------------------------------------------
 ! draws pop number based on migration vector m
 function draw_pop(m)
   integer :: draw_pop,k
   double precision :: r,m(npop),sm
   
   draw_pop=1
   r=rnd()
   sm=0.D0
   do k=1,npop
     sm=sm+m(k)
     if (r<sm) then
       draw_pop=k
       return
     endif
   enddo
   
 end function draw_pop
 !----------------------------------------------------------------------------------------------------------
 
 !----------------------------------------------------------------------------------------------------------
 ! implements a shift-register generator (64-bit version)
 function shift_xor(val,shift)
   integer*8 :: shift_xor
   integer*8, intent(in) :: val, shift
   
   shift_xor=ieor(val,ishft(val,shift))
 end function shift_xor
 !----------------------------------------------------------------------------------------------------------
 
 !----------------------------------------------------------------------------------------------------------
 ! initializes random number generator
 subroutine init_rnd
    
    jsr=shift_xor(shift_xor(shift_xor(jsr, 13_8),-17_8),43_8)
 end subroutine init_rnd
 !----------------------------------------------------------------------------------------------------------
 
 !----------------------------------------------------------------------------------------------------------
 ! returns a random draw from the uniform (0,1) pdf
 double precision function rnd()
   integer*8 :: ojsr
   ojsr=jsr
   jsr=shift_xor(shift_xor(shift_xor(jsr, 13_8),-17_8),43_8)
   rnd=dble(ojsr+jsr)*5.42101086242752217D-20+0.5D0
 end function rnd
 !----------------------------------------------------------------------------------------------------------
 
 !----------------------------------------------------------------------------------------------------------
 ! returns a random draw from the standard gamma pdf (should be used for a>1)
 double precision function rgamma(a)
   double precision, intent(in) :: a
   double precision :: x,v,u,c,d
   
   d=a-dble(1)/dble(3)
   c=dble(1)/dsqrt(dble(9)*d)
   
   do
     do
       x=dsqrt(-dble(2)*dlog(rnd()))*dcos(my2pi*rnd())
       v=(dble(1)+c*x)
       if (v > dble(0)) exit
     enddo
     
     v=v**3
     u=rnd()
     
     if (u<(dble(1)-dble(0.0331)*x**4)) exit
     if (dlog(u)<(dble(0.5)*x**2+d*(dble(1)-v+dlog(v)))) exit
   enddo
   
   rgamma=d*v
 end function rgamma
 !----------------------------------------------------------------------------------------------------------
 
 !----------------------------------------------------------------------------------------------------------
 ! returns a random draw from the standard gamma pdf
 double precision function gammadev(a)
   double precision :: a
  
   if (a<dble(1)) then
     if (a<0.01D0) a=0.01D0
     gammadev=rgamma(a+dble(1))*rnd()**(dble(1)/a)
   else
     gammadev=rgamma(a)
   endif
  
 end function gammadev
 !----------------------------------------------------------------------------------------------------------
 
 !----------------------------------------------------------------------------------------------------------
 ! This routine returns a random draw from the normal pdf (mean m and std.dev. s)
 double precision function normaldev(m,s)
   double precision, intent(in) :: m,s
   double precision :: v1,v2,r
   
   do
     v1 = dble(2) * rnd() - dble(1)
     v2 = dble(2) * rnd() - dble(1)
     r  = v1**2 + v2**2
     if (r < dble(1) .and. r > dble(0)) exit
   enddo
   normaldev = v2 * dsqrt(-dble(2) * dlog(r) / r) * s + m
   
 end function normaldev
 !----------------------------------------------------------------------------------------------------------
 
 !----------------------------------------------------------------------------------------------------------
 ! This routine normalizes environmental factors
 subroutine normalize_factors
   integer :: i,j
   double precision :: mean(nenv),var(nenv)
   
   do i=1,nenv
     mean(i)=(sum(penv(:,:,i))-sum((/(penv(j,j,i),j=1,npop)/)))/dble(npop*(npop-1)) ! removing diagonals
     var(i)=(sum(penv(:,:,i)**2)-sum((/(penv(j,j,i)**2,j=1,npop)/)))/dble(npop*(npop-1))-mean(i)**2 ! removing diagonals
     penv(:,:,i)=(penv(:,:,i)-mean(i))/dsqrt(var(i))
     do j=1,npop
       penv(j,j,i)=0.D0
     enddo
   enddo
   
 end subroutine
 !----------------------------------------------------------------------------------------------------------
 
 !----------------------------------------------------------------------------------------------------------
 ! This routine saves data for BIMr1
 subroutine BIMr1_data_file
   integer :: i,l,tmp_gen(nloc,2)
   character(LEN=255) :: BIMr_file,extension
   
   extension='gen'
   BIMr_file=trim(title)//'.gen'
   
   open (unit=789,file=BIMr_file,action="write",status="replace")
   
   write (789,"('[individuals]= ',i0)") nind
   write (789,"('[populations]= ',i0)") npop
   write (789,"('[loci]= ',i0)") nloc
   write (789,"('[alleles]= ',i0)") nall
   
   write (789,"('[genotypes]= ',i0)")
   do i=1,nind
     tmp_gen=igen(i,:,:)
     do l=1,nloc
       if (tmp_gen(l,1)<0) tmp_gen(l,1)=0
       if (tmp_gen(l,2)<0) tmp_gen(l,2)=0
     enddo
     
     write (789,"('ind',i0,' ,',a,i0,*(a,2i2.2))") i,tab,ipop(i),(tab,tmp_gen(l,1),tmp_gen(l,2),l=1,nloc)
   enddo
   close (789)
   
 end subroutine
 !----------------------------------------------------------------------------------------------------------
 
end program