module nbody_utils
  use omp_lib
  implicit none
  integer, parameter :: pf = kind(1.0d0)
  integer, parameter :: dim = 3
  real(pf), parameter :: G = 6.67430e-11_pf
  real(pf), parameter :: potsoft2 = 1.0e-4_pf
  
contains

    ! Versão sequencial
    subroutine forcas_seq(m, R, forcas, distancias, N)
        real(pf), intent(in) :: m(:), R(:,:)
        real(pf), intent(out) :: forcas(:,:), distancias(:)
        integer, intent(in) :: N
        integer :: i, j, idx
        real(pf) :: dx, dy, dz, r2, rinv, factor
        
        forcas = 0.0_pf
        
        do i = 1, N-1
        do j = i+1, N
            dx = R(j,1) - R(i,1)
            dy = R(j,2) - R(i,2)
            dz = R(j,3) - R(i,3)
            
            r2 = dx*dx + dy*dy + dz*dz + potsoft2
            rinv = G * m(i) * m(j) / (r2 * sqrt(r2))
            factor = rinv
            
            forcas(i,1) = forcas(i,1) + dx * factor
            forcas(i,2) = forcas(i,2) + dy * factor
            forcas(i,3) = forcas(i,3) + dz * factor
            
            forcas(j,1) = forcas(j,1) - dx * factor
            forcas(j,2) = forcas(j,2) - dy * factor
            forcas(j,3) = forcas(j,3) - dz * factor
            
            idx = (i-1)*(2*N-i)/2 + (j-i)
            distancias(idx) = sqrt(r2 - potsoft2)
        end do
        end do
    end subroutine forcas_seq

    ! Versão paralela com OpenMP para GPU
    subroutine forcas_par(m, R, forcas, distancias, N)
        real(pf), intent(in) :: m(:), R(:,:)
        real(pf), intent(out) :: forcas(:,:), distancias(:)
        real(pf) :: forcas_local(3)
        integer, intent(in) :: N
        integer :: i, j, idx, chunk_size
        real(pf) :: dx, dy, dz, r2, rinv, factor
        
        ! Tamanho do bloco para melhor localidade (ajuste conforme sua GPU)
        chunk_size = 64
        
        ! 1. Manter dados na GPU entre chamadas
        !$OMP target data map(to: R, m) map(from: forcas, distancias)
        
        ! 2. Inicialização otimizada
        !$OMP target teams distribute parallel do simd num_teams(N/256) thread_limit(256)
        do i = 1, N
            forcas(i,:) = 0.0_pf
        end do
        
        ! 3. Cálculo principal com janela de trabalho fixa
        !$OMP target teams distribute parallel do private(dx, dy, dz, r2, rinv, factor, idx) &
        !$OMP num_teams(480) thread_limit(256)
        do i = 1, N
            ! Loop interno sem collapse mas com janela fixa
            ! do j = i+1, min(i+chunk_size, N)
            do j = i+1, N
                dx = R(j,1) - R(i,1)
                dy = R(j,2) - R(i,2)
                dz = R(j,3) - R(i,3)
                
                r2 = dx*dx + dy*dy + dz*dz + potsoft2
                rinv = G * m(i) * m(j) / (r2 * sqrt(r2))
                factor = rinv

		!$OMP atomic update
                forcas(i,1) = forcas(i,1) + dx * factor
		!$OMP atomic update
                forcas(i,2) = forcas(i,2) + dy * factor
		!$OMP atomic update
                forcas(i,3) = forcas(i,3) + dz * factor
                
		!$OMP atomic update
                forcas(j,1) = forcas(j,1) - dx * factor
		!$OMP atomic update
                forcas(j,2) = forcas(j,2) - dy * factor
		!$OMP atomic update
                forcas(j,3) = forcas(j,3) - dz * factor
                
                idx = (i-1)*(2*N-i)/2 + (j-i)
                distancias(idx) = sqrt(r2 - potsoft2)
            end do

        end do
        
        !$OMP end target data
    end subroutine forcas_par

    ! subroutine forcas_par_precise(m, R, forcas, distancias, N)
    ! real(pf), intent(in) :: m(:), R(:,:)
    ! real(pf), intent(out) :: forcas(:,:), distancias(:)
    ! integer, intent(in) :: N
    ! integer :: i, j, idx, team_id, local_i, local_j
    ! real(pf) :: dx, dy, dz, r2, rinv, factor
    ! real(pf), allocatable :: forcas_local(:,:,:)
    ! integer :: max_threads

    ! ! Alocação de acumuladores locais por thread
    ! !$OMP declare target
    ! allocate(forcas_local(N, dim, omp_get_max_threads()))
    
    ! max_threads = omp_get_max_threads()
    
    ! ! Inicialização
    ! !$OMP target teams distribute parallel do simd collapse(2) map(tofrom: forcas_local)
    ! do i = 1, N
    !     do j = 1, max_threads
    !     forcas_local(i,:,j) = 0.0_pf
    !     end do
    ! end do
    
    ! ! Cálculo com acumulação local
    ! !$OMP target teams distribute parallel do private(dx, dy, dz, r2, rinv, factor, idx, team_id) &
    ! !$OMP map(to: R, m) map(tofrom: forcas_local, distancias)
    ! do i = 1, N-1
    !     team_id = omp_get_team_num() * omp_get_num_threads() + omp_get_thread_num()
    !     do j = i+1, N
    !     dx = R(j,1) - R(i,1)
    !     dy = R(j,2) - R(i,2)
    !     dz = R(j,3) - R(i,3)
        
    !     r2 = dx*dx + dy*dy + dz*dz + potsoft2
    !     rinv = G * m(i) * m(j) / (r2 * sqrt(r2))
    !     factor = rinv
        
    !     ! Acumulação local por thread
    !     forcas_local(i,1,team_id+1) = forcas_local(i,1,team_id+1) + dx * factor
    !     forcas_local(i,2,team_id+1) = forcas_local(i,2,team_id+1) + dy * factor
    !     forcas_local(i,3,team_id+1) = forcas_local(i,3,team_id+1) + dz * factor
        
    !     forcas_local(j,1,team_id+1) = forcas_local(j,1,team_id+1) - dx * factor
    !     forcas_local(j,2,team_id+1) = forcas_local(j,2,team_id+1) - dy * factor
    !     forcas_local(j,3,team_id+1) = forcas_local(j,3,team_id+1) - dz * factor
        
    !     idx = (i-1)*(2*N-i)/2 + (j-i)
    !     distancias(idx) = sqrt(r2 - potsoft2)
    !     end do
    ! end do
    
    ! ! Redução final
    ! !$OMP target teams distribute parallel do simd map(to: forcas_local) map(from: forcas)
    ! do i = 1, N
    !     forcas(i,:) = sum(forcas_local(i,:,:), dim=2)
    ! end do
    
    ! deallocate(forcas_local)
    ! end subroutine forcas_par_precise

end module nbody_utils

program main
  use nbody_utils
  implicit none
  integer, parameter :: N = 512 * 10  ! Tamanho do problema
  real(pf), allocatable :: m(:), R(:,:), f_seq(:,:), f_par(:,:), distancias(:), distancias1(:)
  real(pf) :: start_time, end_time
  integer :: i, j
  
  ! Aloca memória
  allocate(m(N), R(N,dim), f_seq(N,dim), f_par(N,dim))
  allocate(distancias(N*(N-1)/2),distancias1(N*(N-1)/2))
  
  ! Inicializa posições e massas aleatórias
  call random_number(R)
  call random_number(m)
  R = R * 10.0_pf  ! Escala as posições
  m = m * 1000.0_pf ! Escala as massas
  
  ! Executa versão sequencial
  call cpu_time(start_time)
  call forcas_seq(m, R, f_seq, distancias, N)
  call cpu_time(end_time)
  print '(A,F10.4,A)', 'Tempo sequencial: ', end_time - start_time, ' segundos'
  
  ! Executa versão paralela
  start_time = omp_get_wtime()
!  call cpu_time(start_time)
  call forcas_par(m, R, f_par, distancias1, N)
!  call cpu_time(end_time)
  end_time = omp_get_wtime()
  print '(A,F10.4,A)', 'Tempo paralelo:   ', end_time - start_time, ' segundos'

  ! Verifica resultados
  print '(A,E12.4)', 'Diferença máxima: ', maxval(abs(f_seq - f_par))
  print '(A,E12.4)', 'Diferença máxima dist: ', maxval(abs(distancias1 - distancias))
  
  ! Libera memória
  deallocate(m, R, f_seq, f_par, distancias)
end program main
