module inb
    !! Variables should be prepared by user
    ! unsortped: identification number of all animals in the ped (animal should be ordered in 1 to number of animal), with birth year in the last column
    ! gamma: relationship coefficient matrix bewteen metafounders (see Lagerra et al. 2015)
    implicit none
    integer                 :: n, maxlap, nmf, nanim ! n = number of animal not included metafounder
    integer, allocatable    :: unsortped(:,:), revlink(:), link(:), ordped(:,:), lap(:), lap_s(:), lap_e(:), BY(:), linkmf(:), revlinkmf(:)
    integer, allocatable    :: ped(:,:), ss(:), dd(:), sire_num(:), grouped_ID(:), init_i(:)
    logical(1), allocatable :: indirect_TF(:)
    real(8), allocatable    :: gamma(:,:), F(:), Diag(:)

    contains

    subroutine sort_ped()
        !! Prepare variables:
        ! ordped: ped sorted based on lap
        ! lap: longest ancestral path for each animal in ordped; maxlap = maxval(lap, dim = 1)
        ! lap_s, lap_e: start and end position of each lap block
        ! BY: sorted birth year
        implicit none
        integer :: i, j
        integer, allocatable :: unsortlap(:), nlap(:), nlapi(:) 
        logical :: finish
        
        !! Calculate lap for animal in unsortped
        n = size(unsortped, dim = 1)
        allocate(unsortlap(0:n), source = 0)
        finish = .FALSE. 
        j = 0
        do while (.not.finish)
            j = j+1
            finish = .TRUE.
            do i = 1, n
                if (unsortlap(i)<max(unsortlap(unsortped(i,2)),unsortlap(unsortped(i,3)))+1) then 
                    if(finish) finish = .FALSE.
                    unsortlap(i) = max(unsortlap(unsortped(i,2)),unsortlap(unsortped(i,3))) + 1
                end if
            end do 
        end do

        !! Sort pedigree based on lap number
        maxlap = maxval(unsortlap)
        allocate(nlap(maxlap), source = 0)
        allocate(nlapi(maxlap), source = 0)
        allocate(link(n), source = 0)
        allocate(revlink(n), source = 0)
        allocate(lap_e(0:maxlap), source = 0)
        allocate(lap_s(maxlap), source = 0)
        allocate(lap(0:n), source = 0)


        ! Count number of animal in each lap
        do i = 1, n
            nlap(unsortlap(i)) = nlap(unsortlap(i)) + 1
        end do 
        do i = 1, maxlap
            lap_e(i) = nlap(i) + lap_e(i-1)
        end do
        lap_s(1:maxlap) = lap_e(0:maxlap-1) + 1

        ! Create link list
        do i = 1, n
            nlapi(unsortlap(i)) = nlapi(unsortlap(i)) + 1       ! Current number of animal in this lap
            j = lap_e(unsortlap(i)-1) + nlapi(unsortlap(i))     ! j = sorted identification number
            link(i) = j                                         ! record new identification number for old (unsorted) ped
            revlink(j) = i                                      ! record old identification number for new (sorted) ped
        end do

        ! Sort ped based on lap
        allocate(ordped(n,2), source = 0)
        allocate(BY(n), source = 0)
        do i = 1, n
            if(unsortped(i,2)/=0) ordped(link(i),1) = link(unsortped(i,2))
            if(unsortped(i,3)/=0) ordped(link(i),2) = link(unsortped(i,3))
            BY(link(i)) = unsortped(i,4)
            lap(link(i)) = unsortlap(i)
        end do
    end subroutine

    subroutine add_mf()
        !! Prepare variables:
        ! nmf : number of metafounder
        ! ped: pedigree with metafounders assigned based on birth year of progeny
        implicit none 
        integer :: i, j, base_year

        base_year = minval(BY)
        nmf = maxval(BY) - base_year + 1
        nanim = n + nmf; lap_s(:) = lap_s(:) + nmf; lap_e(:) = lap_e(:) + nmf
        allocate(ped(0:nanim, 3), source = 0) ! i, s, d
        allocate(linkmf(n), source = 0)
        allocate(revlinkmf(nanim), source = 0)
        do i = 1, nanim
            ped(i,1) = i
            if(i>nmf) then
                j = i-nmf 
                revlinkmf(i) = revlink(j)
                linkmf(j) = link(j) + nmf
                if (ordped(j,1)/=0) then 
                    ped(i,2) = ordped(j,1) + nmf
                else 
                    ! unknown parent assigned based on progeny birth year
                    ped(i,2) = BY(j) - base_year + 1
                end if
                if (ordped(j,2)/=0) then 
                    ped(i,3) = ordped(j,2) + nmf
                else 
                    ! unknown parent assigned based on progeny birth year
                    ped(i,3) = BY(j) - base_year + 1
                end if
            end if
        end do
    end subroutine

    subroutine pre_F_mi() 
        !! Prepare variables:
        ! ss, dd = parents with larger and smaller identification number, respectively
        ! sire_num = renumber sire 
        ! grouped_ID = ID grouped by sire ID
        ! init_i = Record the start position of each sire group in the sorted ID 
        ! indirect_TF = Whether the animal should be processed by indirect method (the first progeny for a sire)
        implicit none
        integer :: i, j, is, ns, nindirect
        integer, allocatable :: npro(:), tmpi(:)

        !! Determine large ID in each mating pair
        allocate(ss(nanim), source = 0) 
        allocate(dd(nanim), source = 0)    
        do i = nmf+1, nanim
            if(ped(i,2)==0.or.ped(i,3)==0) stop 'missing genealogy for non-metafounder'
            ss(i) = max(ped(i,2), ped(i,3))
            dd(i) = min(ped(i,2), ped(i,3))
        end do 

        !! Count progeny number for each sire
        allocate(npro(nanim), source = 0) 
        allocate(sire_num(nanim), source = 0) 
        allocate(indirect_TF(nmf+1:nanim))
        indirect_TF(:) = .FALSE.
        do i = nmf+1, nanim
            npro(ped(i,2)) = npro(ped(i,2)) + 1
        end do 
        ns = count(npro>0)

        !! Create renumbered list containing Sires for evaluation
        allocate(init_i(ns+1), source = 0)
        init_i(1) = 1
        j = 0
        do i = 1, nanim
            if(npro(i)>0) then 
                j = j + 1 
                ! Record the start position of each sire group in the grouped_ID
                init_i(j+1) = init_i(j) + npro(i) 
                sire_num(i) = j
            end if 
        end do
        nindirect = init_i(ns+1) - 1

        !! Sort animal ID according to their sire ID
        allocate(tmpi(ns), source = 0)
        allocate(grouped_ID(nindirect+1),  source = 0)
        do i = nmf+1, nanim
            is = sire_num(ped(i,2))
            if(is>0) then 
                if(tmpi(is)==0) then
                    ! Apply indirect method if the animal is the first progeny of a sire
                    indirect_TF(i) = .TRUE. 
                    tmpi(is) = init_i(is) 
                end if
                grouped_ID(tmpi(is)) = i
                tmpi(is) = tmpi(is) + 1
            end if 
        end do
    end subroutine
    
    subroutine get_F_mi_parallel() 
        ! Modified Sargolzaei et al. (2005) indirect method for related metafounders, by Huang and Hirama (2025)
        implicit none
        integer  :: is,i,j,tmpj,g,js,jd,lapi,laps,lapd,geni
        real(8)  :: r,fi
        logical(1), allocatable  :: flag(:)
        integer, allocatable     :: pt(:), fiv(:), liv(:)
        real(8), allocatable     :: l(:)
        
        if(.not.allocated(F)) allocate(F(nanim))
        if(.not.allocated(Diag)) allocate(Diag(nanim))
        F = -5d0; Diag = 0d0

        do i=1, nmf
            F(i)=gamma(i,i)-1d0
        end do

        !$OMP PARALLEL private(is, i, j, g, r, fi, tmpj, js, jd, lapi, laps, lapd, fiv, liv, pt, l, flag) 
        allocate(pt(nanim), source = 0)
        allocate(fiv(0:maxlap), source = 0)
        allocate(liv(0:maxlap), source = 0)
        allocate(l(nanim), source = 0d0)
        allocate(flag(nanim))
        flag = .FALSE.; flag(1:nmf) = .TRUE.
        do geni = 1, maxlap
            !$OMP DO schedule(dynamic)
            do i=lap_s(geni),lap_e(geni)
                is=ped(i,2)
                Diag(i)=0.5-0.25*(F(is)+F(ped(i,3)))
                if(.not.indirect_TF(i)) cycle

                ! Compute isth or ith row of L 
                lapi = lap(is)
                fiv(lapi) = is
                liv(lapi) = is
                l(is)=1d0
                do g = lapi, 0, -1
                    j = fiv(g)

                    do while(j /= 0)
                        js = ss(j); jd = dd(j)
                        r = 0.5d0 * l(j)
    
                        if(js > 0)then
                            l(js) = l(js) + r
                            laps = lap(js)
                            if(fiv(laps) == 0)then
                                fiv(laps) = js
                                liv(laps) = js
                            else
                                if(liv(laps) /= js .and. pt(js) == 0)then
                                    pt(liv(laps)) = js
                                    liv(laps) = js
                                end if
                            end if
    
                            if(jd > 0)then
                                l(jd) = l(jd) + r
                                lapd = lap(jd)
                                if(fiv(lapd) == 0)then
                                    fiv(lapd) = jd
                                    liv(lapd) = jd
                                else
                                    if(liv(lapd) /= jd .and. pt(jd) == 0)then
                                        pt(liv(lapd)) = jd
                                        liv(lapd) = jd
                                    end if
                                end if
                            end if
                        end if
    
                        if(j>nmf) l(j)=l(j)*Diag(j)
                        tmpj=j
                        j=pt(j)
                        pt(tmpj)=0
                    end do
                    fiv(g)=0; liv(g)=0
                end do

                ! Multiply by D for MF
                l(1:nmf)=matmul(gamma,l(1:nmf))
                
                ! Obtain F for the progeny of the sire of i (includes i)
                j = init_i(sire_num(is))
                print "(a,i3,a,16f5.2)", "L of sire", revlinkmf(is), " :", l(linkmf)
                do while(ped(grouped_ID(j),2)==is)
                    ! x and flag need to be passed explicitly because they are private
                    call solve(ped(grouped_ID(j),3), l, flag, fi)
                    F(grouped_ID(j)) = fi * 0.5d0
                    print "(a,i3,a,f5.2,a,i3,a)", "F of", revlinkmf(grouped_ID(j)), " = ", F(grouped_ID(j)), " (Dam =", revlinkmf(ped(grouped_ID(j), 3)),")"
                    j = j + 1
                end do
                call solve(is, l, flag, fi) ! just to flag elements on s-th row of L
                print *

                ! Reset l to zero; flag to .FALSE. (except 1:nmf)
                j = init_i(sire_num(is))
                do while(ped(grouped_ID(j),2)==is)
                    call initialize(ped(grouped_ID(j),3), l, flag)
                    j = j + 1
                end do
                call initialize(is, l, flag)
                l(1:nmf)=0d0
            end do
            !$OMP END DO
        end do
        !$OMP END PARALLEL
        
        contains

        recursive subroutine solve(k, array, finish_flag, res) 
            ! Trace back the pedigree
            implicit none
            integer, intent(in)  :: k
            real(8) :: xs, xd
            real(8), intent(inout) :: res
            real(8), intent(inout) :: array(:)
            integer, allocatable :: idx(:)
            logical(1),  intent(inout) :: finish_flag(:)

            if(finish_flag(k)) then
                res = array(k)
            else
                finish_flag(k) = .TRUE.
                call solve(ped(k,2), array, finish_flag, xs)  
                call solve(ped(k,3), array, finish_flag, xd) 
                array(k) = array(k) + (xs+xd)*0.5d0
                print "(a,16f5.2)", "y: ", array(linkmf)
                res = array(k)
            end if
        end subroutine 

        recursive subroutine initialize(k, array, finish_flag)
            ! Initialize array to zero
            implicit none
            integer, intent(in) :: k
            real(8), intent(inout) :: array(:)
            logical(1), intent(inout) :: finish_flag(:)
            
            if(finish_flag(k) .and. k > nmf) then
                finish_flag(k) = .FALSE.
                call initialize(ped(k,2), array, finish_flag)  
                call initialize(ped(k,3), array, finish_flag)
                array(k) = 0d0 
            end if
        end subroutine
    end subroutine 

    end module inb

    program cal_f
        use inb
        use omp_lib
        integer :: nline,i,io

        open(file = "example.ped", unit = 11, action = "read")
        nline=0
        do
            read(11, *, iostat=io) 
            if(io/=0) exit
            nline=nline+1
        enddo
        rewind(11)

        allocate(unsortped(nline, 4))
        do i = 1, nline
            read(11, *) unsortped(i, :)
        enddo 

        call sort_ped()
        call add_mf()
        call pre_F_mi()
        call omp_set_num_threads(2)
        allocate(gamma(nmf,nmf), source = 0d0)
        call get_F_mi_parallel()
        
        print *, " ID  F"
        do i = nmf+1, nanim
            print "(i3, f5.2)", revlink(i-nmf), F(i)
        end do
    end program