! fortran_dialect=elf
!-----------------------------------------------------------------------
! 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/iterate_mod
! MODULE
!   iterate_mod
! PURPOSE
!   Iteration of SCF equations
! ALGORITHM
!   Newton/Broyden algorithm that uses an approximate calculation 
!   for the initial Jacobian and Broyden updates to the inverse 
!   Jacobian. The approximate initial Jacobian is calculatd from 
!   a corresponding approximate linear response matrix, which is 
!   calculated in response_pd_mod. See the docblock for N_cut for 
!   a discussion of approximate response function. 
!
!   Subroutine iterate_NR can either:
!     i)  Iterate omega in a fixed unit cell     (domain = F)
!     ii) Iterate both omega and unit_cell_param (domain = T)
!   depending upon value of logical module variable domain
! AUTHOR
!   Chris Tyler - multiblock melt version
!   Amit Ranjan - multi-component blend version (2003)
!   David Morse - modifications (1/2004)
!   Jian Qin    - Broyden and backtracking (2006-2007)
!   Kevin J Hou - Added electrostatic interactions (2016)
! SOURCE
!-----------------------------------------------------------------------
module iterate_mod
   use const_mod
   use io_mod
   use response_pd_mod
   implicit none

   private

   ! public procedures
   public :: input_iterate_param   ! input iteration parameters 
   public :: output_iterate_param  ! output iteration variables
   public :: iterate_NR_startup    ! allocates arrays needed by iterate_NR
   public :: iterate_AM_startup    ! allocates arrays needed by iterate_AM
   public :: iterate_NR            ! main Newton-Raphson iteration routine
   public :: iterate_AM            ! main Anderson-Mixing iteration routine

   ! public variables
   public :: max_itr    ! maximum allowable number of iterations
   public :: error_max  ! error tolerance: stop when error < error_max
   public :: domain     ! If true, adjust unit_cell_dimensions
   public :: itr_algo   ! Iteration algorithm 
   public :: N_cut      ! dimension of cutoff response and Jacobian for NR scheme
   public :: N_hist     ! Number of previous histories to use in Anderson-Mixing Scheme
   !***

   !--------------------------------------------------------------------
   !****v iterate_mod/max_itr 
   ! VARIABLE
   !   integer max_itr - maximum allowed number of iterations
   !*** ----------------------------------------------------------------
   !****v iterate_mod/error_max 
   ! VARIABLE
   !   real(long) error_max - maximum error in converged solution
   !                        - Stop when error < error_max
   !*** ----------------------------------------------------------------
   !****v iterate_mod/domain 
   ! VARIABLE
   !   logical domain - If .true.,  iterate variable unit_cell shape
   !                    If .false., iterate omega in fixed unit cell
   !*** ----------------------------------------------------------------
   !****v iterate_mod/iter_algo
   ! VARIABLE
   !   character(10) iter_algo - Iteration algorithm to be used
   !                             For now, must equal 'NR' (Newton-Raphson)
   ! PURPOSE
   !   This variable will eventually allow the user to select from more
   !   than one iteration algorithm. Thus far, however, the only algorithm
   !   that has been implemented is a quasi-Newton / Broyden algorithm, 
   !   for which the value of iter_algo should be 'NR'. Implementation of
   !   a more robust relaxation algorithm is planned. 
   !*** ----------------------------------------------------------------
   !****v iterate_mod/N_cut 
   ! VARIABLE
   !   integer N_cut  - dimension of cutoff response and Jacobian
   ! PURPOSE
   !   The approximate linear response matrix used to construct the
   !   initial Jacobian is obtained by numerically calculating the
   !   response to the first N_cut basis functions to obtain an
   !   N_cut x N_cut submatrix within the subspace spanned by these
   !   functions, and using the analytically calculated response of
   !   a homogenous gas to approximate the response to the remaining
   !   basis functions.
   !*** ----------------------------------------------------------------

   ! public variable declarations
   integer       ::  max_itr       ! maximum # of iterations
   real(long)    ::  error_max     ! error tolerance 
   logical       ::  domain        ! true if domain iteration
   character(10) ::  itr_algo      ! Iteraction algorithm string. For
                                   ! now, must be 'NR' = Newton-Raphson
   integer       ::  N_cut         ! dimension of cutoff response matrix
   integer       ::  N_hist        ! Number of previous histories in AM
   real(long), PARAMETER :: PI = 3.14159265359_long ! Pi

   ! Private variable declarations
   integer                 :: N_residual      ! number of residuals
   integer                 :: N_corner        ! cut-off dimension of Jacobian
   integer, allocatable    :: ipvt(:)         ! Pivots for LU decomp 
   integer                 :: lwork           ! workspace dimension of inverse
   real(long), allocatable :: work(:)         ! workspace for matrix inverse
   real(long), allocatable :: residual(:)     ! residual array (N_residual)
   real(long), allocatable :: delta(:)        ! change in variables (N_residual)
   real(long), allocatable :: Jacobian(:,:)   ! Jacobian matrix
   real(long), allocatable :: J_corner(:,:)   ! Jacobian matrix
   logical                 :: Jacobian_empty  ! if true, Jacobian is empty
   logical                 :: full_inversion  ! if true, full Jacobian is inverted

   ! Varaibles for Anderson Mixing scheme
   real(long),allocatable  :: Umn(:,:)
   real(long),allocatable  :: dev(:,:,:)
   real(long),allocatable  :: omega_hist(:,:,:)
    

contains

   !--------------------------------------------------------------------
   !****p iterate_mod/input_iterate_param
   ! SUBROUTINE
   !   input_iterate_param
   ! PURPOSE
   !   Read parameters max_itr, error_max, domain, itr_algo, N_cut
   ! SOURCE
   !--------------------------------------------------------------------
   subroutine input_iterate_param
   !***
   call input(max_itr,'max_itr')      ! max # of iterations
   call input(error_max,'error_max')  ! error tolerance
   call input(domain,'domain')        ! T -> domain iteration
   call input(itr_algo,'itr_algo')    ! Iteration algorithm

   if (itr_algo=='NR')then
      call input(N_cut,'N_cut')
   elseif(itr_algo=='AM')then
      call input(N_hist,'N_hist')
   endif

   end subroutine input_iterate_param
   !====================================================================


   !--------------------------------------------------------------------
   !****p iterate_mod/output_iterate_param
   ! SUBROUTINE
   !   output_iterate_param
   ! PURPOSE
   !   Read parameters max_itr, error_max, domain, N_cut
   ! SOURCE
   !--------------------------------------------------------------------
   subroutine output_iterate_param
   !***
   call output(max_itr,'max_itr')      ! max # of iterations
   call output(error_max,'error_max')  ! error tolerance
   call output(domain,'domain')        ! T -> domain iteration
   call output(itr_algo,'itr_algo')    ! Iteration algorithm

   if (itr_algo=='NR')then
      call output(N_cut,'N_cut')
   elseif(itr_algo=='AM')then
      call output(N_hist,'N_hist')
   endif

   end subroutine output_iterate_param
   !====================================================================


   !--------------------------------------------------------------------
   !****p iterate_mod/iterate_NR_startup
   ! SUBROUTINE
   !    iterate_NR_startup(N)
   ! PURPOSE
   !   1) Set integer module variable N_residual = # of residuals
   !   2) Allocate arrays needed by iterate_NR
   ! ARGUMENTS
   !   integer N      = # of basis functions
   ! SOURCE
   !---------------------------------------------------------------------
   subroutine iterate_NR_startup(N)
   use unit_cell_mod, only : N_cell_param
   use chemistry_mod, only : N_monomer, ensemble
   integer, intent(IN)  :: N       ! # of basis functions
   !***

   integer              :: info    ! message variable for dgetri

   if (N_cut > N) then
       N_cut = N
   endif
   if (domain) then
      N_residual  = N_monomer*( N - 1 + ensemble ) + N_cell_param
      if( N_cut > 0 ) then
         N_corner = N_monomer*(N_cut- 1 + ensemble ) + N_cell_param
      else
         N_corner = N_cell_param
      endif
   else
      N_residual  = N_monomer*( N - 1 + ensemble )
      if( N_cut > 0 ) then
         N_corner = N_monomer*(N_cut- 1 + ensemble )
      else
         N_corner = 0
      endif
   endif

   if (allocated(ipvt)) deallocate(ipvt)
   if (allocated(work)) deallocate(work)
   if (allocated(residual)) deallocate(residual)
   if (allocated(delta)) deallocate(delta)
   if (allocated(Jacobian)) deallocate(Jacobian)
   if (allocated(J_corner)) deallocate(J_corner)

   allocate(ipvt(N_residual))
   allocate(work(N_residual))
   allocate(residual(N_residual))
   allocate(delta(N_residual))
   allocate(Jacobian(N_residual,N_residual))
   allocate(J_corner(N_corner,N_corner))

   call init_response_pd(N_monomer,N)

   ! estimate the opitimal workspace dimension
   lwork = -1
   call dgetri(N_residual,Jacobian,N_residual,ipvt,work,lwork,info)
   lwork = work(1)
   if (allocated(work)) deallocate(work)
   allocate(work(lwork))

   Jacobian_empty = .true.   ! reset to true after 1st calculation

   ! TESTCODE!
   full_inversion = .true.  ! not invert full Jacobian, see invert_Jacobian

   end subroutine iterate_NR_startup
   !======================================================================


   !---------------------------------------------------------------------
   !****p iterate_mod/iterate_NR
   ! SUBROUTINE
   !    subroutine iterate_NR
   ! PURPOSE
   !    Newton-Raphson iteration of the SCFT, using numerically
   !    calculated Jacobian
   ! COMMENT
   !   1) Algorithm assumes that density_startup, iterate_NR_startup,
   !      and make_unit_cell have been called prior to iterate_NR.
   !
   !   2) Routine will simultaneously iterates omega field and unit 
   !      cell parameters if domain = .true., or iterate the omega
   !      field for a fixed unit cell if domain = .false.
   ! SOURCE
   !--------------------------------------------------------------------
   subroutine iterate_NR(  &
              N,                 &!  # of basis functions
              omega,             &!  chemical potential field (IN/OUT)
              itr,               &!  actual number of interations
              converge,          &!  = .true. if converged
              error,             &!  final error = max(residuals)
              rho,               &!  monomer density field
              f_Helmholtz,       &!  Helmholtz free energy per monomer / kT
              f_components,      &!  Helmholtz free energy components
              pressure,          &!  pressure * monomer volume / kT
              stress,            &!  d(free energy)/d(cell parameters)
              eps,               &! spatially varying dielectric field
              psi,               &! electric potential 
              electrostatic_flag &! true for e. calculations
              )
   !--------------------------------------------------------------------

   use io_mod, only        : output
   use unit_cell_mod, only : N_cell_param, cell_param, &
                             make_unit_cell, G_basis, dGG_basis
   use basis_mod, only     : make_dGsq
   use chemistry_mod, only : N_monomer, ensemble, mu_chain, mu_solvent, &
                             phi_solvent, phi_chain, N_chain, N_solvent, &
                             dielectric_monomer, ionic_radii_monomer, &
                             charge_monomer, dielectric_mixing
   use scf_mod, only       : density, scf_stress, set_omega_uniform, &
                             mu_phi_chain, mu_phi_solvent, free_energy, &
                             calc_dielectric, calc_electric_potential, &
                             free_energy_components, chains ! remove chains
   use grid_mod, only      : make_ksq

   ! TESTCODE
   use field_io_mod

   ! Arguments
   integer, intent(IN)         :: N        
   real(long), intent(INOUT)   :: omega(:,:)  ! omega(N_monomer,N)
   integer, intent(OUT)        :: itr         
   logical, intent(OUT)        :: converge
   real(long), intent(OUT)     :: error
   real(long), intent(OUT)     :: rho(:,:)    ! rho(N_monomer,N)
   real(long), intent(OUT)     :: f_Helmholtz
   real(long), intent(OUT)     :: f_components(:)
   real(long), intent(OUT)     :: pressure
   real(long), intent(OUT)     :: stress(:)   ! stress(N_cell_param)
   real(long), intent(OUT)     :: eps(:)
   real(long), intent(OUT)     :: psi(:)
   logical, intent(IN)         :: electrostatic_flag 

   !***

   ! parameter
   real(long), parameter :: large = 100.0e3

   ! local Variables
   integer    :: i, j, k, alpha, M, info
   integer    :: error_index
   real(long) :: dGsq(N,N_cell_param), q(N_chain), q_solvent(N_solvent)
   real(long) :: old_error
   real(long) :: delta_norm, u(N_residual), v(N_residual), uv

   ! parameter and variables for backtracking line search
   ! backtracking algorithm is as described in Numerical Recipes
   real(long), parameter :: STPMX=100.0_long
   real(long)  :: ow(N_monomer,N)    ! old omega fields
   real(long)  :: ocell(N_cell_param)! old cell param
   real(long)  :: fold, fnew         ! f=x^2/2, to be optimized
   real(long)  :: p(N_residual)      ! update step direction
   real(long)  :: stpmax             ! maximum update step
   real(long)  :: slope              ! line search direction
   real(long)  :: r_old(N_residual)  ! old residual
   real(long)  :: lam, vsum

   real(long)  :: error1, error2
   real(long),parameter :: stress_rescale=1.0d+2

   error_index = N_residual
   if(domain) error_index = error_index - N_cell_param

   M = N - 1 + ensemble 
  
   !  Initialization
   call density(N, omega, rho, q, q_solvent)
   call mu_phi_chain(mu_chain, phi_chain, q)
   if (N_solvent > 0) then
      call mu_phi_solvent(mu_solvent, phi_solvent, q_solvent)
   endif
   call make_correlation(N)
   do i = 1, N_cell_param
      call make_dGsq(dGsq(:,i), dGG_basis(:,:,i))
   enddo
 
   ! Initialize spatially varying dielectric and electric potential 
   ! if electrostatic_flag is set to true.
   if (electrostatic_flag) then
      call calc_dielectric(N,rho,eps)
      call calc_electric_potential(N,rho,eps,psi,.TRUE.)
   endif

   ! TESTCODE
   !print*, "Testing Psi"
   !open(file='out/psi',unit=22,status='replace')
   !call output_single_field(psi,22,'-1')
   !close(22)
   !open(file='out/eps',unit=22,status='replace')
   !call output_single_field(eps,22,'-1')
   !close(22)
   !open(file='out/rho',unit=22,status='replace')
   !call output_field(rho,22,'-1')
   !close(22)
   !stop

   ! Set homogenous coefficients
   if (ensemble == 0) call set_omega_uniform(omega,electrostatic_flag)

   ! calculate the maximum step size
   vsum = 0.0_long
   do alpha = 1, N_monomer
      vsum = vsum + dot_product(omega(alpha,2-ensemble:), &!
                                omega(alpha,2-ensemble:))
   end do
   if(domain) vsum = vsum + dot_product(cell_param(1:N_cell_param),&!
                                        cell_param(1:N_cell_param))
   stpmax = STPMX * max( sqrt(vsum), dble(N_residual) )

   itr = 0

   ! Calculate density/stress and residual/error
   if (domain) stress = scf_stress(N, N_cell_param, dGsq)
   call make_residual(N, domain, rho, stress, omega, residual, electrostatic_flag, psi, eps)

   !error    = maxval(abs(residual(1:error_index)))
   !error     = maxval(abs(residual))
   error1 = maxval(abs(residual(1:error_index)))
   error2 = maxval(abs(residual(error_index+1:N_residual)))
   error  = max(error1, error2*stress_rescale)

   ! Calculate thermodynamic properties
   call free_energy(N,rho,omega,phi_chain,mu_chain,phi_solvent,mu_solvent, &
                    f_Helmholtz,electrostatic_flag,eps,psi,pressure)
   

   iterative_loop : do
    
      write(6,*)
      write(6,"('Iteration ',i3)") itr

      ! TESTCODE
      !print*, '-----residual is:-----'
      !print*, residual
      !print*, '----------------------'

      ! Output to log file
      call output(error,                     'error       =',o=6,f='L')
      if (domain) then
         call output(stress,N_cell_param,    'stress      =',o=6,f='L')
         call output(cell_param,N_cell_param,'cell_param  =',o=6,f='L')
      endif
      call output(f_Helmholtz,               'f_Helmholtz =',o=6,f='L')
      if (ensemble == 1) then
         call output(pressure,               'pressure    =',o=6,f='L')
      endif 

      !if (itr > 0) then
      !   print*, 'EARLY STOP TEST CODE'
      !   converge = .false.
      !   exit iterative_loop
      !endif

      ! Check for convergence/exit conditions
      if ( error < error_max ) then
         converge = .true.
         ! Get components of FE
         call free_energy_components(N,rho,omega,phi_chain,mu_chain,phi_solvent,mu_solvent, &
                          f_components,electrostatic_flag,eps,psi)
         exit iterative_loop
      else if ( itr > max_itr ) then
         converge = .false.
         exit iterative_loop
      else if ( error > 100.0*large ) then
         converge = .false.
         print*, 'ERROR TOO LARGE'
         print*, 'error is: ', error
         print*, 'max_error is: ', (100.0*large)
         exit iterative_loop
      endif


      ! Calculate the inverse of Initial Jacobian/Broyden
      if ( Jacobian_empty ) then

         
         write(6,"('Initializing Jacobian ...')")
         call Jacobian_response(N, N_cut, domain, omega, Jacobian, rho, &
                                electrostatic_flag, eps, psi)
         Jacobian_empty = .false. 

         !print*, '======================    JACOBIAN    ================'
         !print*, 'Maxval: ', Maxval(Abs(Jacobian))
         !print*, 'Minval: ', Minval(Abs(Jacobian))
         !print*, '   '

         !open (unit=41,file='Jacobian.txt', ACTION = "write", STATUS = "replace")
         !do i = 1,N_residual
         !   write(41,*) real( Jacobian(i,:))
         !   write(41,*) "====================="
         !end do
         !close(41)

         !print*, '================== CORNER ==========================='
         !print*, Jacobian(N_residual,N_residual)

         write(6,"('Jacobian ...')")
         write(6,"('Inverting Jacobian ...')")
         if( .not. full_inversion ) then
            call invert_Jacobian(N, N_cut,domain,Jacobian)
         else
            ! inverting the full Jacobian
            call dgetrf(N_residual,N_residual,Jacobian,N_residual,ipvt,info)
            if(info/=0) stop "Full Jacobian LU factorization failed."
            call dgetri(N_residual,Jacobian,N_residual,ipvt,work,lwork,info)
            if(info/=0) stop "Full Jacobian inversion failed."
         end if

         !print*, '=================== JACOBIAN INVERSE ================'
         !print*, 'Maxval: ', Maxval(Abs(Jacobian))
         !print*, 'Minval: ', Minval(Abs(Jacobian))
         !print*, '   '

         !open (unit=41,file='JacobianInv.txt', ACTION = "write", STATUS = "replace")
         !do i = 1,N_residual
         !   write(41,*) real( Jacobian(i,:))
         !   write(41,*) "====================="
         !end do
         !close(41)
      end if

      ! Update Jacobian^(-1) with Broyden's method
      ! To save memory, from now on, Jacobian = Jacobian^(-1)
      if ( itr > 0 ) then
         delta_norm = dot_product(delta, delta)
         delta = delta / delta_norm
         r_old = residual - (1.0_long - lam) * r_old
         u = matmul(Jacobian, r_old)
         do i = 1, N_residual
            v(i) = dot_product(delta, Jacobian(:,i))
         end do
         uv = 1.0_long + dot_product(v, r_old)
         u = u / uv
         do i = 1, N_residual
            Jacobian(:,i) = Jacobian(:,i) - v(i) * u
         end do
      endif
      r_old = residual
  
      ! Solve Newtonian equation: Jacobian * delta = - residual
      p     = - matmul(Jacobian,residual)
      fold  = dot_product(residual, residual) * 0.5_long
      slope = - fold * 2.0_long

      ! Update potential and density fields with/without backtracking
      call update_with_linesearch

      ! Calculate thermodynamic properties
      call mu_phi_chain(mu_chain, phi_chain, q)
      if (N_solvent > 0) then
         call mu_phi_solvent(mu_solvent, phi_solvent, q_solvent)
      endif
      call free_energy(N,rho,omega,phi_chain,mu_chain,phi_solvent,mu_solvent, &
                       f_Helmholtz,electrostatic_flag,eps,psi,pressure)
   


      itr = itr + 1

   end do iterative_loop

   ! If fixed unit cell, calculate residual stress before output
   if (.not.domain) stress = scf_stress(N, N_cell_param, dGsq)

   contains
      !----------------------------------------------------------
      ! Accepting Full Newton/Broyden step unconditionally
      !----------------------------------------------------------
      subroutine update_without_linesearch()
      implicit none

      vsum = sqrt(dot_product(p,p))
      if (vsum > stpmax) then
         write(6,*) "residual rescaled"
         p = p * stpmax / vsum
      end if

      lam   = 1.0_long
      delta = lam * p

      ! Update omega
      k = 0
      do i = 1, M
         do alpha = 1, N_monomer
            k = k + 1
            omega(alpha,i+1-ensemble) = omega(alpha,i+1-ensemble) +  delta(k)
         enddo
      enddo
   
      ! Update unit cell 
      if (domain) then
         do i = 1, N_cell_param
            cell_param(i) = cell_param(i) +  delta(M*N_monomer + i)
         enddo
         call make_unit_cell
         call make_correlation(N)
         call make_ksq(G_basis)
         do i = 1, N_cell_param
            call make_dGsq(dGsq(:,i), dGG_basis(:,:,i))
         enddo
      endif

      ! Update density/stress and residuals/error
      call density(N, omega, rho, q, q_solvent)
      if (domain) stress = scf_stress(N, N_cell_param, dGsq)
      if (electrostatic_flag) then
         call calc_dielectric(N,rho,eps)
         call calc_electric_potential(N,rho,eps,psi,.TRUE.)
      endif 

      call make_residual(N, domain, rho, stress, omega, residual, electrostatic_flag, psi, eps)
!     error = maxval(abs(residual))
      error1 = maxval(abs(residual(1:error_index)))
      error2 = maxval(abs(residual(error_index:N_residual)))
      error  = max(error1, error2*stress_rescale)

      end subroutine update_without_linesearch
      !----------------------------------------------------------


      !----------------------------------------------------------
      ! back tracking line search subroutine
      ! instead of accepting Newton/Broyden step straightforwardly
      ! we look for the optimized step along N/B direction
      !----------------------------------------------------------
      subroutine update_with_linesearch()
      implicit none
      real(long), parameter :: ALF=1.0D-4, TOLX=1.0D-7
      real(long)  :: lam2               ! (0,1], controlling step size
      real(long)  :: lamin, lamax       ! Min and Max step size
      real(long)  :: tmplam             !
      real(long)  :: test, temp         ! auxillary variables
      real(long)  :: rhs1, rhs2, f2     !
      real(long)  :: a, b, disc         !

      vsum = sqrt(dot_product(p,p))
      if (vsum > stpmax) then
         write(6,*) "residual rescaled"
         p = p * stpmax / vsum
         slope = slope * stpmax / vsum
      end if

      ! compute minimum lam
      test = 0.0_long
      k = 0
      do i = 1, M
         do alpha = 1, N_monomer
            k = k + 1
            temp = abs(p(k))/max(abs(omega(alpha,i+1-ensemble)), 1.0_long)
            if( temp > test ) test = temp
         enddo
      enddo
      if (domain) then
         do i = 1, N_cell_param
            k = k + 1
            temp = abs(p(k))/max(abs(cell_param(i)), 1.0_long)
            if( temp > test ) test = temp
         enddo
      endif
      lamin = TOLX / test

      lam   = 1.0_long
      ow    = omega
      ocell = cell_param(1:N_cell_param)

      backtrack_loop : do
         delta = lam * p

         ! Update omega
         k = 0
         do i = 1, M
            do alpha = 1, N_monomer
               k = k + 1
               omega(alpha,i+1-ensemble) = ow(alpha,i+1-ensemble) + delta(k)
            enddo
         enddo
         
         ! Update unit cell 
         if (domain) then
            do i = 1, N_cell_param
               cell_param(i) = ocell(i) + delta(M*N_monomer + i)
            enddo
            call make_unit_cell
            call make_correlation(N)
            call make_ksq(G_basis)
            do i = 1, N_cell_param
               call make_dGsq(dGsq(:,i), dGG_basis(:,:,i))
            enddo
         endif

         ! Calculate trial density/stress and residuals/error
         call density(N, omega, rho, q, q_solvent)

         ! Update phi and eps
         if (electrostatic_flag) then
            call calc_dielectric(N,rho,eps)
            call calc_electric_potential(N,rho,eps,psi,.TRUE.)
         endif 

         ! Error calculations
         if (domain) stress = scf_stress(N, N_cell_param, dGsq)
         call make_residual(N, domain, rho, stress, omega, residual, electrostatic_flag, psi, eps)
         ! error = maxval(abs(residual(1:error_index)))
         !error = maxval(abs(residual))
         error1 = maxval(abs(residual(1:error_index)))
         error2 = maxval(abs(residual(error_index+1:N_residual)))
         error  = max(error1, error2*stress_rescale)


         ! backtracking tests
         fnew  = dot_product(residual, residual) * 0.5_long
         if ( lam < lamin ) then
            write(6,"('Minimum lambda in line search reached.')")
            exit backtrack_loop
         else if( fnew < fold + ALF*lam*slope ) then 
            exit backtrack_loop
         else
            if ( lam == 1.0_long ) then
               tmplam = -slope/(2.0_long*(fnew-fold-slope))
            else
               rhs1 = fnew-fold-lam *slope
               rhs2 = f2  -fold-lam2*slope
               a = (rhs1/lam**2-rhs2/lam2**2)/(lam-lam2)
               b = (-lam2*rhs1/lam**2+lam*rhs2/lam2**2)/(lam-lam2)
               if (a==0.0_long) then
                  tmplam = -slope/(2.0_long*b)
               else
                  disc = b*b-3.0_long*a*slope
                  if (disc < 0.0_long) then
                     tmplam = 0.5_long*lam
                  else if (b < 0.0_long) then
                     tmplam = (-b+sqrt(disc))/(3.0_long*a)
                  else
                     tmplam = -slope/(b+sqrt(disc))
                  endif
               endif
               if (tmplam > 0.5_long*lam) tmplam=0.5_long*lam
            endif
         endif

         lam2 = lam
         f2   = fnew
         lam  = max(tmplam, 0.1_long * lam)

      end do backtrack_loop

         !print*, '-----------Delta--------------'
         !print*, 'lamin: ', lamin
         !print*, 'lambda: ', lam
         !print*, 'delta: ', delta
         !print*, '------------------------------'

      end subroutine update_with_linesearch
      !----------------------------------------------------------
   end subroutine iterate_NR
   !==============================================================

   !--------------------------------------------------------------
   !****ip iterate_mod/invert_Jacobian
   ! SUBROUTINE
   !    subroutine invert_Jacobian(N,N_cut,domain,Jacobian)
   ! PURPOSE
   !    Invert the long wave length block (J_corner) of Jacobian
   ! ARGUMENTS
   !    N             = # of basis function
   !    N_cut         = # of long wave length (cut-off dimension)
   !    domain        = false, no stress related entries are considered
   !                    true, stress related entries are combined with
   !                    long wave length entries of Jacobian
   !    Jacobian(:,:) = self explanatory
   ! COMMENT
   !
   !            a |  | b              a | b |
   !           ---\  |---            ---+---|
   !               \ |        ===>    c | d |
   !           _____\|___            -------\
   !            c |  | d                     \
   !
   !           (Jacobian)         (block diagonal)
   !
   !    Use the sparseness of Jacobian to construct inversion.
   !      1) invert J_corner=(a,b;c,d)
   !      2) shift entries of inversion of J_corner back to Jacobian
   !      3) the diagonal blocks between a and d are then inverted revursively
   !    
   !    Note:
   !      The block diagonal Jacobian is not explicitly needed. The internal
   !      subroutine "mapping" is designed to shift a,b,c,d between Jacobian
   !      and J_corner.
   ! SOURCE
   !--------------------------------------------------------------
   subroutine invert_Jacobian(N,N_cut,domain,Jacobian)
   use unit_cell_mod, only     :  N_cell_param
   use chemistry_mod, only     :  ensemble, N_monomer, N_solvent, l0, lp
   implicit none

   integer,    intent(IN)     ::  N     ! # basis function
   integer,    intent(IN)     ::  N_cut ! basis coeff cutoff
   logical,    intent(IN)     ::  domain
   real(long), intent(INOUT)  ::  Jacobian(:,:)
   !***

   real(long)  ::  diagonal(N_monomer,N_monomer)
   integer     ::  info, i

   if( domain ) then
      call mapping(Jacobian,N_residual,J_corner,N_corner,N_cell_param)
   else
      call mapping(Jacobian,N_residual,J_corner,N_corner,0)
   end if

   ! LU factorization and inversion of J_corner
   call dgetrf(N_corner,N_corner,J_corner,N_corner,ipvt(1:N_corner),info)
   if ( info /=0 ) stop "Jacobian corner LU-factor failed."
   call dgetri(N_corner,J_corner,N_corner,ipvt(1:N_corner),work,lwork,info)
   if ( info /= 0) stop "Jacobian corner inversion failed."

   if( domain ) then
      call mapping(J_corner,N_corner,Jacobian,N_residual,N_cell_param)
   else
      call mapping(J_corner,N_corner,Jacobian,N_residual,0)
   end if
 
   ! recursively invert the other diagonal blocks of Jacobian
   do i = max(1,N_cut) + ensemble, N + ensemble -1
!  do i = N_cut + ensemble, N + ensemble -1
      diagonal = Jacobian( (i-1)*N_monomer+1: i*N_monomer,    &
                           (i-1)*N_monomer+1: i*N_monomer)
      ! LU factorization and inversion of diagonal blocks
      call dgetrf(N_monomer,N_monomer,diagonal,N_monomer,ipvt(1:N_monomer),info)
      if ( info /=0 ) stop "Jacobian diagonal LU-factor failed."
      call dgetri(N_monomer,diagonal,N_monomer,ipvt(1:N_monomer),work,lwork,info)
      if ( info /= 0) stop "Jacobian diagonal inversion failed."

      Jacobian( (i-1)*N_monomer+1: i*N_monomer,             &
                (i-1)*N_monomer+1: i*N_monomer) = diagonal

   end do

   contains
      !---------------------------------------------------------
      subroutine mapping(A,na,B,nb,nc)
      implicit none
      real(long)  ::  A(:,:), B(:,:)
      integer     ::  na, nb, nc, ntmp

      ntmp = min(na,nb) - nc
      if( ntmp < 0 ) stop "Error while mapping Jacobian."
      B(1:ntmp,1:ntmp) = A(1:ntmp,1:ntmp)

      if ( nc > 0 ) then
         B(nb-nc+1:nb,1:ntmp)     = A(na-nc+1:na,1:ntmp)
         B(1:ntmp,nb-nc+1:nb)     = A(1:ntmp,na-nc+1:na) 
         B(nb-nc+1:nb,nb-nc+1:nb) = A(na-nc+1:na,na-nc+1:na) 
      end if

      if ( nb > na ) then
         B(ntmp+1:nb-nc,1:ntmp)     = 0.0_long
         B(1:ntmp,ntmp+1:nb-nc)     = 0.0_long
         B(nb-nc+1:nb,ntmp+1:nb-nc) = 0.0_long
         B(ntmp+1:nb-nc,nb-nc+1:nb) = 0.0_long
      end if

      end subroutine mapping
      !---------------------------------------------------------

   end subroutine invert_Jacobian
   !==============================================================

   !---------------------------------------------------------------------
   !****p iterate_mod/make_residual
   ! SUBROUTINE
   !    make_residual( N, domain, rho, stress, omega, residual )
   ! PURPOSE
   !    Calculate array residual(1:N_residual) 
   !    Uses input values of rho, omega, and stress arrays
   ! ARGUMENTS
   !   integer N                  = # of basis functions
   !   logical domain             = T -> domain iteration, F -> fixed cell
   !   real(long) rho(:,:)        = monomer density field
   !   real(long) stress(:)       = stress components
   !   real(long) omega(:,:)      = chemical potential field
   !   real(long) residual(:)     = residual array, dimension(1:N_residual)
   !   logical electrostatic_flag = True for electrostatic interactions
   !   real(long) psi(:)          = electrostatic potential
   !   real(long) eps(:)          = spatially varying dielectric constant
   ! SOURCE
   !---------------------------------------------------------------------
   subroutine make_residual( N, domain, rho, stress, omega, residual, electrostatic_flag, psi, eps)
   use unit_cell_mod, only: N_cell_param
   use chemistry_mod, only: ensemble, N_monomer, chi, dielectric_monomer, &
                            ionic_radii_monomer, charge_monomer, l0, lp,  &
                            solvent_size, N_monomer, N_solvent, kuhn
   use scf_mod,       only: calc_squared_gradient, field_inverse, field_product
   implicit none

   integer,    intent(IN) :: N
   logical,    intent(IN) :: domain
   real(long), intent(IN) :: rho(:,:)
   real(long), intent(IN) :: stress(:)
   real(long), intent(IN) :: omega(:,:)
   real(long), intent(OUT):: residual(:)
   logical,    intent(IN) :: electrostatic_flag
   real(long), intent(IN) :: psi(:)
   real(long), intent(IN) :: eps(:)
   !***

   ! Local variables
   integer    :: i, alpha, beta, l
   real(long) :: xi(N)      ! Lagrange pressure field
   real(long) :: delpsi_sq(N)   ! Squared gradient of psi
   real(long) :: epsinv(N)   ! 1/eps
   real(long) :: epsinvsq(N) ! 1/eps^2
   real(long) :: phiepssq(N_monomer,N) ! phi_alpha/eps^2
   real(long) :: phipinv(N)   ! 1/(phi_p)
   real(long) :: phipinvsq(N) ! 1/(phi_p)^2
   real(long) :: bornblockterm(2,N) ! phi_alpha/(phi_p)^2
   real(long) :: coulterm(2,N) ! d(eps)/dphi * delpsi
   real(long) :: bornfieldterm(2,N_monomer,N)

   if (electrostatic_flag) then
     ! Get |del(psi)|^2
     delpsi_sq = 0.0_long
     call calc_squared_gradient(psi(:),delpsi_sq(:),.FALSE.)

     ! Get powers of eps
     call field_inverse(eps,epsinv)
     call field_product(epsinv,epsinv,epsinvsq)

     ! Get phieepssq
     do beta = 1, N_monomer
         call field_product(epsinvsq,rho(beta,:),phiepssq(beta,:))
     enddo

     ! Get annoying Born terms
     call field_inverse((rho(1,:)+rho(2,:)), phipinv)
     call field_product(phipinv,phipinv,phipinvsq)
     call field_product(phipinvsq,rho(2,:),bornblockterm(1,:))
     call field_product(phipinvsq,rho(1,:),bornblockterm(2,:))
     
     ! Get Coulomb-Dielectric Terms
     call field_product(bornblockterm(1,:), delpsi_sq(:), coulterm(1,:))
     call field_product(bornblockterm(2,:), delpsi_sq(:), coulterm(2,:))

     ! Not generalizable! 
     do alpha = 1,2
        do beta = 3,4
          call field_product(bornblockterm(alpha,:), phiepssq(beta,:), &
                             bornfieldterm(alpha,beta,:) )
        enddo
     enddo
     

     ! Make sure alpha of one is a monomer and not a salt/solvent
     if (dielectric_monomer(1)<0) stop 'Error: monomer 1 cannot be ion'

     ! Lagrange field 
     xi = omega(1,:)

     ! Psi Terms
     do i = 1, N
        ! FH and Coulomb Terms
        xi(i) = xi(i) - sum( chi(:,1)*rho(:,i) ) & 
                + (lp*(kuhn(1)**2)/(8*PI*l0))*coulterm(1,i) & 
                * (dielectric_monomer(1) - dielectric_monomer(2))

        ! Born Terms 
        do beta = 1, N_monomer
          if (ionic_radii_monomer(beta)>0) then
             xi(i) = xi(i) + (l0/2) &
              * (dielectric_monomer(1) - dielectric_monomer(2)) &
              * ( bornfieldterm(1,beta,i)/ionic_radii_monomer(beta) ) &
              / solvent_size(beta-N_monomer+N_solvent)
          endif
        enddo
     enddo
  
     residual = 0.0_long
     ! Incompressiblity and omega residuals
     l = 0
     ! Skip homogenous N=1 star if ensemble is 0 (canonical ensemble)
     do i = 2-ensemble, N
       l = l + 1
       residual(l) = residual(l) + sum( rho(:,i) )

        do alpha = 2, N_monomer
           l = l + 1
           residual(l) = omega(alpha,i) - xi(i)
           residual(l) = residual(l) - sum( chi(:,alpha)*rho(:,i) )

           ! Monomer
           if (dielectric_monomer(alpha)>0) then
              ! Coulomb Term
              residual(l) = residual(l) & 
                 + (lp*(kuhn(1)**2)/(8*PI*l0))*coulterm(alpha,i) &
                 * ((-1)**(alpha+1))*(dielectric_monomer(1) - dielectric_monomer(2))
              ! Born Term
              do beta = 1, N_monomer
                if (ionic_radii_monomer(beta)>0) then
                  residual(l) = residual(l) + ((-1)**(alpha+1))*(l0/2)  & 
                    * (dielectric_monomer(1)-dielectric_monomer(2)) &
                    * (bornfieldterm(alpha,beta,i)/ionic_radii_monomer(beta))& 
                    / solvent_size(beta-N_monomer+N_solvent)

                   !TESTCODE
                   !if (l == 2) then
                   !     print*, '================='
                   !     print*, 'Problem Element: ', alpha, i
                   !     print*, 'Multiplier:      ', (dielectric_monomer(alpha)*l0/2) &
                   !       * (epsinvsq(i)/ionic_radii_monomer(beta)) & 
                   !       / solvent_size(beta-N_monomer+N_solvent)
                   !     print*, '================='
                   !endif
                endif
              enddo


           endif

           ! Ion
           if (ionic_radii_monomer(alpha)>0) then
              ! Coulomb Term
              residual(l) = residual(l) &
                 - (1/(solvent_size(alpha-N_monomer+N_solvent)) ) &
                 * charge_monomer(alpha)*psi(i)
              ! Born Term
              residual(l) = residual(l) & 
                 - (epsinv(i)*l0/2)/(ionic_radii_monomer(alpha) &
                 * solvent_size(alpha-N_monomer+N_solvent))
               
               !TESTDCODE
               !if (l == 3) then
               !   print*, 'Problem Element: ', alpha, i
               !   print*, 'Psi here is:     ', psi(i)
               !   print*, 'Multiplier:      ', (( (4*PI*l0)/(solvent_size(alpha-N_monomer+N_solvent)) ) &
               !   * charge_monomer(alpha) )
               !   print*, 'Residual is:     ', residual(l)
               !   print*, '================='
               !endif
           endif
           
        enddo
      enddo

   else

    ! Lagrange field 
     xi = omega(1,:)
     do i = 1, N
        xi(i) = xi(i) - sum( chi(:,1)*rho(:,i) )
     enddo
  
     residual = 0.0_long
     ! Incompressiblity and omega residuals
     l = 0
     ! Skip homogenous N=1 star if ensemble is 0 (canonical ensemble)
     do i = 2-ensemble, N
       l = l + 1
       residual(l) = residual(l) + sum( rho(:,i) )

        do alpha = 2, N_monomer
           l = l + 1
           residual(l) = omega(alpha,i) - xi(i)
           residual(l) = residual(l) - sum( chi(:,alpha)*rho(:,i) )
        enddo
      enddo
   endif

   if (ensemble == 1) residual(1) = residual(1) - 1.0_long

   ! Stress residuals
   if (domain) residual(l+1:l+N_cell_param) = stress(:)

   end subroutine make_residual
   !====================================================================


   !--------------------------------------------------------------------
   !****p iterate_mod/Jacobian_response
   ! SUBROUTINE
   !    Jacobian_response(N, N_cut, domain, omega, Jacobian)
   ! PURPOSE
   !    Construct Jacobian from response functions
   ! SOURCE
   !--------------------------------------------------------------------
   subroutine Jacobian_response(N, N_cut, domain, omega, Jacobian, rho, &
                                electrostatic_flag, eps, psi)
   use chemistry_mod, only : ensemble, N_monomer, chi, charge_monomer, &
                             dielectric_monomer, ionic_radii_monomer, l0, lp, &
                             N_solvent, solvent_size, kuhn
   use unit_cell_mod, only : N_cell_param
   use scf_mod,       only: field_inverse, field_product
   implicit none

   
   integer, intent(IN)       :: N
   integer, intent(IN)       :: N_cut
   logical, intent(IN)       :: domain
   real(long), intent(IN)    :: omega(:,:)
   real(long), intent(OUT)   :: Jacobian(:,:)
   real(long), intent(IN)    :: rho(:,:)
   logical, intent(IN)       :: electrostatic_flag
   real(long), intent(IN)    :: eps(:)
   real(long), intent(IN)    :: psi(:)
   !***
   
   ! Local Variables
   real(long)  :: drho_domega(:,:,:,:) 
   real(long)  :: dxi_domega(:,:,:) 
   real(long)  :: dpsi_domega(:,:,:)
   real(long)  :: ddpsisq_domega(:,:,:)
   real(long)  :: deps_domega(:,:,:)
   real(long)  :: dborn_domega(:,:,:,:,:)
   real(long)  :: dcoul_domega(:,:,:,:)
   allocatable :: drho_domega, dxi_domega, dpsi_domega, &
                  ddpsisq_domega, deps_domega, dborn_domega, dcoul_domega
   real(long)  :: drho_dcell(N_monomer, N, N_cell_param)
   real(long)  :: dstress_dcell(N_cell_param,N_cell_param)
   real(long)  :: dxi_dcell(N,N_cell_param) 
   real(long)  :: dpsi_dcell(N,N_cell_param)
   real(long)  :: ddpsisq_dcell(N,N_cell_param)
   real(long)  :: deps_dcell(N,N_cell_param)
   real(long)  :: dborn_dcell(2,N_monomer,N,N_cell_param)
   real(long)  :: dcoul_dcell(2,N,N_cell_param)

   ! Auxillary Variables (try to reduce these)
   real(long) :: epsinv(N)     ! 1/eps
   real(long) :: epsinvsq(N)   ! 1/eps^2
   !real(long) :: epsinvcub(N) ! 1/eps^3
   !real(long) :: temp1(N_monomer,N)
   !real(long) :: temp2(N_monomer,N)
   real(long) :: temp(N)
   real(long) :: tt1, tt2, tt3

   ! indices and temporary variables needed ...
   integer    :: j, beta, l   ! omega field indices
   integer    :: i, alpha, k  ! residual indices
   integer    :: gamma        ! extra index
   !integer   :: n_cut        ! cut + 1, index
   integer    :: imin, imax   ! output indices
   integer    :: info   ! output indices

   allocate(drho_domega(N_monomer,N,N_monomer,N),stat=info)
   if( info /= 0) stop "drho_domega memory allocati on failed in Jacobian_Response"

   allocate(dxi_domega(N,N_monomer,N),stat=info)
   if( info /= 0) stop "dxi_domega memory allocation failed in Jacobian_Response"

   !TESTCODE:
   call cpu_time(tt1)
   ! drho_domega (density response)
   call response_pd_omega(N,N_cut,omega,drho_domega) 
   call cpu_time(tt2)
   print*, 'response_pd_omega (s): ', (tt2-tt1)

   if (electrostatic_flag) then
      allocate(dpsi_domega(N,N_monomer,N),stat=info)
      if( info /= 0) stop "dpsi_domega memory allocation failed in Jacobian_Response"

      allocate(ddpsisq_domega(N,N_monomer,N),stat=info)
      if( info /= 0) stop "ddpsisq_domega memory allocation failed in Jacobian_Response"

      allocate(deps_domega(N,N_monomer,N),stat=info)
      if( info /= 0) stop "deps_domega memory allocation failed in Jacobian_Response"

      allocate(dborn_domega(2,N_monomer,N,N_monomer,N),stat=info)
      if(info /= 0) stop "dborn_domega memory allocation failed in Jacobian_Response"
      
      allocate(dcoul_domega(2,N,N_monomer,N),stat=info)
      if(info /= 0) stop "dborn_domega memory allocation failed in Jacobian_Response"
      
      ! dpsi_domega (psi response)
      call response_psid_omega(N,N_cut,omega,dpsi_domega,ddpsisq_domega, &
                               deps_domega, dborn_domega, dcoul_domega)

      ! Old (incorrect) deps_domega
      !deps_domega = 0.0_long
      !do alpha = 1, N_monomer
      !   if (dielectric_monomer(alpha)>0) then
      !      deps_domega(:,:,:) = deps_domega(:,:,:) &
      !         + dielectric_monomer(alpha)*drho_domega(alpha,:,:,:)
      !   endif
      !enddo

      ! Get powers of epsilon
      call field_inverse(eps,epsinv)
      call field_product(epsinv,epsinv,epsinvsq)
      !call field_product(epsinv,epsinvsq,epsinvcub)
   
      call cpu_time(tt3)
      print*, 'response_psid_omega (s): ', (tt3-tt2)
   endif

   ! dxi_domega (Lagrange response)
   dxi_domega = 0.0_long

   ! Electrostatic Case
   if (electrostatic_flag) then
      ! Error Checks
      if (dielectric_monomer(1)<0) stop "monomer 1 must have dielectric constant"

      do i = 1, N
        do alpha = 1, N_monomer
          ! Born Term Helper
          !temp1 = 0.0_long
          !temp2 = 0.0_long
          !do beta = 1, N_monomer
          !  call field_product(epsinvsq,drho_domega(beta,:,alpha,i),temp1(beta,:))
          !  call field_product(epsinvcub,rho(beta,:),temp2(beta,:))
          !  call field_product(temp2(beta,:),deps_domega(:,alpha,i),temp2(beta,:))
          !enddo

          do j = 1, N
            ! Flory Huggins Term
            dxi_domega(j,alpha,i) = - sum( chi(:,1)*drho_domega(:,j,alpha,i) )

            ! Coulomb Term
            dxi_domega(j,alpha,i) = dxi_domega(j,alpha,i) & 
               + (lp*(kuhn(1)**2)/(8*PI*l0))*dcoul_domega(1,j,alpha,i) &
               * (dielectric_monomer(1) - dielectric_monomer(2))

            ! Born Term
            do beta = 1, N_monomer
               if (ionic_radii_monomer(beta)>0) then
                  dxi_domega(j,alpha,i) = dxi_domega(j,alpha,i) + (l0/2) &
                     * (dielectric_monomer(1) - dielectric_monomer(2)) &
                     * dborn_domega(1,beta,j,alpha,i) &
                     / ( solvent_size(beta-N_monomer+N_solvent) &
                        *ionic_radii_monomer(beta) )
                  !dxi_domega(j,alpha,i) = dxi_domega(j,alpha,i) &
                  !   + ((dielectric_monomer(1)*l0)/(2*ionic_radii_monomer(beta) &
                  !   * solvent_size(beta-N_monomer+N_solvent))) &
                  !   * (temp1(beta,j) - 2*temp2(beta,j))
               endif
            enddo

          enddo
        enddo
        ! Same Field correction
        dxi_domega(i,1,i) = dxi_domega(i,1,i) + 1.0_long
      enddo 

   ! Non-electrostatic Case
   else
      do i = 1, N
        do alpha = 1, N_monomer
          do j = 1, N
            dxi_domega(j,alpha,i) = - sum( chi(:,1)*drho_domega(:,j,alpha,i) )
          enddo
        enddo
        dxi_domega(i,1,i) = dxi_domega(i,1,i) + 1.0_long
      enddo 
   endif

   ! Update Jacobian, outmost loop is omega
   Jacobian = 0.0_long

   ! electrostatic case
   if(electrostatic_flag) then
      do j = 2 - ensemble, N
         do beta = 1, N_monomer
            l = beta + N_monomer*(j-2+ensemble)

            ! Born Term Helpers
            !temp1 = 0.0_long
            !temp2 = 0.0_long
            !do alpha = 1, N_monomer
            !   call field_product(epsinvsq,drho_domega(alpha,:,beta,j),temp1(alpha,:))
            !   call field_product(epsinvcub,rho(alpha,:),temp2(alpha,:))
            !   call field_product(temp2(alpha,:),deps_domega(:,beta,j),temp2(alpha,:))
            !enddo
            temp = 0.0_long
            call field_product(epsinvsq, deps_domega(:,beta,j), temp)

            do i = 2 - ensemble, N
               k = 1 + N_monomer*(i-2+ensemble)
               Jacobian(k,l) = sum( drho_domega(:,i,beta,j) )
               !if ( k == l) then
               !   print *, k
               !   print *, Jacobian(k,k)
               !end if  

               do alpha = 2, N_monomer
                  k = alpha + N_monomer*(i-2+ensemble)
                  Jacobian(k,l) = - sum( chi(:,alpha)*drho_domega(:,i,beta,j) ) &
                                  - dxi_domega(i,beta,j)

                  ! Monomer Case
                  if (dielectric_monomer(alpha)>0) then
                    ! Coulomb Term
                     Jacobian(k,l) = Jacobian(k,l) & 
                        + (lp*(kuhn(1)**2)/(8*PI*l0))*dcoul_domega(alpha,i,beta,j) &
                        * ((-1)**(alpha+1)) &
                        * (dielectric_monomer(1)-dielectric_monomer(2))

                    ! Born Term
                     do gamma = 1, N_monomer 
                        if (ionic_radii_monomer(gamma)>0) then
                           Jacobian(k,l) = Jacobian(k,l) + (l0/2) &
                             * (dielectric_monomer(1) - dielectric_monomer(2)) &
                             * ((-1)**(alpha+1))*dborn_domega(alpha,gamma,i,beta,j) &
                             / ( solvent_size(gamma-N_monomer+N_solvent) &
                                *ionic_radii_monomer(gamma) )
                             !+ ((dielectric_monomer(alpha)*l0)/(2*ionic_radii_monomer(gamma) & 
                             !* solvent_size(gamma-N_monomer+N_solvent))) &
                             !* (temp1(gamma,i) - 2*temp2(gamma,i))
                             !
                             !* ((drho_domega(gamma,i,beta,j)/(eps(i)**2)) &
                             !-  ((2*rho(gamma,i)/(eps(i)**3))*deps_domega(i,beta,j)) )
                        endif
                     enddo
                  endif

                  ! Ion Case
                  if (ionic_radii_monomer(alpha)>0) then
                     ! Coulomb Term
                     Jacobian(k,l) = Jacobian(k,l) & 
                        - (charge_monomer(alpha))*dpsi_domega(i,beta,j) & 
                        / solvent_size(alpha-N_monomer+N_solvent)

                     ! Born Term
                      Jacobian(k,l) = Jacobian(k,l) &
                        + (l0/(2*ionic_radii_monomer(alpha))) &
                        * temp(i) / solvent_size(alpha-N_monomer+N_solvent)
                     !Jacobian(k,l) = Jacobian(k,l) &
                     !   + (l0/(2*ionic_radii_monomer(alpha)* eps(i)**2)) &
                     !   * deps_domega(i,beta,j) &
                     !   / solvent_size(alpha-N_monomer+N_solvent)
                  endif
               enddo
            enddo

            ! Add in the same term corrections (domega_beta/domega_beta)
            if (beta > 1) Jacobian(l,l) = Jacobian(l,l) + 1.0_long
            !Jacobian(l,l) = Jacobian(l,l) + 1.0_long
            !print *, beta,l
            !print *, Jacobian(l,l)
         
         enddo
      enddo

   ! Non-electrostatic Case
   else 
      do j = 2 - ensemble, N
         do beta = 1, N_monomer
            l = beta + N_monomer*(j-2+ensemble)
            ! incompressibility and scf omega response
            do i = 2 - ensemble, N
               k = 1 + N_monomer*(i-2+ensemble)
               Jacobian(k,l) = sum( drho_domega(:,i,beta,j) )
               !if ( k == l) then
               !   print *, k
               !   print *, Jacobian(k,k)
               !end if  

               do alpha = 2, N_monomer
                  k = alpha + N_monomer*(i-2+ensemble)
                  Jacobian(k,l)= - sum( chi(:,alpha)*drho_domega(:,i,beta,j) )&!
                                 - dxi_domega(i,beta,j)
               enddo
            enddo
            if (beta > 1) Jacobian(l,l) = Jacobian(l,l) + 1.0_long
            !Jacobian(l,l) = Jacobian(l,l) + 1.0_long
            !print *, beta,l
            !print *, Jacobian(l,l)
         
         enddo
      enddo
   endif


   ! Variable unit cell case
   if ( domain ) then
      ! Electrostatic Case
      if (electrostatic_flag) then
         ! drho_dcell, dpsi_dcell
         call response_pd_cell(N,omega,drho_dcell,dstress_dcell, &
                               electrostatic_flag,dpsi_dcell, &
                               ddpsisq_dcell, deps_dcell, dborn_dcell, &
                               dcoul_dcell)

            ! Testcode:
            !print*, '------------------------------------------------'
            !print*, 'drho_dcell (Electrostatic)'
            !print*, '     l0 is: ', l0
            !print*, drho_dcell(:,:,:)
            !print*, '============'
            !print*, '(Electrostatic)'
            !print*, dpsi_dcell(:,:)
            !print*, '============'
            !print*, '(Electrostatic)'
            !print*, ddpsisq_dcell(:,:)
            !print*, '============'
            !print*, dstress_dcell
            !print*, '------------------------------------------------'

         ! deps_dcell CALCULATION NOT CORRECT
         !deps_dcell(:,:) = 0.0_long
         !do alpha = 1, N_monomer
         !  if (dielectric_monomer(alpha)>0) then
         !      deps_dcell(:,:) = deps_dcell(:,:) &
         !         + dielectric_monomer(alpha)*drho_dcell(alpha,:,:)
         !   endif
         !enddo

         ! dxi_dcell (Lagrangian response)
         do j = 1, N_cell_param 
            !temp1 = 0.0_long
            !temp2 = 0.0_long
            !do beta = 1, N_monomer
            !   call field_product(epsinvsq,drho_dcell(beta,:,j),temp1(beta,:))
            !   call field_product(epsinvcub,rho(beta,:),temp2(beta,:))
            !   call field_product(temp2(beta,:),deps_dcell(:,j),temp2(beta,:))
            !enddo

            do i = 1, N
               ! Flory Huggins Term
               dxi_dcell(i,j) = - sum ( chi(:,1)*drho_dcell(:,i,j) )

               ! Coulomb Term
               dxi_dcell(i,j) = dxi_dcell(i,j) & 
                  + (lp*(kuhn(1)**2)/(8*PI*l0))*dcoul_dcell(1,i,j) &
                  * (dielectric_monomer(1) - dielectric_monomer(2))
               
               !dxi_dcell(i,j) = dxi_dcell(i,j) &
               !   + (lp*(kuhn(1)**2)/(8*PI*l0))*dielectric_monomer(1)*ddpsisq_dcell(i,j)

               ! Born Term
               do beta = 1, N_monomer
                  if (ionic_radii_monomer(beta)>0) then
                     dxi_dcell(i,j) = dxi_dcell(i,j) + (l0/2) &
                        * (dielectric_monomer(1)-dielectric_monomer(2)) &
                        * dborn_dcell(1,beta,i,j) &
                        / ( solvent_size(beta-N_monomer+N_solvent) &
                           *ionic_radii_monomer(beta))   
                        !+ ((dielectric_monomer(1)*l0)/(2*ionic_radii_monomer(beta) &
                        !* solvent_size(beta-N_monomer+N_solvent))) &
                        !* (temp1(beta,i) - 2*temp2(beta,i))
                        !* ((drho_dcell(beta,i,j)/(eps(i)**2)) &
                        !- ((2*rho(beta,i)/(eps(i)**3))*deps_dcell(i,j)) )
                  endif
               enddo

            enddo
         enddo

         ! TODO (Does this need a change?)
         ! Stress Response to omega, by transposition of drho_dcell
         do i = 1, N_cell_param
            k = i + N_monomer*(N-1+ensemble)
            do j = 2 - ensemble, N
               do alpha = 1, N_monomer
                  l = alpha + N_monomer*(j-2+ensemble) 
                  !Jacobian(k,l) = chain_length * drho_dcell(alpha,j,i)
                  Jacobian(k,l) = drho_dcell(alpha,j,i)
               enddo
            enddo
         enddo

      
         ! Populate the Jacobian 
         do j = 1, N_cell_param
            l = j + N_monomer * (N-1+ensemble)

            !temp1 = 0.0_long
            !temp2 = 0.0_long
            !do beta = 1, N_monomer
            !   call field_product(epsinvsq,drho_dcell(beta,:,j),temp1(beta,:))
            !   call field_product(epsinvcub,rho(beta,:),temp2(beta,:))
            !   call field_product(temp2(beta,:),deps_dcell(:,j),temp2(beta,:))
            !enddo
            temp = 0.0_long
            call field_product(epsinvsq,deps_dcell(:,j),temp)

            do i = 2 - ensemble, N
               ! Incompressibilty Response
               k = 1 + N_monomer * (i-2+ensemble)
               Jacobian(k,l) = sum( drho_dcell(:,i,j) )

               ! Omega Response
               do alpha = 2, N_monomer
                  k = alpha + N_monomer*(i-2+ensemble)

                  ! Flory Huggins/Incompressibility Term
                  Jacobian(k,l) = - sum( chi(:,alpha)*drho_dcell(:,i,j) )&!
                                  - dxi_dcell(i,j)

                  ! Monomer Case 
                  if (dielectric_monomer(alpha)>0) then
                     ! Coulomb Term
                     Jacobian(k,l) = Jacobian(k,l) &
                         + (lp*(kuhn(1)**2)/(8*PI*l0))*dcoul_dcell(alpha,i,j) &
                         * ((-1)**(alpha+1)) &
                         * (dielectric_monomer(1) - dielectric_monomer(2))
                     
                     !Jacobian(k,l) = Jacobian(k,l) & 
                     !   + (lp*(kuhn(1)**2)/(8*PI*l0))*dielectric_monomer(alpha) &
                     !   * ddpsisq_dcell(i,j)

                     ! Born Term
                     do gamma = 1, N_monomer 
                        if (ionic_radii_monomer(gamma)>0) then
                           Jacobian(k,l) = Jacobian(k,l) + (l0/2) &
                             * (dielectric_monomer(1) - dielectric_monomer(2)) &
                             * ((-1)**(alpha+1))*dborn_dcell(alpha,gamma,i,j) &
                             / ( solvent_size(gamma-N_monomer+N_solvent) &
                                *ionic_radii_monomer(gamma) )                       
                             !+ ((dielectric_monomer(alpha)*l0)/(2*ionic_radii_monomer(gamma) &
                             !* solvent_size(gamma-N_monomer+N_solvent))) &
                             !* (temp1(gamma,i) - 2*temp2(gamma,i))
                             !* ((drho_dcell(gamma,i,j)/(eps(i)**2)) &
                             !-  ((2*rho(gamma,i)/(eps(i)**3))*deps_dcell(i,j)) )
                        endif
                     enddo
                  endif 

                  ! Ion Case
                  if (ionic_radii_monomer(alpha)>0) then
                     ! Coulomb Term
                     Jacobian(k,l) = Jacobian(k,l) & 
                        - (charge_monomer(alpha))*dpsi_dcell(i,j) &
                        / solvent_size(alpha-N_monomer+N_solvent)

                     ! Born Term
                     Jacobian(k,l) = Jacobian(k,l) &
                        + (l0/(2*ionic_radii_monomer(alpha))) &
                        * temp(i) / solvent_size(alpha-N_monomer+N_solvent)
                  endif 
                  !print*, '(', k, ',', l, '):   ', Jacobian(k,l)
               enddo
               !print*, '(', k, ',', l, '):   ', Jacobian(k,l)
            enddo
            ! Stress Response (dcell param /dcell param?)
            Jacobian(k+1:k+N_cell_param,l) = dstress_dcell(:,j)
            !print*, '====== stress corner ========'
            !print*, '(k,l):   (',(k+1),',',l,')'
            !print*, dstress_dcell 
         enddo

      ! Non-electrostatic case   
      else
         call response_pd_cell(N,omega,drho_dcell,dstress_dcell)

            ! Testcode:
            !print*, '------------------------------------------------'
            !print*, 'drho_dcell (Neutral)'
            !print*, drho_dcell(:,:,:)
            !print*, '============'
            !print*, dstress_dcell
            !print*, '------------------------------------------------'

         ! dxi_dcell (Lagrangian response)
         do j = 1, N_cell_param
            do i = 1, N
               dxi_dcell(i,j) = - sum ( chi(:,1)*drho_dcell(:,i,j) )
            enddo
         enddo

         ! Stress Response to omega, by transpostion of drho_dcell
         do i = 1, N_cell_param
            k = i + N_monomer*(N-1+ensemble)
            do j = 2 - ensemble, N
               do alpha = 1, N_monomer
                  l = alpha + N_monomer*(j-2+ensemble) 
                  !Jacobian(k,l) = chain_length * drho_dcell(alpha,j,i)
                  Jacobian(k,l) = drho_dcell(alpha,j,i)
               enddo
            enddo
         enddo


         do j = 1, N_cell_param
            l = j + N_monomer * (N-1+ensemble)

            do i = 2 - ensemble, N
               ! Incompressibilty Response
               k = 1 + N_monomer * (i-2+ensemble)
               Jacobian(k,l) = sum( drho_dcell(:,i,j) )

               ! Omega Response
               do alpha = 2, N_monomer
                  k = alpha + N_monomer*(i-2+ensemble)
                  Jacobian(k,l) = - sum( chi(:,alpha)*drho_dcell(:,i,j) )&!
                                  - dxi_dcell(i,j)

                  !print*, '(', k, ',', l, '):   ', Jacobian(k,l)
               enddo

               ! TESTCODE
               !print*, '(', k, ',', l, '):   ', Jacobian(k,l)
            enddo

            ! Stress Response
            Jacobian(k+1:k+N_cell_param,l) = dstress_dcell(:,j)

         enddo
      endif
   end if

   if( allocated(drho_domega) ) deallocate(drho_domega)
   if( allocated(dxi_domega) ) deallocate(dxi_domega)
   if( allocated(dpsi_domega) ) deallocate(dpsi_domega)
   if( allocated(ddpsisq_domega) ) deallocate(ddpsisq_domega)
   if( allocated(deps_domega) ) deallocate(deps_domega)
   if( allocated(dborn_domega) ) deallocate(dborn_domega)
   if( allocated(dcoul_domega) ) deallocate(dcoul_domega)

   end subroutine Jacobian_response
   !===================================================================

   !--------------------------------------------------------------------
   !****p iterate_mod/iterate_AM_startup
   ! SUBROUTINE
   !    iterate_AM_startup(N)
   ! PURPOSE
   !   1) Allocate arrays needed by iterate_AM
   ! ARGUMENTS
   !   integer N      = # of basis functions
   ! SOURCE
   !---------------------------------------------------------------------
   subroutine iterate_AM_startup(N)
   use unit_cell_mod, only : N_cell_param
   use chemistry_mod, only : N_monomer, ensemble
   
   integer, intent(IN)       :: N        ! # of basis functions 
   integer                   :: info     ! message variable for dgetri
!!   real(long), allocatable   :: Umn(:,:) ! Local Umn variable to calculate the
                                         ! work space

   if (allocated(ipvt)) deallocate(ipvt)
   if (allocated(work)) deallocate(work)
   if (allocated(dev)) deallocate(dev)
   if (allocated(omega_hist)) deallocate(omega_hist)
   if (allocated(Umn)) deallocate(Umn)

   allocate(ipvt(N_hist))
   allocate(work(N_hist))
   allocate(dev(N_hist+1,N-1,N_monomer))
   allocate(omega_hist(N_hist+1,N-1,N_monomer))
   allocate(Umn(N_hist,N_hist))


   ! estimate the opitimal workspace dimension
   lwork = -1
   call dgetri(N_hist,Umn,N_hist,ipvt,work,lwork,info)
   lwork = work(1)
   if (allocated(work)) deallocate(work)
   allocate(work(lwork))

   ! Umn changes size after every iteration so will be alloated in main (next)
   ! subroutine, however, here we allocated it to calculate the maximum work
   ! space needed for the AM scheme.
   if (allocated(Umn)) deallocate(Umn)
   
   end subroutine iterate_AM_startup


   !---------------------------------------------------------------------
   !****p iterate_mod/iterate_AM
   ! SUBROUTINE
   !    subroutine iterate_AM
   ! PURPOSE
   !    Anderson-Mixing iteration of the SCFT
   !    Note: Presently AM scheme is applicable only for 
   !          the canonical ensemble (ensemble == 0)
   ! COMMENT
   !   1) Algorithm assumes that density_startup, iterate_AM_startup,
   !      and make_unit_cell have been called prior to iterate_AM.
   !
   !   2) If domain = .true., routine iterates both omega fields and 
   !      unit cell parameters. However, in contrast to the NR scheme,
   !      the fields and unit cell parameters are iterated sequentially, 
   !      and not simultaneously
   ! SOURCE
   !--------------------------------------------------------------------
   subroutine iterate_AM(  &
              N,           &!  # of basis functions
              omega,       &!  chemical potential field (IN/OUT)
              itr,         &!  actual number of interations
              converge,    &!  = .true. if converged
              error,       &!  final error = max(residuals)
              rho,         &!  monomer density field
              f_Helmholtz, &!  Helmholtz free energy per monomer / kT
              f_components,&! Helmholtz free energy components
              pressure,    &!  pressure * monomer volume / kT
              stress,      &!  d(free energy)/d(cell parameters)
              eps,               &! spatially varying dielectric field
              psi,               &! electric potential 
              electrostatic_flag &! true for e. calculations
              )
   !--------------------------------------------------------------------

   use io_mod, only        : output
   use unit_cell_mod, only : N_cell_param, cell_param, &
                             make_unit_cell, G_basis, dGG_basis
   use basis_mod, only     : make_dGsq
   use chemistry_mod, only : N_monomer, ensemble, chi, mu_chain, mu_solvent, &
                             phi_solvent, phi_chain, N_chain, N_solvent 
   use scf_mod, only       : density, scf_stress, set_omega_uniform, &
                             mu_phi_chain, mu_phi_solvent, free_energy
   use grid_mod, only      : make_ksq

   ! Arguments
   integer, intent(IN)         :: N        
   real(long), intent(INOUT)   :: omega(:,:)  ! omega(N_monomer,N)
   integer, intent(OUT)        :: itr         
   logical, intent(OUT)        :: converge
   real(long), intent(OUT)     :: error
   real(long), intent(OUT)     :: rho(:,:)    ! rho(N_monomer,N)
   real(long), intent(OUT)     :: f_Helmholtz
   real(long), intent(OUT)     :: f_components(:)
   real(long), intent(OUT)     :: pressure
   real(long), intent(OUT)     :: stress(:)   ! stress(N_cell_param)
   real(long), intent(OUT)     :: eps(:)
   real(long), intent(OUT)     :: psi(:)
   logical, intent(IN)         :: electrostatic_flag 
   !***

   ! parameter
   real(long), parameter :: large = 100.0
   real(long), parameter :: stress_rescale=1.0d+2
   real(long), parameter :: STPMX=100.0_long ! Original: 100

   ! local Variables
   integer    :: i, k, alpha, M, info, itr_stress, itr_relax
   integer    :: error_index
   real(long) :: dGsq(N,N_cell_param), q(N_chain), q_solvent(N_solvent)
   real(long) :: vsum, stpmax
   real(long) :: delta_norm, u(N_residual), v(N_residual), uv
   real(long) :: dstress_dcell(N_cell_param,N_cell_param), p(N_cell_param)

   ! Parameters for Anderson-Mixing
   real(long), allocatable  :: Vm(:,:), Cn(:,:)
   real(long)               :: WW(N-1), DD(N-1)
   real(long)               :: elm(N-1)
   real(long)               :: dev_2, omega_2
   real(long)               :: lam_AM
   integer                  :: mm,nn, itr_prev

   dev=0
   omega_hist=0

   M = N - 1 + ensemble 

   ! Initialization
   call density(N, omega, rho, q, q_solvent)
   ! call mu_phi(mu_chain,phi_chain,q,mu_solvent,phi_solvent,q_solvent)
   call mu_phi_chain(mu_chain, phi_chain, q)
   if (N_solvent > 0) then
      call mu_phi_solvent(mu_solvent, phi_solvent, q_solvent)
   endif
   !!   call make_correlation(N)
   do i = 1, N_cell_param
      call make_dGsq(dGsq(:,i), dGG_basis(:,:,i))
   enddo

   ! set the homogeneous coefficents
   if (ensemble == 0) call set_omega_uniform(omega,electrostatic_flag)

   itr = 0
   
   ! AM scheme needs atleast two previous iterations to mix and produce
   ! output for the next. Here, we use first iteration as a guess supplied
   ! and then create two (although need one) more iterations by linear 
   ! combinations of the previous iterations    
   write(6,*)'Creating History for Anderson-Mixing Scheme'
   do i=1,3
      itr = itr + 1
      do alpha=1,N_monomer
         omega(alpha,2:) = 0.9**(itr-1)*omega(alpha,2:)
      end do
      call density(N, omega, rho, q, q_solvent)
      do alpha=1,N_monomer
         if(N_monomer==2)then
            do k=1,N-1
               dev(itr,k,alpha) = sum(chi(:,alpha)*rho(:,k+1))&!
                        + sum(omega(:,k+1))/N_monomer - omega(alpha,k+1)
            end do

         elseif(N_monomer==3)then
            do k=1,N-1
               dev(itr,k,alpha) = sum(chi(:,alpha)*rho(:,k+1)) &!
                      + ( sum(omega(:,k+1)) + chi(1,2)*rho(3,k+1) &!
             +  chi(2,3)*rho(1,k+1) + chi(1,3)*rho(2,k+1) )/N_monomer &!
                      - omega(alpha,k+1)
            end do
         endif
         omega_hist(itr,:,alpha) = omega(alpha,2:)

      end do
   end do   
   itr = itr - 1   !  i=1 in the previous loop is 0th iteration

    !  Calculating error in omega
    dev_2   = 0.0
    omega_2 = 0.0
    do alpha = 1,N_monomer
       dev_2 = dev_2 + norm_2(dev(itr,:,alpha))**2.0
       omega_2 = omega_2 + norm_2(omega(alpha,2:))**2.0
    end do
    error   = (dev_2/omega_2)**0.5
   
   ! Calculate thermodynamic properties
   call free_energy(N,rho,omega,phi_chain,mu_chain,phi_solvent,mu_solvent, &
                    f_Helmholtz,electrostatic_flag,eps,psi,pressure)
   if (domain) stress = scf_stress(N, N_cell_param, dGsq)


   itr_stress = 0
   stress_loop : do 

      itr_stress = itr_stress + 1
       
   iterative_loop : do

      itr = itr + 1
      if(itr<N_hist+1)then
         itr_prev = itr-1
         lam_AM   = 1-0.9**(itr-1)
      else
         lam_AM   = 1.0
         itr_prev = N_hist
      end if

      write(6,*)
      write(6,"('Iteration ',i3)") itr

      ! Output to log file
      call output(error,                     'error       =',o=6,f='L')
      if (domain) then
         call output(stress,N_cell_param,    'stress      =',o=6,f='L')
         call output(cell_param,N_cell_param,'cell_param  =',o=6,f='L')
      endif
      call output(f_Helmholtz,               'f_Helmholtz =',o=6,f='L')
      if (ensemble == 1) then
         call output(pressure,               'pressure    =',o=6,f='L')
      endif 

      ! Convergence criteria
      if ( error < error_max ) then
         converge = .true.
         write(*,*)'1'
         exit iterative_loop
      else if ( itr > max_itr ) then
         converge = .false.
         write(*,*)'2'
         exit iterative_loop
      else if ( error > 100.0*large ) then
         converge = .false.
         write(*,*)'3'
         exit iterative_loop
      endif

      if(itr==3) write(6,*)'Anderson-Mixing Starts Now '

      WW  = 0.0
      DD  = 0.0

      if(itr<N_hist+2)then
         allocate(Umn(itr_prev,itr_prev)) 
         allocate(Vm(itr_prev,1)) 
         allocate(Cn(itr_prev,1)) 
      endif

      Umn = 0
      Vm  = 0 
      Cn  = 0

      do mm=1,itr_prev
         do nn=mm,itr_prev
            
           do alpha=1,N_monomer 
              elm = 0
              elm = (dev(itr_prev+1,:,alpha) - dev(itr_prev+1-mm,:,alpha))&!
                  *(dev(itr_prev+1,:,alpha) - dev(itr_prev+1-nn,:,alpha))
              Umn(mm,nn) = Umn(mm,nn) + sum(elm)
           end do
           Umn(nn,mm) = Umn(mm,nn)       
     
        end do

          do alpha=1,N_monomer 
              elm = 0
              elm = (dev(itr_prev+1,:,alpha) - dev(itr_prev+1-mm,:,alpha))&!
                  *dev(itr_prev+1,:,alpha)
              Vm(mm,1) = Vm(mm,1) + sum(elm)   
          end do

      end do  

      ! Inverting Umn
      call dgetrf(itr_prev,itr_prev,Umn,itr_prev,ipvt,info)
      if(info/=0) stop "Full Umn LU factorization failed."
      call dgetri(itr_prev,Umn,itr_prev,ipvt,work,lwork,info)
      if(info/=0) stop "Full Umn inversion failed."
      Cn = matmul(Umn,Vm)
       
      do alpha = 1,N_monomer
         WW = omega(alpha,2:)
         DD = dev(itr_prev+1,:,alpha) 
         do nn=1,itr_prev
       
            WW = WW + Cn(nn,1)*(omega_hist(itr_prev+1-nn,:,alpha) - omega_hist(itr_prev+1,:,alpha))         
            DD = DD + Cn(nn,1)*(dev(itr_prev+1-nn,:,alpha) - dev(itr_prev+1,:,alpha))
      
        end do

        omega(alpha,2:) = WW + lam_AM*DD

      end do      
   
     ! Note that Umn, Vm, Cn change sizes after every iteration 
     ! until N_hist number of iterations are completed, so have 
     ! to deallocate them everytime and allocate at the start of step 
      if(itr<N_hist+1)then 
         deallocate(Umn)
         deallocate(Vm)
         deallocate(Cn)
      endif
 
      call density(N, omega, rho, q, q_solvent)
      if (domain) stress = scf_stress(N, N_cell_param, dGsq)

      if(itr<N_hist+1)then

         dev_2   = 0.0
         omega_2 = 0.0
         do alpha=1,N_monomer

            if(N_monomer==2)then
               do k=1,N-1
                  dev(itr+1,k,alpha) = sum(chi(:,alpha)*rho(:,k+1))&!
                       + sum(omega(:,k+1))/N_monomer - omega(alpha,k+1)
               end do
               
            elseif(N_monomer==3)then
               do k=1,N-1
                  dev(itr+1,k,alpha) = sum(chi(:,alpha)*rho(:,k+1))&!
                         + ( sum(omega(:,k+1)) + chi(1,2)*rho(3,k+1)&!
                         + chi(2,3)*rho(1,k+1) + chi(1,3)*rho(2,k+1))/N_monomer&!
                         - omega(alpha,k+1)
               end do
            endif
            omega_hist(itr+1,:,alpha) = omega(alpha,2:) 
            dev_2 = dev_2 + norm_2(dev(itr+1,:,alpha))**2.0
            omega_2 = omega_2 + norm_2(omega(alpha,2:))**2.0

         end do
         error   = (dev_2/omega_2)**0.5

      else

         dev_2   = 0.0
         omega_2 = 0.0
         do alpha=1,N_monomer

            dev(1:N_hist,:,alpha) = dev(2:N_hist+1,:,alpha)   
            omega_hist(1:N_hist,:,alpha) = omega_hist(2:N_hist+1,:,alpha)
            if(N_monomer==2)then
               do k=1,N-1
                  dev(N_hist+1,k,alpha) = sum(chi(:,alpha)*rho(:,k+1))&!
                         + sum(omega(:,k+1))/N_monomer - omega(alpha,k+1)
               end do

            elseif(N_monomer==3)then
               do k=1,N-1
                  dev(N_hist+1,k,alpha) = sum(chi(:,alpha)*rho(:,k+1))&!
                          + ( sum(omega(:,k+1)) + chi(1,2)*rho(3,k+1) &!
                          + chi(2,3)*rho(1,k+1) + chi(1,3)*rho(2,k+1) )/N_monomer&! 
                          - omega(alpha,k+1)
               end do
            endif
            omega_hist(N_hist+1,:,alpha) = omega(alpha,2:)
            dev_2 = dev_2 + norm_2(dev(N_hist+1,:,alpha))**2.0
            omega_2 = omega_2 + norm_2(omega(alpha,2:))**2.0

         end do
         error = (dev_2/omega_2)**0.5

      endif

      ! Calculate thermodynamic properties
      call mu_phi_chain(mu_chain, phi_chain, q)
      if (N_solvent > 0) then
         call mu_phi_solvent(mu_solvent, phi_solvent, q_solvent)
      endif
      call free_energy(N,rho,omega,phi_chain,mu_chain,phi_solvent,mu_solvent, &
                       f_Helmholtz,electrostatic_flag,eps,psi,pressure)

   end do iterative_loop


   ! AM scheme first resolve omega fields for a fixed unt cell dimenions 
   ! and then take one step of unit cell dimension. Then again resolves 
   ! the omega fields for updated unit cell dimensions and conitnue this
   ! to iterate fields and unit cell parameterssequentially  until both converges.
   if(domain) then

      if ( maxval(abs(stress)*stress_rescale) < error_max ) then
         converge = .true.
         write(*,*)'1'
         exit stress_loop
      else if ( itr_stress > 20 ) then
         converge = .false.
         write(*,*)'2'
         exit stress_loop
      else if ( maxval(abs(stress))*stress_rescale > large ) then
         converge = .false.
         write(*,*)'3'
         exit stress_loop
      endif

      write(6,*)'Stress relaxation iteration  = ',itr_stress
      
          call response_dstress_dcell(N,omega,dstress_dcell )
    
          if(N_cell_param==1)then
             p = -stress(1)/dstress_dcell(1,1)
          else
             call dgetrf(N_cell_param,N_cell_param,dstress_dcell,N_cell_param,ipvt,info)
             if(info/=0) stop "Full dstress_dcell LU factorization failed."
             call dgetri(N_cell_param,dstress_dcell,N_cell_param,ipvt,work,lwork,info)
             if(info/=0) stop "Full dstress_dcell inversion failed."
            
             p = - matmul(dstress_dcell,stress)
          endif
    
          vsum=0.0_long
          vsum = vsum + dot_product(cell_param(1:N_cell_param),&!
                                      cell_param(1:N_cell_param))
          stpmax = STPMX * max( sqrt(vsum), dble(N_cell_param) )
    
          vsum = sqrt(dot_product(p,p))
          if (vsum > stpmax) then
             write(6,*) "residual rescaled"
             p = p * stpmax / vsum
          end if
          
          do i=1,N_cell_param
             cell_param(i) = cell_param(i) + p(i)
          end do

          call make_unit_cell
          call make_ksq(G_basis)
          do i = 1, N_cell_param
             call make_dGsq(dGsq(:,i), dGG_basis(:,:,i))
          end do

          stress = scf_stress(N, N_cell_param, dGsq)
          
          write(6,*)'Updated Unit Cell dimensions = ',cell_param(1:N_cell_param)
          write(6,*)'stress = ',stress
          write(6,*)''
     
!!      end do stress_relax

      call density(N, omega, rho, q, q_solvent)

      if(itr<N_hist+1)then

         dev_2   = 0.0
         omega_2 = 0.0
         do alpha=1,N_monomer

            dev(1:2,:,alpha) = dev(itr-1:itr,:,alpha)   
            omega_hist(1:2,:,alpha) = omega_hist(itr-1:itr,:,alpha)
            
            if(N_monomer==2)then 
               do k=1,N-1
                  dev(3,k,alpha) = sum(chi(:,alpha)*rho(:,k+1))&!
                        + sum(omega(:,k+1))/N_monomer - omega(alpha,k+1)
               end do

            elseif(N_monomer==3)then
               do k=1,N-1
                  dev(3,k,alpha) = sum(chi(:,alpha)*rho(:,k+1))&!
                          + ( sum(omega(:,k+1)) + chi(1,2)*rho(3,k+1) &!
                          + chi(2,3)*rho(1,k+1) + chi(1,3)*rho(2,k+1) )/N_monomer &!
                          - omega(alpha,k+1)
               end do
            endif
            omega_hist(3,:,alpha) = omega(alpha,2:)
            dev_2 = dev_2 + norm_2(dev(3,:,alpha))**2.0
            omega_2 = omega_2 + norm_2(omega(alpha,2:))**2.0

         end do
         error   = (dev_2/omega_2)**0.5

         itr = 2  ! Number of previous iterations becuase current iteration will
                  ! be updated at the begining of itr loop, itr = itr + 1

      else

         dev_2   = 0.0
         omega_2 = 0.0
         do alpha=1,N_monomer

            dev(1:N_hist,:,alpha) = dev(2:N_hist+1,:,alpha)   
            omega_hist(1:N_hist,:,alpha) = omega_hist(2:N_hist+1,:,alpha)
            
            if(N_monomer==2)then 
               do k=1,N-1
                  dev(N_hist+1,k,alpha) = sum(chi(:,alpha)*rho(:,k+1))&!
                        + sum(omega(:,k+1))/N_monomer - omega(alpha,k+1)
               end do

            elseif(N_monomer==3)then
               do k=1,N-1
                  dev(N_hist+1,k,alpha) = sum(chi(:,alpha)*rho(:,k+1))&!
                          + ( sum(omega(:,k+1)) + chi(1,2)*rho(3,k+1) &!
                          + chi(2,3)*rho(1,k+1) + chi(1,3)*rho(2,k+1) )/N_monomer &!
                          - omega(alpha,k+1)
               end do
            endif
            omega_hist(N_hist+1,:,alpha) = omega(alpha,2:)
            dev_2 = dev_2 + norm_2(dev(N_hist+1,:,alpha))**2.0
            omega_2 = omega_2 + norm_2(omega(alpha,2:))**2.0

         end do
         error   = (dev_2/omega_2)**0.5

         itr = N_hist

     endif

     if(allocated(Umn)) deallocate(Umn)
     if(allocated(Vm))  deallocate(Vm)
     if(allocated(Cn))  deallocate(Cn)


     ! Calculate thermodynamic properties
     call free_energy(N,rho,omega,phi_chain,mu_chain,phi_solvent,mu_solvent, &
                      f_Helmholtz,electrostatic_flag,eps,psi,pressure)
  
     stress = scf_stress(N, N_cell_param, dGsq)

   else
      
      exit stress_loop

   endif

   end do stress_loop

   deallocate(dev)
   deallocate(omega_hist)
   ! If fixed unit cell, calculate residual stress before output
   if (.not.domain) stress = scf_stress(N, N_cell_param, dGsq)

contains

      function norm_2(x)
          
         implicit none
         
         intrinsic              :: dot_product, sqrt
         real(long)             :: norm_2
         real(long), intent(IN) :: x(:)
   
         norm_2 = sqrt(dot_product(x,x))
   
      end function norm_2
        

   end subroutine iterate_AM





end module iterate_mod
