!-----------------------------------------------------------------------
! PSCF - Polymer Self-Consistent Field Theory
! Copyright (2002-2016) Regents of the University of Minnesota
! contact: David Morse, morse012@umn.edu
!
! 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. A copy of this license is included in
! the LICENSE file in the top-level PSCF directory. 
!----------------------------------------------------------------------
!****m scf/scf_mod
! PURPOSE
!   Calculate monomer concentrations, free energies, and stresses.
!   Solve the modified diffusion equation for chains by a pseudo-spectral 
!   algorithm. Use a simple Boltzmann weight on a grid for solvent.
! AUTHOR 
!   Jian Qin - Implemented pseudo-spectral algorithm (2005-2006)
!   Raghuram Thiagarajan - Added small molecule solvent (2007)
!   Kevin J Hou - Added electrostatic interactions (2016)
! SOURCE
!------------------------------------------------------------------------
module scf_mod 
   use const_mod
   use chemistry_mod
   use fft_mod
   use grid_mod
   use grid_basis_mod 
   use chain_mod
   use step_mod
   implicit none

   ! Internal Variables
   private aminusbvecs              ! qa-qb (in fft)
   private qadotqb                  ! qa dot qb
   private qaminbdotqa              ! (qb-qa) dot qb

   ! public procedures
   public:: density_startup         ! allocates arrays needed by density
   public:: density                 ! scf calculation of rho & q
   public:: scf_stress              ! calculates d(free energy)/d(cell_param)
   public:: mu_phi_chain            ! calculates mu from phi (canonical)
                                    ! or phi from mu (grand) for chains
   public:: mu_phi_solvent          ! calculates mu from phi (canonical)
                                    ! or phi from mu (grand) for solvents
   public:: free_energy             ! calculates helmholtz free energy 
                                    ! (optionally calculates pressure)
   public:: free_energy_components  ! Returns the components of the helmholtz FE
   public:: free_energy_chain       ! Split chain FE components (conf/transl)
   public:: free_energy_homog       ! Helmholtz energy of homogenous phase
   public:: set_omega_uniform       ! sets k=0 component of omega (canonical)
   public:: elec_startup            ! Initializes aminusb and qadotqb arrays
   public:: calc_dielectric         ! calculates eps from rho
   public:: calc_electric_potential ! calculates psi from eps and rho
   public:: calc_squared_gradient   ! Given field A, calculates |\del A|^2
   public:: field_inverse           ! Given field A, calculates 1/A
   public:: field_product           ! Given fields A and B, calculates A*B

   ! Public Test Functions
   public:: check_poisson_hermitian ! checks if poisson's green func. is hermitian
   public:: check_poisson_duality   ! checks if poisson's green func. has duality

   !# ifdef DEVEL
   public:: divide_energy     ! calculates different components of free energy
   !# endif

   ! public module variable 
   public:: plan              ! module variable, used in iterate_mod
   !***

   type(fft_plan)                     :: plan
   type(chain_grid_type),allocatable  :: chains(:)
   integer                            :: extrap_order
   real(long), allocatable            :: aminusbvecs(:,:,:)
   real(long), allocatable            :: qadotqb(:,:)
   real(long), allocatable            :: qaminbdotqa(:,:)

   !****v scf_mod/plan -------------------------------------------------
   ! VARIABLE
   !     type(fft_plan) plan - Plan of grid sizes etc. used for FFTs
   !                           (Public because its used in iterate_mod)
   !*** ----------------------------------------------------------------

contains

   !--------------------------------------------------------------------
   !****p scf_mod/density_startup
   ! SUBROUTINE
   !    subroutine density_startup(N_grids,extr_order,chain_step,update_chain)
   !
   ! PURPOSE
   !    Initialize FFT_plan, grid_data.
   !    Allocate or update/re-allocate memory for chains
   !
   ! ARGUMENTS
   !    N_grids      = grid dimensions
   !    extr_order   = Richardson extrapolation order
   !    chain_step   = the discretized chain segment length
   !    update_chain = true if simply the chain memory need to be re-allocated
   !
   ! SOURCE
   !--------------------------------------------------------------------
   subroutine density_startup(N_grids, extr_order, chain_step, update_chain)
   implicit none

   integer, intent(IN)    :: N_grids(3) ! # of grid points in each direction
   integer, intent(IN)    :: extr_order
   real(long), intent(IN) :: chain_step
   logical, intent(IN)    :: update_chain
   !***

   integer :: i, nblk, error
   integer :: nx,ny,nz    
   if ( .NOT. update_chain ) then
      call create_fft_plan(N_grids,plan)

      if (N_chain > 0) then
         allocate(chains(N_chain),STAT=error)
         if(error /= 0) STOP "chains allocation error in scf_mod/density_startup"
         do i=1, N_chain
           nblk      = N_block(i)
           call null_chain_grid(chains(i))
           call make_chain_grid(chains(i),plan,nblk,&
                       block_length(1:nblk,i),chain_step,allocate_q=.TRUE.)
         end do
      end if
      call init_step(N_grids)
      extrap_order = extr_order   ! set up global variable
   else
      do i=1, N_chain
        nblk = N_block(i)
        call destroy_chain_grid(chains(i))
        call make_chain_grid(chains(i),plan,nblk,&
                    block_length(1:nblk,i),chain_step,allocate_q=.TRUE.)
      end do
   end if
   
   end subroutine density_startup
   !====================================================================


   !--------------------------------------------------------------------
   !****p scf_mod/elec_startup
   ! SUBROUTINE
   !    subroutine elec_startup()
   !
   ! PURPOSE
   !    Allocate and Initialize aminusbvecs, qadotqb
   !    Save time on repeat calculations for Poisson Operator
   !
   ! ARGUMENTS
   !
   ! SOURCE
   !--------------------------------------------------------------------
   subroutine elec_startup(elec_init_flag)
   use unit_cell_mod
   use basis_mod

   ! Inputs
   logical, intent(IN)    :: elec_init_flag !True if initialization needed


   integer    :: info               ! Allocation Check
   integer    :: abz(3), bbz(3)     ! First Brilloun Zone indices for a and b
   integer    :: afft(3),bfft(3)    ! (0:N(1)-1,...), natural order of fft
   real(long) :: avec(3),bvec(3)    ! reciprocal basis vectors for alpha, beta
   integer    :: aminusb(3)         ! avec-bvec, dvec is converted to real grid
   integer    :: aplusb(3)          ! avec+bvec, evec is converted to real grid
   integer    :: a1,a2,a3,b1,b2,b3  ! loop indices
   integer    :: alpha,beta,i,j     ! More indices
   real(long) :: qdotprod, qmindotprod
   integer    :: dfft(3),dbz(3)
   real(long) :: dvec(3)

   ! Quit if elec_init_flag is false
   if (.not.elec_init_flag) then
      return
   endif 

   ! Check for old arrays, clear if they're there
   if(allocated(aminusbvecs)) deallocate(aminusbvecs)
   if(allocated(qadotqb)) deallocate(qadotqb)
   if(allocated(qaminbdotqa)) deallocate(qaminbdotqa)

   ! Initialize Arrays
   allocate(aminusbvecs(N_wave,N_wave,3), stat = info )
   if (info.ne.0) stop 'Error allocating aminusbvecs'
   allocate(qadotqb(N_wave,N_wave), stat = info )
   if (info.ne.0) stop 'Error allocating qadotqb'
   allocate(qaminbdotqa(N_wave,N_wave), stat = info )
   if (info.ne.0) stop 'Error allocating qaminbdotqa'

   qadotqb(:,:) = 0.0_long
   qaminbdotqa(:,:) = 0.0_long

   ! Loop over q_alpha
   do alpha = 1, N_wave
      abz = wave(:,alpha)
      avec(:) = 0.0_long
      do i = 1, dim
         do j = 1, dim
            avec(j) = avec(j) + abz(i)*G_basis(i,j) 
         enddo 
      enddo

      ! Loop over q_beta
      do beta = 1, N_wave
         bbz = wave(:,beta)
         bvec(:) = 0.0_long
         do i = 1, dim
            do j = 1, dim
               bvec(j) = bvec(j) + bbz(i)*G_basis(i,j) 
            enddo 
         enddo

         ! qDot calculation
         qdotprod = 0.0_long
         do i = 1, dim
             if ( (MOD(ngrid(i),2)==0).AND.( (abz(i)==ngrid(i)/2) &
                  .OR.(bbz(i)==ngrid(i)/2)) )  then
                qdotprod = qdotprod + ABS(avec(i)*bvec(i))
             else
                qdotprod = qdotprod + avec(i)*bvec(i)
             endif
         end do

         ! difference vector
         afft = G_to_fft(abz)
         bfft = G_to_fft(bbz)
         aminusb = afft - bfft
         dfft = G_to_fft(aminusb)
         
         dbz = G_to_bz(dfft)
         dvec(:) = 0.0_long
         do i = 1, dim
            do j = 1, dim
               dvec(j) = dvec(j) + dbz(i)*G_basis(i,j) 
            enddo 
         enddo
        
         ! qmindotprod calculation
         qmindotprod = 0.0_long
         do i = 1, dim
             if ( (MOD(ngrid(i),2)==0).AND.( (dbz(i)==ngrid(i)/2) &
                  .OR.(abz(i)==ngrid(i)/2)) )  then
                qmindotprod = qmindotprod + ABS(dvec(i)*avec(i))
             else
                qmindotprod = qmindotprod + dvec(i)*avec(i)
             endif
         end do

         aminusbvecs(alpha,beta,1) = dfft(1)
         aminusbvecs(alpha,beta,2) = dfft(2)
         aminusbvecs(alpha,beta,3) = dfft(3)
         qadotqb(alpha,beta) = qdotprod
         qaminbdotqa(alpha,beta) = qmindotprod
      enddo
   enddo
 
   end subroutine elec_startup


   !-----------------------------------------------------------------
   !****p scf_mod/density
   ! SUBROUTINE
   !    density(N,omega,rho,qout,q_solvent)
   !
   ! PURPOSE
   !    Main SCFT calculation. Solve the modified diffusion equation
   !    for all polymer species, and calculate monomer density field
   !    for all monomer types. 
   !
   ! ARGUMENTS
   !    N                    = # of basis functions
   !    omega(N_monomer,N)   = chemical potential
   !    rho(N_monomer,N)     = monomer density fields
   !    qout(N_chain)        = partition functions
   !    q_solvent(N_solvent) = partition functions of solvent molecules
   !
   ! COMMENT   
   !      density_startup should be called prior to density to
   !      allocate arrays used by density and scf_stress.
   !
   ! SOURCE
   !------------------------------------------------------------------
   subroutine density( &
       N,              & ! # of basis functions
       omega,          & ! (N_monomer,N)chemical potential field
       rho,            & ! (N_monomer,N) monomer density field
       qout,           & ! (N_chain) 1-chain partition functions
       q_solvent       & ! (N_solvent) solvent partition functions
   )
   implicit none

   integer,    intent(IN)            :: N
   real(long), intent(IN)            :: omega(:,:)
   real(long), intent(OUT)           :: rho(:,:)
   real(long), intent(OUT), optional :: qout(N_chain)
   real(long), intent(OUT), optional :: q_solvent(N_solvent)
   !***
   
   ! local variables
   !complex(long)  :: kgrid(0:plan%n(1)/2, &
   !                        0:plan%n(2)-1, &
   !                        0:plan%n(3)-1)
   complex(long),allocatable  :: kgrid(:,:,:)

   real(long)     :: rnodes       ! number of grid points
   real(long)     :: partion      ! partion of single chain
   real(long)     :: bigQ_solvent ! partition of solvent
   integer        :: i_chain      ! index to chain
   integer        :: i_blk        ! index to block
   integer        :: alpha        ! index to monomer
   integer        :: i            ! dummy variable
   real(long)     :: Ns           ! number of solvent molecules in a reference volume  
   integer        :: info
   
   allocate( kgrid(0:plan%n(1)/2, 0:plan%n(2)-1, 0:plan%n(3)-1), stat=info)
   if( info /= 0 ) stop "density/kgrid(:,:,:) allocation error"
  
   rnodes=dble( plan%n(1) * plan%n(2) * plan%n(3) )

   ! Transform omega fields onto a grid
   do alpha = 1, N_monomer
     call basis_to_kgrid(omega(alpha,:),kgrid)
     call ifft(plan,kgrid,omega_grid(:,:,:,alpha))
   end do
   
   ! loop over chains
   do i_chain = 1, N_chain
     call chain_density(i_chain,chains(i_chain),ksq_grid,omega_grid)
     if(present(qout)) qout(i_chain) = chains(i_chain)%bigQ
   end do
  
   ! takes into account solvent monomer densities
   rho_grid = 0.0_long
   do i=1, N_solvent
      alpha = solvent_monomer(i)
      Ns  = solvent_size(i)   ! = solvent volume / reference volume
      CALL solvent_density(alpha,Ns,omega_grid,rho_grid,bigQ_solvent)
      if(present(q_solvent)) q_solvent(i) = bigQ_solvent
   end do
  
   if (ensemble ==1) then 
      if (present(qout)) then
         call mu_phi_chain(mu_chain,phi_chain,qout)
      end if
      if (present(q_solvent)) then
         call mu_phi_solvent(mu_solvent,phi_solvent,q_solvent)
      end if
   end if
 
   ! calculate monomer densities
   do i_chain = 1, N_chain
     do i_blk = 1, N_block(i_chain)
         alpha = block_monomer(i_blk,i_chain)
         rho_grid(:,:,:,alpha) = rho_grid(:,:,:,alpha) &
             + phi_chain(i_chain) * chains(i_chain)%rho(:,:,:,i_blk)
     end do
   end do
  
   ! project monomer densities onto basis functions
   do alpha=1, N_monomer
     call fft(plan,rho_grid(:,:,:,alpha),kgrid)
     call kgrid_to_basis(kgrid,rho(alpha,:))
     rho(alpha,:)=rho(alpha,:)/rnodes
   end do

   if( allocated(kgrid) ) deallocate(kgrid)

   end subroutine density
   !=======================================================
  

   !--------------------------------------------------------------------------
   !****p scf_mod/solvent_density
   ! SUBROUTINE 
   !    solvent_density(monomer,s_size,omega,rho_grid,bigQ_solvent)
   !
   ! PURPOSE
   !    to calculate the density profile of a  solvent specie
   !
   ! ARGUMENTS
   !    monomer      - monomer type of the solvent species
   !    s_size       - volume occupied by solvent molecule / reference volume
   !                   (volume in units where reference volume = 1)
   !    omega        - omega fields on grid, per reference volume
   !    rho_grid     - density fields on grid    
   !    bigQ_solvent - spatial average of Boltzmann factor exp(-s_size*omega)
   !
   ! SOURCE
   !--------------------------------------------------------------------------
   subroutine solvent_density(monomer,s_size,omega,rho_grid,bigQ_solvent)
   implicit none
   
   real(long),intent(IN)              :: s_size
   real(long),intent(IN)              :: omega(0:,0:,0:,:)
   integer,intent(IN)                 :: monomer
   real(long),intent(INOUT)           :: rho_grid(0:,0:,0:,:)
   real(long),intent(OUT)             :: bigQ_solvent          
   !***
   
   real(long):: rnodes

   integer   :: ix,iy,iz,i    ! loop indices
   integer   :: solvent       ! solvent species index in phi array
   integer   :: error
 
   rnodes = dble(ngrid(1) * ngrid(2) * ngrid(3))
  
   ! calculating bigQ_solvent
   bigQ_solvent = 0.0_long  
   do iz=0, ngrid(3)-1 
      do iy=0, ngrid(2)-1
         do ix=0, ngrid(1)-1
            
            bigQ_solvent = bigQ_solvent + EXP((-s_size)&
                                              * omega(ix,iy,iz,monomer))
          
         end do
      end do
   end do     

   bigQ_solvent = bigQ_solvent/dble(rnodes)
      
   ! calculating the index of the solvent in the phi array
   do i=1, N_solvent
      if (solvent_monomer(i)==monomer) then
         solvent = i
      end if
      if ( ensemble == 1 )   phi_solvent(solvent) = bigQ_solvent*exp(mu_solvent(solvent))
   end do
 
   rho_grid(:,:,:,monomer) = rho_grid(:,:,:,monomer) + phi_solvent(solvent) * &
                             EXP((-s_size) * omega(:,:,:,monomer))/bigQ_solvent
 
   end subroutine solvent_density
   !====================================================================
  

   !--------------------------------------------------------------------
   !****p scf_mod/chain_density
   ! SUBROUTINE
   !    chain_density(i_chain, chain, ksq, omega)
   !
   ! PURPOSE
   !    solve the PDE for a single chain
   !    evaluate the density for each block
   !
   ! ARGUMENTS
   !    i_chain - index to the chain 
   !    chain   - chain_grid_type, see chain_mod
   !    ksq     - k^2 on grid, initialized in grid_mod
   !    omega   - omega fields on grid
   ! SOURCE   
   !--------------------------------------------------------------------
   subroutine chain_density(i_chain, chain, ksq, omega)
   implicit none

   integer,intent(IN)                   :: i_chain
   type(chain_grid_type),intent(INOUT)  :: chain
   real(long),intent(IN)                :: ksq(0:,0:,0:)
   real(long),intent(IN)                :: omega(0:,0:,0:,:)
   !***

   integer   :: chain_end, i_blk
   integer   :: istep, ibgn, iend
   real(long):: ds, b
   integer   :: i, j, monomer
   integer   :: ix, iy, iz

   !real(long):: wscale
   !real(long):: scale_offset = 5.0_long

   ! Stiff: Find omega-offset (Scale to -5)
   !wscale = 0.0_long
   !do i_blk = 1, N_block(i_chain)
   !   monomer = block_monomer(i_blk, i_chain)
   !   if (MINVAL(omega(:,:,:,monomer)) < wscale - scale_offset) then
   !       wscale = MINVAL(omega(:,:,:,monomer)) + scale_offset
   !   endif
   !end do

   ! Calculate qf, by integratin forward from s=0
   chain%qf(:,:,:,:) = 0.0_long
   chain%qf(:,:,:,1) = 1.0_long
  
   do i_blk = 1, N_block(i_chain)
      monomer = block_monomer(i_blk, i_chain)
      ds = chain%block_ds(i_blk)
      b  = kuhn( monomer )
      !call make_propg(ds, b, ksq, (omega(:,:,:,monomer)-wscale) )
      call make_propg(ds, b, ksq, omega(:,:,:,monomer) )

      do istep = chain%block_bgn(i_blk), chain%block_bgn(i_blk+1)-1
         ! print*, '(iblk,istep):  (', i_blk, ',', istep, ')'
          !print*, 'zero element', chain%qf(0,0,0,istep)
          !print*, 'min element', MINVAL(chain%qf(:,:,:,istep))
          call step_exp(chain%qf(:,:,:,istep), &
                       chain%qf(:,:,:,istep+1), plan)
      end do
   end do

   ! Calculate qr, by integrating backward from s = chain_end
   chain_end = chain%block_bgn(N_block(i_chain)+1)
   chain%qr(:,:,:,:) = 0.0_long
   chain%qr(:,:,:,chain_end) = 1.0_long

   do i_blk = N_block(i_chain), 1, -1
      monomer = block_monomer(i_blk,i_chain)
      ds = chain%block_ds(i_blk)
      b  = kuhn( monomer )
      !call make_propg(ds, b, ksq, (omega(:,:,:,monomer)-wscale) )
      call make_propg(ds, b, ksq, omega(:,:,:,monomer) )

      do istep = chain%block_bgn(i_blk+1), chain%block_bgn(i_blk)+1, -1
         call step_exp(chain%qr(:,:,:,istep), &
                       chain%qr(:,:,:,istep-1), plan)
      end do
   end do

   ! Calculate single chain partition function chain%bigQ
   chain%bigQ = sum(chain%qf(:,:,:,chain_end)) &
          / dble(size(chain%qf(:,:,:,chain_end))) 
   
   !print*, 'wscale: ', wscale
   !print*, 'Big Q:  ', chain%bigQ

   ! Error Catching
   if (MINVAL(chain%qf(:,:,:,:)) < 0.0 ) then
      print*, 'ERROR: Chain propagator qf contains negative value'
      print*, MINVAL(chain%qf(:,:,:,:))
      !print*, 'Q: ', chain%bigQ
      !print*, chain%qf(0,0,0,:)
      CALL EXIT
   endif
   if (MINVAL(chain%qr(:,:,:,:)) < 0.0 ) then
      print*, 'ERROR: Chain propagator qr contains negative value'
      print*, MINVAL(chain%qr(:,:,:,:))
      CALL EXIT
   endif

   ! Calculate monomer concentration fields, using Simpson's rule
   ! to evaluate the integral \int ds qr(r,s)*qf(r,s)
   chain%rho = 0.0_long
   do i = 1, N_block(i_chain)
      ! Chain ends: Add qf(r,ibgn)*qr(r,ibgn) & qf(r,iend)*qr(r,iend)
      ibgn=chain%block_bgn(i)
      iend=chain%block_bgn(i+1)

!     chain%rho(:,:,:,i)=chain%qf(:,:,:,ibgn)*chain%qr(:,:,:,ibgn)
      do iz=0,ngrid(3)-1
      do iy=0,ngrid(2)-1
      do ix=0,ngrid(1)-1
        chain%rho(ix,iy,iz,i)=chain%qf(ix,iy,iz,ibgn)*chain%qr(ix,iy,iz,ibgn)
      end do
      end do
      end do

!     chain%rho(:,:,:,i)=chain%rho(:,:,:,i)+chain%qf(:,:,:,iend)*  &
!                               chain%qr(:,:,:,iend)
      do iz=0,ngrid(3)-1
      do iy=0,ngrid(2)-1
      do ix=0,ngrid(1)-1
        chain%rho(ix,iy,iz,i)=chain%rho(ix,iy,iz,i) + &
            chain%qf(ix,iy,iz,iend)*chain%qr(ix,iy,iz,iend)
      end do
      end do
      end do

      ! Odd indices: Sum values of qf(i)*qr(i)*4.0 with i odd
      do j=ibgn+1,iend-1,2
!        chain%rho(:,:,:,i)=chain%rho(:,:,:,i)+chain%qf(:,:,:,j)*  &
!                          chain%qr(:,:,:,j)*4.0_long
         do iz=0,ngrid(3)-1
         do iy=0,ngrid(2)-1
         do ix=0,ngrid(1)-1
           chain%rho(ix,iy,iz,i)=chain%rho(ix,iy,iz,i) + &
               chain%qf(ix,iy,iz,j)*chain%qr(ix,iy,iz,j)*4.0_long
         end do
         end do
         end do

      end do

      ! Even indices: Sum values of qf(i)*qr(i)*2.0 with i even
      do j=ibgn+2,iend-2,2
!        chain%rho(:,:,:,i)=chain%rho(:,:,:,i)+chain%qf(:,:,:,j)*  &
!                            chain%qr(:,:,:,j)*2.0_long
         do iz=0,ngrid(3)-1
         do iy=0,ngrid(2)-1
         do ix=0,ngrid(1)-1
           chain%rho(ix,iy,iz,i)=chain%rho(ix,iy,iz,i) + &
               chain%qf(ix,iy,iz,j)*chain%qr(ix,iy,iz,j)*2.0_long
         end do
         end do
         end do

      end do

      ! Multiply sum by ds/3
      chain%rho(:,:,:,i)=chain%rho(:,:,:,i)*chain%block_ds(i)/3.0_long  
   end do

   chain%rho=chain%rho/chain_length(i_chain)/chain%bigQ
   
   ! Renormalize bigQ
   !chain%bigQ = chain%bigQ/EXP(wscale)
   !print*, 'BigQ: ', chain%bigQ

   end subroutine chain_density
   !====================================================================


   !--------------------------------------------------------------------
   !****p scf_mod/scf_stress
   ! FUNCTION
   !    scf_stress(N, size_dGsq, dGsq )
   !
   ! RETURN
   !    real(long) array of dimension(size_dGsq) containing
   !    derivatives of free energy with respect to size_dGsq 
   !    cell parameters or deformations
   !
   ! ARGUMENTS
   !    N         = number of basis functions
   !    size_dGsq = number of cell parameters or deformations
   !    dGsq      = derivatives of |G|^2 w.r.t. cell parameters
   !                dGsq(i,j) = d |G(i)|**2 / d cell_param(j)
   ! COMMENT
   !    Requires previous call to density, because scf_stress
   !    uses module variables computed in density.
   !
   ! SOURCE
   !--------------------------------------------------------------------
   function scf_stress(N, size_dGsq, dGsq )
   implicit none

   integer,    intent(IN) :: N
   integer,    intent(IN) :: size_dGsq
   real(long), intent(IN) :: dGsq(:,:)
   !***

   real(long)  :: scf_stress(size_dGsq)

   ! ngrid(3) was obtained by association
   ! Local Variables

   real(long)      :: dQ(size_dGsq)    ! change in q
   real(long)      :: qf_basis(N),qr_basis(N),q_swp(N)
   !complex(long)   :: kgrid(0:ngrid(1)/2,0:ngrid(2)-1,0:ngrid(3)-1)
   complex(long),allocatable   :: kgrid(:,:,:)

   real(long)      :: rnodes, normal
   real(long)      :: ds0, ds, b
   real(long)      :: increment
   integer         :: i, alpha, beta   ! summation indices
   integer         :: monomer             ! monomer index
   integer         :: sp_index            ! species index
   integer         :: ibgn,iend
   integer         :: info

   allocate( kgrid(0:ngrid(1)/2, 0:ngrid(2)-1, 0:ngrid(3)-1), stat=info )
   if ( info /= 0 ) stop "scf_mod/scf_stress/kgrid(:,:,:) allocation error"

   ! number of grid points
   rnodes = dble( ngrid(1) * ngrid(2) * ngrid(3) )

   ! normal = rnodes  * &! normalization of bigQ, divided by volume
   normal = rnodes   *  &! fft normal of forward partition
            rnodes   *  &! fft normal of backward partition
            3.0_long *  &! normal simpson's rule
            6.0_long     ! b**2/6

   scf_stress = 0.0_long

   !print*, '==== TESTCODE (scfstress) ===='
   !print*, '---Iteration 0--'
   !print*, 'dQ: ', dQ
   !print*, 'big Q: ', chains(1)%bigQ
   !print*, '-----'

   ! Loop over chain species
   do sp_index = 1, N_chain
      dQ = 0.0_long

      ! Loop over blocks
      do alpha = 1,  N_block(sp_index) 
         monomer = block_monomer(alpha,sp_index)
               b = kuhn(monomer)
             ds0 = chains(sp_index)%block_ds(alpha)

            ibgn = chains(sp_index)%block_bgn(alpha)
            iend = chains(sp_index)%block_bgn(alpha+1)

            !print*, '---simple stuff--'
            !print*, 'ds0: ',   chains(1)%block_ds(alpha)
            !print*, 'ibgn: ', chains(1)%block_bgn(alpha)
            !print*, 'iend: ', chains(1)%block_bgn(alpha+1)
            !print*, '-----'

         do i = ibgn, iend
            ! rgrid=dcmplx( chains(sp_index)%qf(:,:,:,i), 0.0_long)
            call fft(plan, chains(sp_index)%qf(:,:,:,i), kgrid )
            call kgrid_to_basis( kgrid, qf_basis )

            ! rgrid=dcmplx( chains(sp_index)%qr(:,:,:,i), 0.0_long)
            call fft(plan, chains(sp_index)%qr(:,:,:,i), kgrid )
            call kgrid_to_basis( kgrid, qr_basis )

            ds = ds0
            if ( i/= ibgn .and. i/= iend) then
               if (modulo(i,2) == 0) then
                  ds = 4.0_long * ds
               else
                  ds = 2.0_long * ds
               end if
            end if   ! Simpson's rule quadrature

            do beta = 1, size_dGsq
               q_swp     = qr_basis * dGsq(:,beta)
               increment = dot_product(q_swp, qf_basis)
               increment = increment * b**2 * ds / normal
               dQ(beta)  = dQ(beta) - increment
            end do

         end do      ! loop over nodes of single block
      end do         ! loop over blocks


      ! Note the mixing rule
      ! stress(total) = \sum_alpha \phi_alpha \cdot~stress(\alpha)


      select case(ensemble)
      case (0)
         scf_stress = scf_stress - (dQ / chains(sp_index)%bigQ)*  &
                      phi_chain(sp_index)/chain_length(sp_index)
      case (1)
         scf_stress = scf_stress - (dQ / chains(sp_index)%bigQ)*  &
                      exp(mu_chain(sp_index))*chains(sp_index)%bigQ  / &
                      chain_length(sp_index)
      end select
      
      !print*, 'dQ: ', dQ
      !print*, 'big Q: ', chains(1)%bigQ
      !print*, 'stress: ', scf_stress 
      !print*, '========  END TEST  ======='

   end do

   if ( allocated(kgrid) ) deallocate( kgrid )

   end function scf_stress
   !===================================================================


   !# ifdef DEVEL
   !-------------------------------------------------------------------
   !****p scf_mod/divide_energy
   ! SUBROUTINE
   !    divide_energy(rho, omega, phi_chain, phi_solvent, Q, f_comp, ovlap)
   ! PURPOSE   
   !    Divide free energy into components arising from binary
   !    interaction free energy and from chain entropy
   ! ARGUMENTS
   !    rho         = density fields
   !  omega         = potential fields
   !    phi_chain  = volume fraction of species (chain)
   !    phi_solvent = volume fraction of species (solvent)
   !      Q         = partion function of species (chain)
   !  f_tot         = total free energy
   ! f_comp         = components of free energy (see below)
   !  ovlap         = overlap integrals
   ! COMMENT
   !
   !    a) Components of f_comp array:
   !       f_comp(1) = overall interaction energy
   !       f_comp(2) = conformational energy of first block
   !       f_comp(3) = conformational energy of last  block
   !       f_comp(4) = junction translational energy (diblock)
   !
   !    b) Calculation of junction translational entropy is correct
   !       only for diblocks, for which there is only one junction
   !
   !    c) Components of overlap integral array ovlap can be used
   !       to divide interaction energy into components arising from
   !       interactions between specific pairs of monomer types.
   !
   ! SOURCE
   !----------------------------------------------------------------
   subroutine divide_energy(rho, omega, phi_chain, phi_solvent, Q, f_tot, f_comp, ovlap)
   implicit none
   real(long), intent(IN)  :: rho(:,:)        ! monomer vol. frac fields
   real(long), intent(IN)  :: omega(:,:)      ! chemical potential field
   real(long), intent(IN)  :: phi_chain(:)    ! molecule vol. frac of chain mol
   real(long), intent(IN)  :: phi_solvent(:)  ! molecule vol. frac of solvent mol
   real(long), intent(IN)  :: Q(:)            ! chain partition functions
   real(long), intent(IN)  :: f_tot           ! components of free energy
   real(long), intent(OUT) :: f_comp(:)       ! components of free energy
   real(long), intent(OUT) :: ovlap(:,:)      ! overlap integrals,N_monomer**2
   !***

   real(long) :: rnodes
   real(long) :: enthalpy    ! interaction energy
   real(long) :: fhead       ! head block energy
   real(long) :: ftail       ! tail block energy
   real(long) :: fjct        ! junction translational entropy
   real(long) :: ftmp        ! junction translational entropy (temporary)
   integer    :: alpha, beta ! monomer indices
   integer    :: i, nh, nt   ! loop indices
   integer    :: ix,iy,iz    ! loop indices

   rnodes = dble( ngrid(1) * ngrid(2) * ngrid(3) )

   enthalpy = 0.0_long
   ovlap    = 0.0_long
   do alpha = 1, N_monomer
      do beta = alpha+1, N_monomer
         ovlap(alpha,beta) = dot_product(rho(alpha,:),rho(beta,:))
         ovlap(beta,alpha) = ovlap(alpha,beta)
         enthalpy = enthalpy + chi(alpha,beta) * ovlap(alpha,beta)
      end do
   end do

   fhead = 0.0_long
   ftail = 0.0_long
   do i=1, N_chain
      nh = chains(i)%block_bgn(2)
      do iz = 0, ngrid(3)-1
      do iy = 0, ngrid(2)-1
      do ix = 0, ngrid(1)-1
         if ( chains(i)%qf(ix,iy,iz,nh) > 0.0_long .AND. &
              chains(i)%qr(ix,iy,iz,nh) > 0.0_long ) then
            fhead = fhead - chains(i)%qf(ix,iy,iz,nh)   &
                          * chains(i)%qr(ix,iy,iz,nh)   &
                          / Q(i)                        &
                          * log( chains(i)%qf(ix,iy,iz,nh) ) & 
                          * phi_chain(i) / chain_length(i)
         end if
      end do
      end do
      end do

      nt = chains(i)%block_bgn(N_block(i))
      do iz = 0, ngrid(3)-1
      do iy = 0, ngrid(2)-1
      do ix = 0, ngrid(1)-1
         if ( chains(i)%qf(ix,iy,iz,nt) > 0.0_long .AND. &
              chains(i)%qr(ix,iy,iz,nt) > 0.0_long ) then
            ftail = ftail - chains(i)%qf(ix,iy,iz,nt)   &
                          * chains(i)%qr(ix,iy,iz,nt)   &
                          / Q(i)                        &
                          * log( chains(i)%qr(ix,iy,iz,nt) ) &
                          * phi_chain(i) / chain_length(i)
         end if
      end do
      end do
      end do
   end do
   fhead = fhead / rnodes
   ftail = ftail / rnodes

   ! When monomer types in the middle blocks are different from either
   ! head or tail block, the subtraction below is correct.
   do i=1, N_chain
      beta=block_monomer(1,i)
      fhead = fhead - dot_product(omega(beta,:),rho(beta,:)) * phi_chain(i)

      beta=block_monomer(N_block(i),i)
      ftail = ftail - dot_product(omega(beta,:),rho(beta,:)) * phi_chain(i)
   end do

   ! --------------------------------------------
   ! The following block was used to calculate
   ! junction entropy contribution to free
   ! energy of diblocks. 
   ! Since it is not universal, we now instead
   ! calculate by subtraction, which can be
   ! interpreted by excess entropies for arbitrary
   ! molecular (linear) architecture.
   ! --------------------------------------------
   !## fjct = 0.0_long
   !## do i=1, N_chain
   !##    nh = chains(i)%block_bgn(2)
   !##    do iz = 0, ngrid(3)-1
   !##    do iy = 0, ngrid(2)-1
   !##    do ix = 0, ngrid(1)-1
   !##       if ( chains(i)%qf(ix,iy,iz,nh) > 0.0_long .AND. &
   !##            chains(i)%qr(ix,iy,iz,nh) > 0.0_long ) then
   !##          ftmp  = chains(i)%qf(ix,iy,iz,nh) &
   !##                * chains(i)%qr(ix,iy,iz,nh) &
   !##                / Q(i) 
   !##          fjct = fjct + ftmp * log( ftmp ) * phi_chain(i) / chain_length(i)
   !##       end if
   !##    end do
   !##    end do
   !##    end do
   !## end do
   !## fjct  = fjct  / rnodes
   ! --------------------------------------------
   fjct = 0.0_long
   do i=1, N_chain
      fjct = fjct + phi_chain(i) / chain_length(i)
   end do
   fjct = f_tot + fjct - enthalpy - fhead - ftail

   f_comp(1) = enthalpy   ! overall interaction energy
   f_comp(2) = fhead      ! conformational energy of first block
   f_comp(3) = ftail      ! conformational energy of last  block
   f_comp(4) = fjct       ! junction translational energy (diblock)

   end subroutine divide_energy
   !=============================================================
   !# endif


   !------------------------------------------------------------
   !****p scf_mod/set_omega_uniform
   ! SUBROUTINE
   !    set_omega_uniform(omega)
   ! PURPOSE
   !   Sets uniform (k=0) component of field omega to convention
   !   corresponding to vanishing Lagrange multiplier field
   ! SOURCE
   !------------------------------------------------------------
   subroutine set_omega_uniform(omega, electrostatic_flag)

   real(long), intent(INOUT) :: omega(:,:)
   logical, intent(IN)       :: electrostatic_flag

   !***

   integer    :: i, j, alpha, beta  ! loop indices
   real(long) :: phi_mon(N_monomer) ! average monomer vol. frac.
   real(long) :: phi_eps_sum        ! Normalizes the dielectric average
   real(long) :: eps_avg            ! Average (homogenous) dielectric constant
   real(long) :: borne_mon          ! l0/2eps^2 (l0 is the Bjerrum length)
   real(long) :: borne_ion          ! l0/2eps

   phi_mon = 0.0_long

   do i = 1, N_chain
      do j = 1, N_block(i)
         alpha = block_monomer(j,i)
         phi_mon(alpha) = phi_mon(alpha) &
                        + phi_chain(i)*block_length(j,i)/chain_length(i)
      end do
   end do

   do i = 1, N_solvent
      alpha = solvent_monomer(i)
      phi_mon(alpha) = phi_mon(alpha) + phi_solvent(i)
   end do

   do alpha = 1, N_monomer
      omega(alpha,1) = 0.0_long
      do beta = 1, N_monomer
         omega(alpha,1) = omega(alpha,1) &
              + chi(alpha,beta) * phi_mon(beta)
      end do
   end do

   if (electrostatic_flag) then  
      ! Calculate the homogenous dielectric constant
      phi_eps_sum = 0.0_long
      eps_avg = 0.0_long
      do alpha = 1, N_monomer
         if (dielectric_monomer(alpha) > 0) then
            phi_eps_sum = phi_eps_sum + phi_mon(alpha)
            eps_avg = eps_avg & 
                 + dielectric_monomer(alpha)* phi_mon(alpha)
         endif
      enddo

      eps_avg = eps_avg/phi_eps_sum

      borne_mon = l0/(2*eps_avg*eps_avg)
      borne_ion = l0/(2*eps_avg)

      do alpha = 1, N_monomer
         ! Loop over neutral monomers
         if (dielectric_monomer(alpha) > 0) then
            do beta = 1, N_monomer
               if (ionic_radii_monomer(beta) > 0) then
                  omega(alpha,1) = omega(alpha,1) &
                    + ((-1)**alpha)*borne_mon &
                    * (dielectric_monomer(1) - dielectric_monomer(2)) &
                    * ( phi_mon(1+MOD(alpha,2)) / ( (phi_mon(1)+phi_mon(2))**2) ) &
                    * (phi_mon(beta)/ionic_radii_monomer(beta)) &
                    * (1/solvent_size(beta-N_monomer+N_solvent))             
                  ! - dielectric_monomer(alpha)*borne_mon &
                  ! * (phi_mon(beta)/ionic_radii_monomer(beta)) &
                  ! * (1/solvent_size(beta-N_monomer+N_solvent))
               endif
            enddo
         endif

         ! Loop over ions
         if (ionic_radii_monomer(alpha) > 0) then
            omega(alpha,1) = omega(alpha,1) &
               + borne_ion*(1/ionic_radii_monomer(alpha)) &
               * (1/solvent_size(alpha-N_monomer+N_solvent))
         endif
      end do
   endif

   end subroutine set_omega_uniform
   !================================================================


   !-------------------------------------------------------------
   !****p scf_mod/mu_phi_chain
   ! SUBROUTINE
   !    mu_phi_chain(mu, phi, q)
   ! PURPOSE
   !    If ensemble = 0 (canonical), calculate mu from phi
   !    If ensemble = 1 (grand), calculate phi from mu
   ! ARGUMENTS
   !    mu(N_chain)  = chain chemical potentials (units kT=1)
   !    phi(N_chain) = chain molecular volume fractions 
   !    q(N_chain)   = single chain partition functions
   !
   ! SOURCE
   !-------------------------------------------------------------
   subroutine mu_phi_chain(mu, phi, q)
   real(long), intent(INOUT) :: mu(N_chain)
   real(long), intent(INOUT) :: phi(N_chain) 
   real(long), intent(IN)    :: q(N_chain)
   !***

   integer :: i
   select case(ensemble)
   case (0)
      do i = 1, N_chain
         mu(i) = log( phi(i) / q(i) )
      end do
   case (1)
      do i = 1, N_chain
         phi(i) = q(i)*exp(mu(i))
      end do
   end select
   end subroutine mu_phi_chain
   !================================================================


   !-------------------------------------------------------------
   !****p scf_mod/mu_phi_solvent
   ! SUBROUTINE
   !    mu_phi_solvent(mu, phi, q)
   ! PURPOSE
   !    If ensemble = 0 (canonical), calculate mu from phi
   !    If ensemble = 1 (grand can), calculate phi from mu
   ! ARGUMENTS
   !    mu(N_solvent)  = solvent chemical potentials 
   !    phi(N_solvent) = solvent volume fractions
   !    q(N_solvent)   = solvent partition functions
   !
   ! SOURCE
   !-------------------------------------------------------------
   subroutine mu_phi_solvent(mu, phi, q)
   real(long), intent(INOUT) :: mu(N_solvent)
   real(long), intent(INOUT) :: phi(N_solvent)
   real(long), intent(IN)    :: q(N_solvent) 
   !***

   integer :: i
   select case(ensemble)
   case (0)
      do i = 1, N_solvent
         mu(i) = log(phi(i) / q(i))  
      end do
   case (1)
      do i = 1, N_solvent
         phi(i) = q(i)*exp(mu(i))
      end do
   end select
   end subroutine mu_phi_solvent
   !================================================================


   !--------------------------------------------------------------------
   !****p scf_mod/free_energy
   ! SUBROUTINE
   !    free_energy( N, rho, omega, phi_chain, mu_chain, phi_solvent,
   !                 mu_solvent, f_Helmholtz, electrostatic_flag, 
   !                 eps, psi, [pressure] ) 
   ! PURPOSE   
   !    Calculates Helmholtz free energy / monomer and (optionally)
   !    the pressure, given phi, mu, and omega and rho fields.
   !    Additionally, calculates electrostatic free energy terms
   !    if eps and psi are present (dielectric flag is true)
   ! SOURCE
   !--------------------------------------------------------------------
   subroutine free_energy(N, rho, omega, phi_chain, mu_chain, &
                          phi_solvent, mu_solvent, f_Helmholtz, &
                          electrostatic_flag, eps, psi, pressure)
   integer, intent(IN)    :: N              ! # of basis functions
   real(long), intent(IN) :: rho(:,:)       ! monomer vol. frac fields
   real(long), intent(IN) :: omega(:,:)     ! chemical potential field
   real(long), intent(IN) :: phi_chain(:)   ! molecule vol. frac of chain species
   real(long), intent(IN) :: mu_chain(:)    ! chemical potential of chain species
   real(long), intent(IN) :: phi_solvent(:) ! molecule vol. fraction of solvent species 
   real(long), intent(IN) :: mu_solvent(:)  ! chemical potential of solvent species
   real(long), intent(OUT):: f_Helmholtz    ! free energy/monomer
   logical,    intent(IN) :: electrostatic_flag ! True for electrostatic calculations
   real(long), intent(IN) :: eps(:)         ! dielectric field
   real(long), intent(IN) :: psi(:)         ! electric potential field
   real(long), intent(OUT), optional :: pressure 
   !***
 
   integer    :: i, alpha, beta, j   ! loop indices
   real(long) :: delpsi_sq(N)        ! Squared gradient of psi, container
   real(long) :: rho_q(N)            ! Charge density
   real(long) :: epsinv(N)           ! 1/eps

   f_Helmholtz = 0.0_long

   ! Chain mixing entropy
   do i = 1, N_chain
      if ( phi_chain(i) > 1.0E-8 ) then
         f_Helmholtz = f_Helmholtz &
                     + phi_chain(i)*( mu_chain(i) - 1.0_long )/chain_length(i)
      end if
   end do

   ! Solvent Mixing Entropy
   do i=1, N_solvent
      if ( phi_solvent(i) > 1.0E-8) then
         f_Helmholtz = f_Helmholtz &
                     + phi_solvent(i)*( mu_solvent(i) - 1.0_long)/solvent_size(i)
      end if
   end do

   ! Flory Huggins 
   do i = 1, N
      do alpha = 1, N_monomer
         do beta = alpha+1, N_monomer
            f_Helmholtz = f_Helmholtz &
                        + rho(alpha,i)*chi(alpha,beta)*rho(beta,i)
         end do
         ! Subtract field terms (Entropic)
         f_Helmholtz = f_Helmholtz - omega(alpha,i) * rho(alpha,i)
      end do
   end do

   ! Electrostatic Terms 
   if (electrostatic_flag) then

      ! Coulomb energy
      rho_q(:) = 0.0_long

      ! Calculate charge density
      j = 0
      do alpha = 1, N_monomer
         if (ABS(charge_monomer(alpha))>0) then
            j = j + 1
            rho_q = rho_q + charge_monomer(alpha)*rho(alpha,:)/solvent_size(j)
         endif
      end do

      call calc_squared_gradient(psi(:),delpsi_sq,.FALSE.)
      call field_inverse(eps,epsinv)

      do i = 1, N
         f_Helmholtz = f_Helmholtz + (rho_q(i)*psi(i) &
                     - (lp*(kuhn(1)**2)*eps(i)/(8*3.14159265359*l0))*delpsi_sq(i))
      enddo

      ! Loop over charged species
      do alpha = 1, N_monomer
         if (ionic_radii_monomer(alpha) > 0.0_long) then
            do i = 1, N
               f_Helmholtz = f_Helmholtz + (l0 / 2) &
                             * (rho(alpha,i)*epsinv(i)/ionic_radii_monomer(alpha)) &
                             / solvent_size(alpha-N_monomer+N_solvent)
            enddo
         endif 
      enddo

   endif

   
   if (present(pressure)) then
      pressure = -f_Helmholtz
      do i = 1, N_chain
         pressure = pressure + mu_chain(i)*phi_chain(i)/chain_length(i)
      end do
      do i = 1, N_solvent
         pressure = pressure + mu_solvent(i)*phi_solvent(i)/solvent_size(i)
      end do
   end if
 
   end subroutine free_energy
   !====================================================================


   !--------------------------------------------------------------------
   !****p scf_mod/free_energy_components
   ! SUBROUTINE
   !    free_energy( N, rho, omega, phi_chain, mu_chain, phi_solvent,
   !                 mu_solvent, f_Helmholtz, electrostatic_flag, 
   !                 eps, psi, [pressure] ) 
   ! PURPOSE   
   !    Calculates Helmholtz free energy / monomer and (optionally)
   !    the pressure, given phi, mu, and omega and rho fields. Splits
   !    up and returns the different contributions to free energy due
   !    to different energetic terms (ideal, flory-huggins, and
   !    electrostatics if they are turned on)
   !
   !    Only correct for single chain - needs to be adjusted in the
   !    future.
   !
   !    Normalized against homogeneous phase
   !
   ! SOURCE
   !--------------------------------------------------------------------
   subroutine free_energy_components(N, rho, omega, phi_chain, mu_chain, &
                                     phi_solvent, mu_solvent, f_components, &
                                     electrostatic_flag, eps, psi)
   integer, intent(IN)    :: N              ! # of basis functions
   real(long), intent(IN) :: rho(:,:)       ! monomer vol. frac fields
   real(long), intent(IN) :: omega(:,:)     ! chemical potential field
   real(long), intent(IN) :: phi_chain(:)   ! molecule vol. frac of chain species
   real(long), intent(IN) :: mu_chain(:)    ! chemical potential of chain species
   real(long), intent(IN) :: phi_solvent(:) ! molecule vol. fraction of solvent species 
   real(long), intent(IN) :: mu_solvent(:)  ! chemical potential of solvent species
   real(long), intent(OUT):: f_components(:)    ! free energy/monomer
   logical,    intent(IN) :: electrostatic_flag ! True for electrostatic calculations
   real(long), intent(IN) :: eps(:)         ! dielectric field
   real(long), intent(IN) :: psi(:)         ! electric potential field
   !***
 
   integer    :: i, alpha, beta, j ! loop indices
   real(long) :: delpsi_sq(N)   ! Squared gradient of psi, container
   real(long) :: rho_q(N)       ! Charge density
   real(long) :: epsinv(N)     ! 1/eps

   ! Homogenous Normalization
   real(long) :: phi_mon(N_monomer)    ! Average vol frac
   real(long) :: phi_eps_sum           ! normalize dielectric avg
   real(long) :: eps_avg               ! avg dielectric constant
   integer    :: i_ion, i_block, i_mon ! extra indices
  
   ! Components
   real(long) :: U_Flory 
   real(long) :: U_Coulomb
   real(long) :: U_Born
   real(long) :: S_chain(N_chain)
   real(long) :: S_components(N_chain*3) 
   real(long) :: S_ion(N_solvent)
   real(long) :: Q_chain(N_chain)

   ! Set things to zero
   U_Flory = 0.0_long
   U_Coulomb = 0.0_long
   U_Born = 0.0_long
   S_chain(:) = 0.0_long
   S_ion(:) = 0.0_long
   S_components = 0.0_long
   f_components = 0.0_long
   phi_mon(:) = 0.0_long

   ! Chain mixing entropy
   do i = 1, N_chain
      if ( phi_chain(i) > 1.0E-8 ) then
         S_chain(i) = S_chain(i) &
                     + phi_chain(i)*( mu_chain(i) - 1.0_long )/chain_length(i) &
                     - phi_chain(i)*(log(phi_chain(i))-1.0_long)/chain_length(i)
      end if
      do i_block = 1, N_block(i)
         i_mon = block_monomer(i_block,i)
         ! Calculate the volume fraction of each monomer species
         phi_mon(i_mon) = phi_mon(i_mon) &
                    + phi_chain(i)*block_length(i_block,i)/chain_length(i)
      end do
   end do

   ! Solvent Mixing Entropy
   do i=1, N_solvent
      if ( phi_solvent(i) > 1.0E-8) then
         S_ion(i) = S_ion(i) &
                     + phi_solvent(i)*( mu_solvent(i) - 1.0_long)/solvent_size(i) &
                     - phi_solvent(i)*( log(phi_solvent(i)) - 1.0_long)/solvent_size(i)
      end if
      i_mon = solvent_monomer(i)
      ! Update volume fractions
      phi_mon(i_mon) = phi_mon(i_mon) + phi_solvent(i)
   end do

   ! Field Term Corrections
   do i = 1,N
      do alpha = 1, N_monomer
         if (alpha <= 2) then
            S_chain(1) = S_chain(1) - omega(alpha,i) * rho(alpha,i)
         else
            S_ion(alpha-N_monomer+N_solvent) = S_ion(alpha-N_monomer+N_solvent) &
                                               - omega(alpha,i) * rho(alpha,i) 
         endif
      enddo
   enddo

   ! Flory Huggins 
   do i = 1, N
      do alpha = 1, N_monomer
         do beta = alpha+1, N_monomer
            U_Flory = U_Flory &
                        + rho(alpha,i)*chi(alpha,beta)*rho(beta,i)
         end do
      end do
   end do
  
   ! Subtract the homogeneous part
   do alpha = 1, N_monomer-1
       do beta = alpha+1, N_monomer
           U_Flory = U_Flory - phi_mon(alpha)*chi(alpha,beta)*phi_mon(beta)
       end do
   end do

   ! Electrostatic Terms 
   if (electrostatic_flag) then
      ! Calculate the homogeneous dielectric constant
      phi_eps_sum = 0.0_long
      eps_avg = 0.0_long
      do i_mon = 1, N_monomer
         if (dielectric_monomer(i_mon) > 0) then
            phi_eps_sum = phi_eps_sum + phi_mon(i_mon)
            eps_avg = eps_avg + dielectric_monomer(i_mon)*phi_mon(i_mon)
         endif
      enddo
      eps_avg = eps_avg/phi_eps_sum  

      ! Coulomb energy
      rho_q(:) = 0.0_long

      ! Calculate charge density
      j = 0
      do alpha = 1, N_monomer
         if (ABS(charge_monomer(alpha))>0) then
            j = j + 1
            rho_q = rho_q + charge_monomer(alpha)*rho(alpha,:)/solvent_size(j)
         endif
      end do

      call calc_squared_gradient(psi(:),delpsi_sq,.FALSE.)
      call field_inverse(eps,epsinv)

      do i = 1, N
         U_Coulomb = U_Coulomb + (rho_q(i)*psi(i) &
                   - (lp*(kuhn(1)**2)*eps(i)/(8*3.14159265359*l0))*delpsi_sq(i))
      enddo

      ! Born Energy

      ! Loop over charged species
      do alpha = 1, N_monomer
         if (ionic_radii_monomer(alpha) > 0.0_long) then
            do i = 1, N
               U_Born = U_Born + ( l0 / 2) &
                             * (rho(alpha,i)*epsinv(i)/ionic_radii_monomer(alpha)) &
                             / solvent_size(alpha-N_monomer+N_solvent)
            enddo
            
            ! Homogeneous normalization
            U_born = U_born - (l0/(2*eps_avg)) &
                     * (phi_mon(alpha)/ionic_radii_monomer(alpha)) &
                     / solvent_size(alpha - N_monomer+N_solvent)
         endif 
      enddo

   endif

   ! Get Chain Components
   Q_chain = 0.0_long
   do i = 1, N_chain
       Q_chain(i) = phi_chain(i)/EXP(mu_chain(i))
   enddo
   call free_energy_chain(rho,omega,phi_chain,phi_solvent,Q_chain,S_components)

 
   ! Assign to f_components
   f_components(1) = U_Flory
   f_components(2) = U_Coulomb
   f_components(3) = U_Born
   f_components(4:(4+N_chain-1)) = S_chain
   f_components((4+N_chain):(4+4*N_chain-1)) = S_components
   f_components((4+4*N_chain):(4+4*N_chain+N_solvent-1)) = S_ion

   end subroutine free_energy_components
   !====================================================================


   !-------------------------------------------------------------------
   !****p scf_mod/free_energy_chain
   ! SUBROUTINE
   !    free_energy_chain(omega, phi_chain, phi_solvent, Q, s_comp)
   ! PURPOSE   
   !    Divides the total chain entropy into conformational and
   !    translation components
   ! ARGUMENTS
   !       omega = potential fields
   !   phi_chain = volume fraction of species (chain)
   ! phi_solvent = volume fraction of species (solvent)
   !           Q = partion function of species (chain)
   !      s_comp = components of chain entropy(see below)
   ! COMMENT
   !
   !    a) Components of f_comp array:
   !       s_comp(1) = conformational energy of first block
   !       s_comp(2) = conformational energy of last  block
   !       s_comp(3) = junction translational energy (diblock)
   !
   !    b) Calculation of junction translational entropy is correct
   !       only for diblocks, for which there is only one junction
   !
   !    c) Currently only correct for linear diblock
   !
   ! SOURCE
   !----------------------------------------------------------------
   subroutine free_energy_chain(rho, omega, phi_chain, phi_solvent, Q, s_comp)
   implicit none
   real(long), intent(IN)  :: rho(:,:)        ! density field
   real(long), intent(IN)  :: omega(:,:)      ! chemical potential field
   real(long), intent(IN)  :: phi_chain(:)    ! molecule vol. frac of chain mol
   real(long), intent(IN)  :: phi_solvent(:)  ! molecule vol. frac of solvent mol
   real(long), intent(IN)  :: Q(:)            ! chain partition functions
   real(long), intent(OUT) :: s_comp(:)       ! components of chain entropy

   real(long) :: rnodes
   real(long) :: fhead       ! head block energy
   real(long) :: ftail       ! tail block energy
   real(long) :: fjct        ! junction translational entropy
   real(long) :: ftmp        ! junction translational entropy (temporary)
   integer    :: alpha, beta ! monomer indices
   integer    :: i, nh, nt   ! loop indices
   integer    :: ix,iy,iz    ! loop indices

   rnodes = dble( ngrid(1) * ngrid(2) * ngrid(3) )

   fhead = 0.0_long
   ftail = 0.0_long
   do i=1, N_chain
      nh = chains(i)%block_bgn(2)
      do iz = 0, ngrid(3)-1
      do iy = 0, ngrid(2)-1
      do ix = 0, ngrid(1)-1
         if ( chains(i)%qf(ix,iy,iz,nh) > 0.0_long .AND. &
              chains(i)%qr(ix,iy,iz,nh) > 0.0_long ) then
            fhead = fhead - chains(i)%qf(ix,iy,iz,nh)   &
                          * chains(i)%qr(ix,iy,iz,nh)   &
                          / Q(i)                        &
                          * log( chains(i)%qf(ix,iy,iz,nh) ) & 
                          * phi_chain(i) / chain_length(i)
         end if
      end do
      end do
      end do

      nt = chains(i)%block_bgn(N_block(i))
      do iz = 0, ngrid(3)-1
      do iy = 0, ngrid(2)-1
      do ix = 0, ngrid(1)-1
         if ( chains(i)%qf(ix,iy,iz,nt) > 0.0_long .AND. &
              chains(i)%qr(ix,iy,iz,nt) > 0.0_long ) then
            ftail = ftail - chains(i)%qf(ix,iy,iz,nt)   &
                          * chains(i)%qr(ix,iy,iz,nt)   &
                          / Q(i)                        &
                          * log( chains(i)%qr(ix,iy,iz,nt) ) &
                          * phi_chain(i) / chain_length(i)
         end if
      end do
      end do
      end do
   end do
   fhead = fhead / rnodes
   ftail = ftail / rnodes

   ! TODO: Not sure how to generalize for multiple chains
   do i=1, N_chain
      beta=block_monomer(1,i)
      fhead = fhead - dot_product(omega(beta,:),rho(beta,:)) !* phi_chain(i) 

      beta=block_monomer(N_block(i),i)
      ftail = ftail - dot_product(omega(beta,:),rho(beta,:)) !* phi_chain(i)
   end do

   ! --------------------------------------------
   ! The following block was used to calculate
   ! junction entropy contribution to free
   ! energy of diblocks. 
   ! Since it is not universal, we now instead
   ! calculate by subtraction, which can be
   ! interpreted by excess entropies for arbitrary
   ! molecular (linear) architecture.
   ! --------------------------------------------
   fjct = 0.0_long
   do i=1, N_chain
      nh = chains(i)%block_bgn(2)
      do iz = 0, ngrid(3)-1
      do iy = 0, ngrid(2)-1
      do ix = 0, ngrid(1)-1
         if ( chains(i)%qf(ix,iy,iz,nh) > 0.0_long .AND. &
              chains(i)%qr(ix,iy,iz,nh) > 0.0_long ) then
            ftmp  = chains(i)%qf(ix,iy,iz,nh) &
                  * chains(i)%qr(ix,iy,iz,nh) &
                  / Q(i) 
            fjct = fjct + ftmp * log( ftmp ) * phi_chain(i) / chain_length(i)
         end if
      end do
      end do
      end do
   end do
   fjct  = fjct  / rnodes
   ! --------------------------------------------
   
   !fjct = 0.0_long
   !do i=1, N_chain
   !   fjct = fjct - phi_chain(i) / chain_length(i)
   !end do
   !fjct = f_tot + fjct - enthalpy - fhead - ftail

   s_comp(1) = fhead      ! conformational energy of first block
   s_comp(2) = ftail      ! conformational energy of last  block
   s_comp(3) = fjct       ! junction translational energy (diblock)

   end subroutine free_energy_chain
   !=============================================================


   !--------------------------------------------------------------------
   !****p scf_mod/free_energy_homog
   ! FUNCTION
   !    real(long) function free_energy_homog(phi_chain,
   !    electrostatic_flag, phi_solvent)
   ! RETURN
   !    Helmholtz free energy per monomer, in energy units (kT = 1)
   !    for a homogeneous mixture of the specified composition.
   ! ARGUMENTS
   !    phi_chain(N_chain)     = molecular volume fractions of chains
   !    electrostatic_flag     = True/False depending on electrostatics
   !    phi_solvent(N_solvent) = molecular volume fractions of solvents
   ! SOURCE
   !--------------------------------------------------------------------
   real(long) function free_energy_homog(phi_chain,electrostatic_flag,phi_solvent)
   ! INPUT DEFINITIONS 
   real(long), intent(IN)           :: phi_chain(:)       ! molecule vol. frac of chain species
   logical,    intent(IN)           :: electrostatic_flag ! True for electrostatics
   real(long), intent(IN), optional :: phi_solvent(:)     ! molecule vol. fraction of solvent species


   ! Local variables
   integer      :: i_ion, i_block, i_mon, i, j ! loop indices
   real(long)   :: phi_mon(N_monomer)          ! average vol frac of each species
   real(long)   :: phi_eps_sum                 ! normalize the dielectric average
   real(long)   :: eps_avg                     ! average (homogenous) dielectric constant

   ! Initialize return value to zero
   phi_mon = 0.0
   free_energy_homog = 0.0

   ! Ideal Terms - Chain Entropy
   do i = 1, N_chain
      if ( phi_chain(i) > 1.0E-8 ) then
           free_energy_homog = free_energy_homog + & 
                            (phi_chain(i)/chain_length(i))*(log(phi_chain(i))-1)
      end if
      do i_block = 1, N_block(i)
         i_mon = block_monomer(i_block,i)
         ! Calculate the volume fraction of each monomer species
         phi_mon(i_mon) = phi_mon(i_mon) & 
                    + phi_chain(i)*block_length(i_block,i)/chain_length(i)
      end do
   end do

   ! Ideal Terms - Solvent Entropy 
   if (present(phi_solvent)) then
      do i=1, N_solvent
         if ( phi_solvent(i) > 1.0E-8 ) then
              free_energy_homog = free_energy_homog + &
                         (phi_solvent(i)/solvent_size(i))*(log(phi_solvent(i))-1)
         end if
         i_mon = solvent_monomer(i)
         ! Update volume fractions
         phi_mon(i_mon) = phi_mon(i_mon) + phi_solvent(i)
      end do
   end if

   ! Interactions

   ! Flory-Huggins Calculation
   do i = 1, N_monomer - 1
      do j = i+1, N_monomer
         free_energy_homog = free_energy_homog + chi(i,j)*phi_mon(i)*phi_mon(j)
      end do
   end do

   ! Electrostatic terms
   if (electrostatic_flag) then
      ! No coulomb terms because psi is homogenous 

      ! Calculate the homogenous dielectric constant
      phi_eps_sum = 0.0_long
      eps_avg = 0.0_long
      do i_mon = 1, N_monomer
         if (dielectric_monomer(i_mon) > 0) then
            phi_eps_sum = phi_eps_sum + phi_mon(i_mon)
            eps_avg = eps_avg + dielectric_monomer(i_mon)*phi_mon(i_mon) 
         endif
      enddo
      eps_avg = eps_avg/phi_eps_sum 

      ! Born Terms 
      do i_mon = 1, N_monomer
         if (ionic_radii_monomer(i_mon) > 0.0_long) then
            free_energy_homog = free_energy_homog + (l0/(2*eps_avg)) &
                                * (phi_mon(i_mon)/ionic_radii_monomer(i_mon)) &
                                / solvent_size(i_mon - N_monomer + N_solvent)
         endif
      enddo

   end if

   end function free_energy_homog
   !=============================================================


   !--------------------------------------------------------------------
   !****p scf_mod/calc_dielectric
   ! SUBROUTINE
   !    calc_dielectric(N, rho, eps)
   ! PURPOSE   
   !    Calculates the spatially varying dielectric constant
   !    given the density fields of monomer (rho).
   ! SOURCE
   !--------------------------------------------------------------------
   subroutine calc_dielectric(N, rho, eps)
   implicit none 
   integer, intent(IN)       :: N              ! # of basis functions
   real(long), intent(IN)    :: rho(:,:)       ! monomer vol. frac fields
   real(long), intent(OUT)   :: eps(:)         ! dielectric field

   ! Local Variables
   complex(long)   :: eps_k_grid(0:ngrid(1)/2,&
                             0:ngrid(2)-1,&
                             0:ngrid(3)-1) 
   real(long)      :: eps_r_grid(0:ngrid(1)-1,& 
                             0:ngrid(2)-1,&
                             0:ngrid(3)-1)
   complex(long)   :: rho_k_grid(0:ngrid(1)/2,&
                             0:ngrid(2)-1,&
                             0:ngrid(3)-1) 
   real(long)      :: rho_r_grid(0:ngrid(1)-1,& 
                             0:ngrid(2)-1,&
                             0:ngrid(3)-1,&
                             N_monomer)
   integer :: ix, iy, iz, i_monomer            ! Loop indices
   type(fft_plan)     :: plan                  ! fft plan
   real(long)         :: volfrac               ! volume frac of contributors
   real(long), PARAMETER :: twopi = 4.0_long*acos(0.0_long) 


   ! Initialize eps
   eps(:) = 0.0_long

   ! Select dielectric_mixing type 
   select case(trim(dielectric_mixing))
   
   case ('linear')
      do i_monomer = 1, N_monomer
         if (dielectric_monomer(i_monomer) > 0) then
            eps = eps + dielectric_monomer(i_monomer)*rho(i_monomer,:)
         endif
      end do

   case default
      write(6,*) 'Error: Invalid dielectric_mixing string'

   end select
   
   ! Convert rho and eps to real space
   call create_fft_plan(ngrid,plan)

   do i_monomer=1, N_monomer
      call basis_to_kgrid( rho(i_monomer,:), rho_k_grid )
      call ifft(plan,rho_k_grid,rho_r_grid(:,:,:,i_monomer))
   enddo

   call basis_to_kgrid(eps(:), eps_k_grid)
   call ifft(plan,eps_k_grid,eps_r_grid)

   ! Loop over the rgrid and adjust
   do iz=0, ngrid(3)-1
     do iy=0, ngrid(2)-1
         do ix=0, ngrid(1)-1
            ! Normalize by the volume fraction of species that
            ! contribute to the dielectric constant
            volfrac = 0.0_long;
            do i_monomer = 1, N_monomer
               if (dielectric_monomer(i_monomer) > 0) then
                  volfrac = volfrac + rho_r_grid(ix,iy,iz,i_monomer)
               endif
            enddo
            eps_r_grid(ix,iy,iz) = eps_r_grid(ix,iy,iz)/volfrac
         enddo
      enddo
   enddo

   ! Convert eps back to basis representation 
   eps_k_grid(:,:,:) = (0.0_long, 0.0_long)
   call fft(plan,eps_r_grid,eps_k_grid)
   eps_k_grid = eps_k_grid/dble(plan%n(1)*plan%n(2)*plan%n(3))
   call kgrid_to_basis(eps_k_grid, eps)

   end subroutine calc_dielectric
   !====================================================================


   !--------------------------------------------------------------------
   !****p scf_mod/calc_electric_potential
   ! SUBROUTINE
   !    calc_electric_potential(N, rho, eps, psi)
   ! PURPOSE   
   !    Calculates the electric potential field given the
   !    monomer density fields and the dielectric field.
   !    Output psi is (e/eps_0)*psi 
   ! SOURCE
   !    PoissonGreenFunc.opp
   !    Jian Qin, 2010
   !--------------------------------------------------------------------
   subroutine calc_electric_potential(N, rho, eps, psi, elec_init_flag)
   
   ! Modules
   use basis_mod
   use unit_cell_mod
   use group_mod,       only : operator(.dot.)
   !use field_io_mod     ! TEST CODE ONLY REMOVE LATER

   ! Input Parameters
   integer, intent(IN)       :: N               ! # of basis functions
   real(long), intent(IN)    :: rho(:,:)        ! monomer vol. frac fields
   real(long), intent(IN)    :: eps(:)          ! dielectric field
   real(long), intent(OUT)   :: psi(:)          ! electric potential field
   logical, intent(IN)       :: elec_init_flag  ! TRUE if elec_init_flag needs init

   ! Internal/subroutine variables
   integer    :: i, j, k            ! loop indices
   integer    :: a1, a2, a3         ! loop indices for q_alpha
   integer    :: b1, b2, b3         ! loop indices for q_beta
   integer    :: alpha, beta        ! more loop indices
   integer    :: info               ! Allocation error check
   integer    :: abz(3), bbz(3)     ! First Brilloun Zone indices for a and b
   integer    :: afft(3),bfft(3)    ! (0:N(1)-1,...), natural order of fft
   real(long) :: avec(3),bvec(3)    ! reciprocal basis vectors for alpha, beta
   integer    :: aminusb(3),dvec(3) ! avec-bvec, dvec is converted to real grid
   real(long) :: qdotprod           ! stores q_alpha * q_beta
   complex(long) :: epskval         ! stores temporary value for dielectric term
   complex(long) :: phi_qkval       ! stores temporary value for phiq
   integer       :: counts(N-1)     ! stores row multiplicity for conv to basis

   ! Arrays
   real(long), allocatable   :: phi_q(:)
   complex(long),allocatable :: eps_kgrid(:,:,:)
   complex(long),allocatable :: A_kspace(:,:)
   complex(long),allocatable :: A_basis(:,:)
   real(long),allocatable :: A_basis_inv(:,:)

   ! Auxillary vars for matrix inversion (LAPACK)
   integer :: lpdim, lapackinfo
   complex(long),allocatable :: work(:)
   integer,allocatable       :: ipiv(:)

   ! Allocate and initialize phi_q
   allocate(phi_q(N), stat = info )
   if (info.ne.0) stop 'Error allocating charge density'
   phi_q(:) = 0.0_long

   ! Calculate charge density
   j = 0
   do k = 1, N_monomer
      if (ABS(charge_monomer(k)) > 0) then
         j = j + 1
         phi_q = phi_q + charge_monomer(k)*rho(k,:)/solvent_size(j)
      endif
   end do
 
   ! Allocate some Arrays
   allocate( eps_kgrid(0:ngrid(1)/2,0:ngrid(2)-1,0:ngrid(3)-1), stat=info)
   if (info.ne.0) stop 'Error allocating dielectric k-grid'
   allocate(A_kspace(N_wave-1,N_wave-1), stat = info)
   if (info.ne.0) stop 'Error allocating A_kspace'
 
   ! Converte eps(basis) to kgrid (use full grid)
   call basis_to_kgrid(eps,eps_kgrid)

   ! Set Array 
   A_kspace(:,:) = (0.0_long,0.0_long)

   ! Initialize/Recalculate Gradient Vectors if needed
   call elec_startup(elec_init_flag)

   ! Calculate the Green's Function for the Poisson Eq in k-space
   ! Loop over all possible q_alpha 
   ! Skip unwanted zero frequency mode (constant)
   do beta = 2, N_wave
     do alpha = 2, N_wave
         dvec = aminusbvecs(alpha,beta,:)
         if ((dvec(1).GT.(ngrid(1)/2)) ) then
            dvec(1) = -dvec(1)
            dvec(2) = -dvec(2)
            dvec(3) = -dvec(3) 
            dvec = G_to_fft(dvec)
            epskval = CONJG(eps_kgrid(dvec(1),dvec(2),dvec(3)))
         else
            epskval = eps_kgrid(dvec(1),dvec(2),dvec(3))
         endif
         
         ! Note that A_kspace is shifted one index left
         ! Drops the zero-frequency mode
         A_kspace(alpha-1,beta-1) = - qadotqb(alpha,beta) * epskval
      enddo
   enddo
 
   if(allocated(eps_kgrid)) deallocate(eps_kgrid)
   allocate(A_basis(N-1,N-1), stat = info)
   if (info.ne.0) stop 'Error allocating A_basis'
   allocate(A_basis_inv(N-1,N-1), stat = info)
   if (info.ne.0) stop 'Error allocating A_basis_inv'

   A_basis(:,:) = (0.0_long,0.0_long)
   counts(:) = 0

   ! Basis Representation 
   do alpha = 2, N_wave
      i = star_of_wave(alpha) - 1
      counts(i) = counts(i) + 1

      do beta = 2, N_wave
         j = star_of_wave(beta) - 1
         A_basis(j,i) =  A_basis(j,i) & 
         + coeff(alpha)*A_kspace(beta-1,alpha-1)/coeff(beta)
      end do
   end do

   if(allocated(A_kspace)) deallocate(A_kspace)
   
   ! Normalize by Rows
   do i = 1, N-1
      A_basis(i,:) = A_basis(i,:)/real(counts(i))
   enddo

   ! Free up work and ipiv (from kspace only)
   if(allocated(work)) deallocate(work)
   if(allocated(ipiv)) deallocate(ipiv)

   ! Invert the green's function in basis representation
   allocate(work(N-1), stat = info)
   if (info.ne.0) stop 'Error allocating work'
   allocate(ipiv(N-1), stat = info)
   if (info.ne.0) stop 'Error allocating ipiv'

   A_basis_inv = REALPART(A_basis)
   lpdim = size(A_basis_inv,1)
   if(allocated(A_basis)) deallocate(A_basis)

   call dgetrf(lpdim,lpdim,A_basis_inv,lpdim,ipiv,lapackinfo)
   if (lapackinfo.ne.0) stop 'Green Matrix is singular'
   call dgetri(lpdim,A_basis_inv,lpdim,ipiv,work,lpdim,lapackinfo)
   if (lapackinfo.ne.0) stop 'Green Function Inversion Failed'

   psi(:) = 0.0_long
   ! Calculate the electrostatic potential
   do i = 2, N
      do j = 2, N 
         psi(i) = psi(i) + A_basis_inv(i-1,j-1) &
                 *(-4*3.14159265359*l0*phi_q(j)) / (lp*kuhn(1)**2) 
      end do
   end do

   ! Free up memory 
   if(allocated(phi_q)) deallocate(phi_q)
   if(allocated(A_basis_inv)) deallocate(A_basis_inv)
   if(allocated(work)) deallocate(work)
   if(allocated(ipiv)) deallocate(ipiv)

   end subroutine calc_electric_potential 
   !====================================================================


   !-------------------------------------------------------------
   !****p scf_mod/calc_squared_gradient
   ! SUBROUTINE
   !    calc_squared_gradient(A, Asqgrad)
   ! PURPOSE
   !    Calculates |grad(A)|^2 for input field A in basis repr.
   ! ARGUMENTS
   !    A(N_Star)         = input field, basis representation
   !    Asqgrad(N_star)   = sq-grad of input field, basis repr.
   ! SOURCE
   !-------------------------------------------------------------
   subroutine calc_squared_gradient(A, Asqgrad,elec_init_flag)
   use basis_mod
   use unit_cell_mod
   use group_mod,       only : operator(.dot.)

   ! Subroutine Arguments
   real(long), intent(IN)  :: A(:)
   real(long), intent(OUT) :: Asqgrad(:)
   logical, intent(IN)     :: elec_init_flag

   ! k-space representations
   complex(long)::       A_kgrid(0:ngrid(1)-1,& 
                                 0:ngrid(2)-1,&
                                 0:ngrid(3)-1)
   complex(long):: Asqgrad_kgrid(0:ngrid(1)-1,& 
                                 0:ngrid(2)-1,&
                                 0:ngrid(3)-1)

   ! Internal variables
   integer       :: i, j                     ! loop index
   real(long)    :: rnodes                   ! number of grid points
   integer       :: alpha ,beta              ! alpha (a1, a2, a3), beta (b1, b2, b3)
   integer       :: abz(3), bbz(3)           ! First Brilloun Zone indices for a and b
   integer       :: afft(3) ,bfft(3)         ! (0:N(1)-1,...), natural order of fft
   real(long)    :: avec(3),bvec(3)          ! reciprocal basis vectors for alpha, beta
   integer       :: aplusb(3), abfft(3)      ! avec+bvec
   complex(long) :: Akval1, Akval2           ! stores temporary value for psi term

   ! Re-initalize if needed
   call elec_startup(elec_init_flag)

   ! Convert input to kgrid (use full grid)
   A_kgrid = 0.0_long
   call basis_to_kgrid(A,A_kgrid)
   rnodes=dble( plan%n(1) * plan%n(2) * plan%n(3) )

   ! Set Array 
   Asqgrad_kgrid(:,:,:) = (0.0_long,0.0_long)

   ! Loop over all alpha and beta
   do alpha = 1, N_wave
      abz = wave(:,alpha)
      afft = G_to_fft(abz)

      do beta = 1, N_wave
         bbz = wave(:,beta)
         bfft = G_to_fft(bbz)

         abfft = aminusbvecs(beta,alpha,:)
         ! If abfft is outside of the FBZ then we can get
         ! the value of eps by taking the complex conj
         ! of it's mirror term (epsilon field is real)
         if ((abfft(1).GT.(ngrid(1)/2)) ) then
            abfft(1) = -abfft(1)
            abfft(2) = -abfft(2)
            abfft(3) = -abfft(3) 
            abfft = G_to_fft(abfft)
            Akval1 = A_kgrid(abfft(1),abfft(2),abfft(3))
         else
            Akval1 = CONJG(A_kgrid(abfft(1),abfft(2),abfft(3)))
         endif

         ! If bfft is outside of the FBZ then we can get
         ! the value of eps by taking the complex conj
         ! of it's mirror term (epsilon field is real)
         if ((bfft(1).GT.(ngrid(1)/2)) ) then
            bfft(1) = -bfft(1)
            bfft(2) = -bfft(2)
            bfft(3) = -bfft(3) 
            bfft = G_to_fft(bfft)
            Akval2 = CONJG(A_kgrid(bfft(1),bfft(2),bfft(3)))
         else
            Akval2 = A_kgrid(bfft(1),bfft(2),bfft(3))
         endif

         ! Update Value
         Asqgrad_kgrid(afft(1),afft(2),afft(3)) = &
            Asqgrad_kgrid(afft(1),afft(2),afft(3)) &
            + qaminbdotqa(beta,alpha)*Akval1*Akval2
      enddo
   enddo

   call kgrid_to_basis(Asqgrad_kgrid,Asqgrad)
 
   end subroutine calc_squared_gradient
   !================================================================

   !--------------------------------------------------------------------
   !****p scf_mod/field_inverse
   ! SUBROUTINE
   !    real function field_inverse(A)
   ! PURPOSE
   !    Given input field A in basis representation, computes the
   !    inverse of A in real space and returns that in basis rep. 
   ! ARGUMENTS
   !    A  =  input field, basis representation
   ! SOURCE
   !--------------------------------------------------------------------
   subroutine field_inverse(A,Ainv)

   real(long), intent(IN)  :: A(:)
   real(long), intent(OUT) :: Ainv(:) 

   ! k-space, real-space representations
   complex(long)::       A_kgrid(0:ngrid(1)/2,& 
                                 0:ngrid(2)-1,&
                                 0:ngrid(3)-1)
   complex(long)::    Ainv_kgrid(0:ngrid(1)/2,& 
                                 0:ngrid(2)-1,&
                                 0:ngrid(3)-1)
   real(long)   ::       A_rgrid(0:ngrid(1)-1,& 
                                 0:ngrid(2)-1,&
                                 0:ngrid(3)-1)
   real(long)   ::    Ainv_rgrid(0:ngrid(1)-1,& 
                                 0:ngrid(2)-1,&
                                 0:ngrid(3)-1)

   ! Initialize
   Ainv = 0.0_long

   ! Convert A into real space
   call basis_to_kgrid(A,A_kgrid)
   call ifft(plan,A_kgrid,A_rgrid)

   ! Invert
   Ainv_rgrid = 1/A_rgrid

   ! Conver Ainv to basis representation
   call fft(plan,Ainv_rgrid, Ainv_kgrid)
   Ainv_kgrid = Ainv_kgrid/dble(plan%n(1)*plan%n(2)*plan%n(3))
   call kgrid_to_basis(Ainv_kgrid, Ainv)

   end subroutine field_inverse
   !=============================================================


   !--------------------------------------------------------------------
   !****p scf_mod/field_product
   ! SUBROUTINE
   !    subroutine field_product(A,B)
   ! PURPOSE
   !    Given input fields A and B in basis representation, computes the
   !    product of A and B in real space and returns that in basis rep.
   ! ARGUMENTS
   !    A  =  input field, basis representation
   !    B  =  input field, basis representation
   ! SOURCE
   !--------------------------------------------------------------------
   subroutine field_product(A,B,AB)

   real(long), intent(IN)  :: A(:)
   real(long), intent(IN)  :: B(:)
   real(long), intent(OUT) :: AB(:)

   ! k-space, real-space representations
   complex(long)::       A_kgrid(0:ngrid(1)/2,& 
                                 0:ngrid(2)-1,&
                                 0:ngrid(3)-1)
   complex(long)::       B_kgrid(0:ngrid(1)/2,& 
                                 0:ngrid(2)-1,&
                                 0:ngrid(3)-1)
   complex(long)::      AB_kgrid(0:ngrid(1)/2,& 
                                 0:ngrid(2)-1,&
                                 0:ngrid(3)-1)
   real(long)   ::       A_rgrid(0:ngrid(1)-1,& 
                                 0:ngrid(2)-1,&
                                 0:ngrid(3)-1)
   real(long)   ::       B_rgrid(0:ngrid(1)-1,& 
                                 0:ngrid(2)-1,&
                                 0:ngrid(3)-1)
   real(long)   ::      AB_rgrid(0:ngrid(1)-1,& 
                                 0:ngrid(2)-1,&
                                 0:ngrid(3)-1)

   ! Initialize
   AB = 0.0_long

   ! Convert A and B into real space
   call basis_to_kgrid(A,A_kgrid)
   call basis_to_kgrid(B,B_kgrid)
   call ifft(plan,A_kgrid,A_rgrid)
   call ifft(plan,B_kgrid,B_rgrid)

   ! Multiply
   AB_rgrid = A_rgrid*B_rgrid 

   ! Conver AB to basis representation
   call fft(plan,AB_rgrid, AB_kgrid)
   AB_kgrid = AB_kgrid/dble(plan%n(1)*plan%n(2)*plan%n(3))
   call kgrid_to_basis(AB_kgrid,AB)
   
   end subroutine field_product
   !=============================================================

   !--------------------------------------------------------------------
   !****p scf_mod/check_poisson_hermitian
   ! FUNCTION
   !    logical function check_poisson_hermitian(A)
   ! RETURN
   !    k
   ! ARGUMENTS
   !    A  =  array of complex numbers to check
   ! SOURCE
   !--------------------------------------------------------------------
   logical function check_poisson_hermitian(A)

   complex(long), intent(IN) :: A(:,:)
   integer :: i, j
   logical :: checkvalue

   checkvalue = .TRUE.

   do i = 1, size(A,1)
      do j = 1, size(A,2)
         if (ABS(A(i,j)-CONJG(A(j,i))) > 1.0E-9_long) then
            checkvalue = .FALSE.
         endif
      enddo
   enddo

   check_poisson_hermitian = checkvalue

   end function check_poisson_hermitian
   !=============================================================
   
   
   !--------------------------------------------------------------------
   !****p scf_mod/check_poisson_duality
   ! FUNCTION
   !    logical function check_poisson_duality(A)
   ! RETURN
   !    k
   ! ARGUMENTS
   !    A  =  array of complex numbers to check
   ! SOURCE
   !--------------------------------------------------------------------
   logical function check_poisson_duality(A)
   use basis_mod,   only : wave, which_wave
   
   complex(long), intent(IN) :: A(:,:)
   integer :: i, j, k, x, y
   integer :: xvec(3), yvec(3), xt(3), yt(3)
   logical :: checkvalue

   checkvalue = .TRUE.

   do i = 1, size(A,1)
      do j = 1, size(A,2)
         xvec = wave(:,i+1)
         yvec = wave(:,j+1)

         do k = 1, 3
            if (k>dim) then
               xvec(k) = 0
               yvec(k) = 0
            else 
               xvec(k) = -xvec(k)
               yvec(k) = -yvec(k)
            endif
         enddo

         xt = -xvec;
         yt = -yvec;
         xvec = G_to_bz(xvec)
         yvec = G_to_bz(yvec)

         x = which_wave(xvec(1),xvec(2),xvec(3)) - 1
         y = which_wave(yvec(1),yvec(2),yvec(3)) - 1


         if ( ABS(A(x,y)-CONJG(A(i,j))) > 1.0E-9_long ) then
            checkvalue = .FALSE.
            print*, '================='
            print*, '(', x, ',', y, ')'
            print*, xvec
            print*, yvec
            print*, '(', i, ',', j, ')'
            print*, xt
            print*, yt
            print*, (ABS(A(x,y)-CONJG(A(i,j))))
         endif
      enddo
   enddo

   check_poisson_duality = checkvalue

   end function check_poisson_duality
   !=============================================================


end module scf_mod
