module globals

  ! Parameters for q1c: n = 64, a = 5, q = 1.0, sigma = 0.2
  ! Parameter changes for q2a: sigma = 0.4
  ! Parameter changes for q2b: n = 256
  ! Parameter changes for q2c: n = 64, sigma = 0.6
  ! Parameter changes for q2d: sigma = 0.4, m = 2,3,5
  integer, parameter   :: n = 64
  integer, parameter   :: m = 4
  real, parameter      :: a = 5.43
  real, parameter      :: d = real(m*a) / real(n)
  real, parameter      :: q = 1.0
  real, parameter      :: sigma = 0.3

end module globals


module utils

  real, parameter :: pi = acos(-1.0)
  real, parameter :: sq2pi = sqrt(2.0 * pi)

end module utils


subroutine cfft2(c,n) 
  implicit none

  integer, intent(in) :: n 
  complex, dimension(n,n), intent(inout) :: c

  real, dimension(4*n+15) :: wrk
  integer :: i, j, i2, j2
  complex :: tmp

  call cffti(n,wrk)

  do i = 1, n
    call cfftf(n,c(i,:),wrk)
  enddo
  do i = 1, n
    call cfftf(n,c(:,i),wrk)
  enddo

  do i = 1, n
    do j = 1, (n/2)
      tmp = c(i,j)
      i2 = mod(i + n/2 - 1, n) + 1
      j2 = mod(j + n/2 - 1, n) + 1
      c(i,j) = c(i2,j2)
      c(i2,j2) = tmp
    enddo
  enddo

end subroutine cfft2


subroutine init_one_atom(n, density, d, q, sigma)
  use utils
  implicit none

  integer, intent(in) :: n
  complex, dimension(n,n), intent(inout) :: density
  real, intent(in) :: d, q, sigma

  integer :: i, j
  real :: x, y, center, sqdist

  center = real(n)*d/2.0

  do i = 1, n
    do j = 1, n
      x = real(i) * d
      y = real(j) * d
      sqdist = (x - center) ** 2.0 + (y - center) ** 2.0
      density(i,j) = (q / (sigma * sq2pi)) * exp(-sqdist / (2.0 * sigma**2.0))
    enddo
  enddo

end subroutine init_one_atom


subroutine mkreseau(m, a, r)
  use utils
  implicit none

  integer, intent(in) :: m
  real, intent(in)    :: a
  real, dimension(m**2,2), intent(out) :: r

  integer :: i,j

  do i = 1, m
    do j = 1, m
      r((j-1) * m + i, 1) = (real(i)-0.75)*a
      r((j-1) * m + i, 2) = (real(j)-0.75)*a
    enddo
  enddo

end subroutine mkreseau


subroutine mkreseau2(m, a, r)
  use utils
  implicit none

  integer, intent(in) :: m
  real, intent(in)    :: a
  real, dimension(2*(m**2),2), intent(out) :: r

  integer :: i, j, ind, modi

  ind = 1
  do i = 1, 2*m
    modi = mod(i-1,2)
    do j = 1, m
      r(ind, 1) = (real((i-1)/2)+0.25+real(modi)*0.50)*a
      r(ind, 2) = (real(j-1)+0.25+real(modi)*0.50)*a
      ind = ind+1
    enddo
  enddo

end subroutine mkreseau2


subroutine mkdens(n, nat, r, d, rho, sigma)
  use utils
  implicit none

  integer, intent(in) :: n, nat
  real, dimension(nat,2), intent(in) :: r
  real, intent(in) :: d, sigma
  complex, dimension(n,n), intent(out) :: rho

  integer :: i, j, im
  real :: x, y, sqdist

  do i = 1, n
    do j = 1, n
      x = real(i) * d
      y = real(j) * d
      rho(i,j) = 0
      do im = 1, nat 
        sqdist = (x - r(im,1))**2.0 + (y - r(im,2))**2.0
        rho(i,j) = rho(i,j) + exp(-sqdist / (2.0 * sigma**2.0))
      enddo
    enddo
  enddo

end subroutine mkdens


! Main program
program main
  use globals
  use utils
  implicit none

  complex, dimension(n,n) :: density
  real, dimension(n,n)    :: intens
  real, dimension(2*(m**2),2) :: r

  integer :: i, j
  real :: x, y, qx, qy

  call mkreseau2(m, a, r)
  call mkdens(n, 2*(m**2), r, d, density, sigma)

  open(1, file='results/q3a_density.out')
  do i = 1, n
    do j = 1, n
      x = real(i) * d
      y = real(j) * d
      write(1,*) x, y, cabs(density(i,j))
    enddo
    write(1,*) ''
  enddo
  close(1)


  call cfft2(density, n)
  intens = cabs(density)**2.0

  open(2, file='results/q3a_intensity.out')
  do i = 1, n
    do j = 1, n
      qx = 2.0 * pi * real(i - n/2) / (real(m) * a)
      qy = 2.0 * pi * real(j - n/2) / (real(m) * a)
      write(2,*) qx, qy, intens(i,j)
    enddo
    write(2,*) ''
  enddo
  close(2)

end program main
