!===============================================================================
! MODULE: scf_addons
!===============================================================================
!
! DESCRIPTION:
!   The scf_addons module provides specialized functionality to enhance SCF
!   convergence for challenging electronic systems. It implements three major
!   techniques:
!       - pseudo-Fractional Occupation Numbers (pFON),
!       - Maximum Overlap Method (MOM), and
!       - level shifting.
!   These methods help with systems exhibiting near-degeneracies, state flipping,
!   or convergence difficulties.
!
! MEMBERS:
!   - pfon_t [TYPE]: Encapsulates pFON functionality for managing fractional
!                    occupations based on temperature-dependent Fermi-Dirac
!                    distributions.
!
! DEPENDENCIES:
!   - precision: Provides `dp` for double precision real numbers.
!   - io_constants: Provides `iw` for output unit.
!   - mathlib: For matrix operations including pack_matrix, unpack_matrix.
!   - messages: For error handling.
!
! PUBLIC INTERFACES:
!   - pfon_t: Type for managing pseudo-Fractional Occupation Numbers.
!   - apply_mom: Implements Maximum Overlap Method for orbital tracking.
!   - level_shift_fock: Applies level shifting to the Fock matrix.
!
! NOTES:
!   - The module's functionality is designed to be used within SCF iterations.
!   - Methods work with RHF, UHF, and ROHF wavefunctions.
!   - The pFON implementation follows the approach described in:
!     https://doi.org/10.1063/1.478177
!
! HISTORY:
!   - [2025] Initial Module Creation - Konstantin Komarov
!     Established this module by extracting and refactoring auxiliary SCF
!     functionality from the main `scf`` module for better code organization
!     and maintainability. Implemented the `pfon_t`` type as a proper object
!     to encapsulate the pFON functionality.
!   - [January 2025] pFON Implementation - Alireza Lashkaripour
!     Developed the pseudo-Fractional Occupation Number (pFON) method
!     functionality that was later integrated into this module.
!   - [2023-2025] Advanced Convergence Methods - Konstantin Komarov
!     Implemented the Maximum Overlap Method (MOM) and level shifting
!     techniques to improve convergence for challenging electronic systems.
!
!===============================================================================

!===============================================================================
! TYPE: pfon_t - PSEUDO-FRACTIONAL OCCUPATION NUMBERS
!===============================================================================
!
! DESCRIPTION:
!   The `pfon_t` type encapsulates functionality for managing fractional
!   occupation numbers in SCF calculations using a temperature-dependent
!   Fermi-Dirac distribution. This technique smooths convergence for systems
!   with near-degeneracies by allowing partial orbital occupations.
!
! MEMBERS:
!   active         [LOGICAL]: Whether pFON is currently enabled.
!   temp           [REAL(dp)]: Current temperature for Fermi-Dirac distribution.
!   beta           [REAL(dp)]: Inverse temperature parameter (1/(kB*T)).
!   last_cooled_temp [REAL(dp)]: Last temperature at which cooling occurred.
!   cooling_rate   [REAL(dp)]: Rate of temperature decrease per iteration.
!   nsmear         [INTEGER]: Number of orbitals to smear around the Fermi level.
!   occ_a          [REAL(dp), POINTER]: Alpha orbital occupations array.
!   occ_b          [REAL(dp), POINTER]: Beta orbital occupations array.
!   scf_type       [INTEGER]: SCF calculation type (1=RHF, 2=UHF, 3=ROHF).
!   nelec          [INTEGER]: Total number of electrons.
!   nelec_a        [INTEGER]: Number of alpha electrons.
!   nelec_b        [INTEGER]: Number of beta electrons.
!   nbf            [INTEGER]: Number of basis functions.
!
! METHODS:
!   init                 - Initializes pFON parameters based on control settings.
!   adjust_temperature   - Dynamically adjusts temperature during iterations.
!   compute_occupations  - Wrapper for helper `pfon_occupations` function.
!                          Calculates fractional occupations from orbital energies.
!   build_density        - Wrapper for helper `build_pfon_density` function.
!                          Constructs density matrices using fractional occupations.
!
! HELPER FUNCTIONS:
!   pfon_occupations     - Standalone function that computes fractional occupations
!   build_pfon_density   - Constructs density matrices from MO coefficients and
!                          fractional occupations.
!
! ALGORITHM:
!   1. Start with high temperature (typically 2000K) to allow significant
!      fractional occupation and smooth energy surface
!   2. Gradually decrease temperature during iterations (cooling_rate parameter)
!   3. Compute Fermi level as average of HOMO and LUMO energies
!   4. Calculate occupations using Fermi-Dirac distribution:
!      n_i = 2/(1+exp((ε_i-εF)/kT)) for RHF
!      n_i = 1/(1+exp((ε_i-εF)/kT)) for UHF/ROHF
!   5. Normalize occupations to preserve total electron count
!   6. Use occupations to build weighted density matrices
!   7. Set temperature to 1K for final iteration to obtain integer occupations
!
! USAGE NOTES:
!   - Temperature gradually decreases during SCF iterations to facilitate convergence
!   - Final iteration typically uses T=0K to obtain integer occupations
!   - Works with all SCF types (RHF, UHF, ROHF) with appropriate occupation patterns
!   - Particularly effective for systems with small HOMO-LUMO gaps or
!     near-degenerate orbital energies
!
!===============================================================================

!===============================================================================
! TYPE: pfon_t - PSEUDO-FRACTIONAL OCCUPATION NUMBERS
!===============================================================================
!
! DESCRIPTION:
!   The `pfon_t` type encapsulates functionality for managing fractional
!   occupation numbers in SCF calculations using a temperature-dependent
!   Fermi-Dirac distribution. This technique smooths convergence for systems
!   with near-degeneracies by allowing partial orbital occupations.
!
! MEMBERS:
!   active         [LOGICAL]: Whether pFON is currently enabled.
!   temp           [REAL(dp)]: Current temperature for Fermi-Dirac distribution.
!   beta           [REAL(dp)]: Inverse temperature parameter (1/(kB*T)).
!   last_cooled_temp [REAL(dp)]: Last temperature at which cooling occurred.
!   cooling_rate   [REAL(dp)]: Rate of temperature decrease per iteration.
!   nsmear         [INTEGER]: Number of orbitals to smear around the Fermi level.
!   occ_a          [REAL(dp), POINTER]: Alpha orbital occupations array.
!   occ_b          [REAL(dp), POINTER]: Beta orbital occupations array.
!   scf_type       [INTEGER]: SCF calculation type (1=RHF, 2=UHF, 3=ROHF).
!   nelec          [INTEGER]: Total number of electrons.
!   nelec_a        [INTEGER]: Number of alpha electrons.
!   nelec_b        [INTEGER]: Number of beta electrons.
!   nbf            [INTEGER]: Number of basis functions.
!
! METHODS:
!   init                 - Initializes pFON parameters based on control settings.
!   adjust_temperature   - Dynamically adjusts temperature during iterations.
!   compute_occupations  - Calculates fractional occupations from orbital energies.
!   build_density        - Constructs density matrices using fractional occupations.
!
! NOTES:
!   - Works with RHF, UHF, ROHF with appropriate occupation patterns.
!   - Works with Second-Order SCF convergence method.
!
!===============================================================================

!===============================================================================
! SUBROUTINE: apply_mom - MAXIMUM OVERLAP METHOD
!===============================================================================
!
! DESCRIPTION:
!   Implements the Maximum Overlap Method (MOM) to maintain consistent orbital
!   ordering between SCF iterations. This helps prevent oscillations and state
!   flipping during convergence, especially for open-shell systems or cases
!   with near-degeneracies.
!
! PARAMETERS:
!   infos         [TYPE(information)]: System information.
!   v_prev        [REAL(dp)]: Previous iteration's MO coefficients.
!   e_prev        [REAL(dp)]: Previous iteration's orbital energies.
!   v_curr        [REAL(dp)]: Current iteration's MO coefficients (reordered on output).
!   e_curr        [REAL(dp)]: Current iteration's orbital energies (reordered on output).
!   s_ao          [REAL(dp)]: Overlap matrix in AO basis.
!   n_occ         [INTEGER]: Number of occupied orbitals.
!   spin_label    [CHARACTER(*)]: Identifier for spin channel ("Alpha" or "Beta").
!   work          [REAL(dp)]: Work array for intermediate calculations.
!   s_mo          [REAL(dp)]: Work array for MO overlap matrix.
!
! ALGORITHM:
!   1. Computes overlap between previous and current MOs: S_MO = V_prev^T * S * V_curr
!   2. For each orbital (occupied+1), finds maximum overlap match
!   3. Reorders current orbitals to maximize consistency with previous iteration
!   4. Ensures proper tracking of HOMO/LUMO and other important orbitals
!
! HELPER ROUTINES:
!   reorder_orbitals - Internal subroutine that performs the actual orbital swapping
!
!===============================================================================

!===============================================================================
! SUBROUTINE: level_shift_fock - VIRTUAL ORBITAL SHIFTING
!===============================================================================
!
! DESCRIPTION:
!   Applies level shifting to the virtual orbitals in the Fock matrix to increase
!   the HOMO-LUMO gap and improve SCF convergence. This technique is particularly
!   useful for systems with small HOMO-LUMO gaps or near-degeneracies.
!
! PARAMETERS:
!   fock_ao       [REAL(dp)]: Fock matrix in AO basis (triangular format).
!   mo_coefs      [REAL(dp)]: MO coefficients.
!   smat_full     [REAL(dp)]: Full overlap matrix.
!   nocc          [INTEGER]: Number of occupied orbitals.
!   nbf           [INTEGER]: Number of basis functions.
!   vshift        [REAL(dp)]: Level shift parameter value.
!   work1, work2  [REAL(dp)]: Work arrays for intermediate calculations.
!
! ALGORITHM:
!   1. Transforms Fock from AO to MO basis: F_MO = C^T * F_AO * C
!   2. Adds shift to diagonal elements corresponding to virtual orbitals
!   3. Transforms modified Fock back to AO basis for use in SCF
!
! USAGE NOTES:
!   - Typically applied in early iterations and gradually reduced
!   - Often combined with DIIS for optimal convergence
!
! NOTES:
!   - Works with RHF and UHF calculations. The ROHF case is handled through
!   the `form_rohf_fock` function in the `scf` module.
!
!===============================================================================
module scf_addons
  use precision, only: dp

  character(len=*), parameter :: module_name = "scf_addons"

  private

  public :: pfon_t
  public :: apply_mom
  public :: level_shift_fock
  public :: fock_jk
  public :: calc_fock
  public :: scf_energy_t
  public :: get_solver_name
  public :: compute_energy
  public :: calc_jk_xc
  public :: get_response_packed
  public :: get_scf_name
  integer, parameter, public :: scf_rhf  = 1  ! Restricted HF
  integer, parameter, public :: scf_uhf  = 2  ! Unrestricted HF
  integer, parameter, public :: scf_rohf = 3  ! ROHF
  integer, parameter, public :: scf_diis = 0, scf_bfgs = 1, scf_trah = 2

  !> @brief Type to encapsulate pFON (pseudo-Fractional Occupation Number) functionality
  !> @detail Provides methods for managing fractional occupation numbers in SCF calculations,
  !>         including temperature control, occupation computation, and density building.
  type :: pfon_t
    private
    logical :: active = .false.                   ! Whether pFON is enabled
    real(kind=dp), public :: temp                 ! Current temperature
    real(kind=dp), public :: beta                 ! Inverse temperature (1/(kB * temp))
    real(kind=dp) :: last_cooled_temp = 0.0_dp    ! Last temperature at which cooling occurred
    real(kind=dp) :: cooling_rate = 50.0_dp       ! Temperature cooling rate
    integer :: nsmear = 0                         ! Number of smearing steps
    real(kind=dp), pointer, public :: occ_a(:) => null()  ! Alpha occupations
    real(kind=dp), pointer, public :: occ_b(:) => null()  ! Beta occupations
    integer :: scf_type = 1  ! SCF calculation type (1=RHF, 2=UHF, 3=ROHF)
    integer :: nelec = 0     ! Total number of electrons
    integer :: nelec_a = 0   ! Number of alpha electrons
    integer :: nelec_b = 0   ! Number of beta electrons
    integer :: nbf = 0       ! Number of basis functions
  contains
    procedure :: init => pfon_init
    procedure :: adjust_temperature => pfon_adjust_temperature
    procedure :: compute_occupations => pfon_compute_occupations
    procedure :: build_density => pfon_build_density
  end type pfon_t

  type :: scf_energy_t
    real(kind=dp) :: ehf      ! Electronic energy (HF part)
    real(kind=dp) :: ehf1     ! One-electron energy
    real(kind=dp) :: nenergy  ! Nuclear repulsion energy
    real(kind=dp) :: etot     ! Total SCF energy
    real(kind=dp) :: e_old    ! Energy from previous iteration
    real(kind=dp) :: psinrm   ! Wavefunction normalization
    real(kind=dp) :: vne      ! Nucleus-electron potential energy
    real(kind=dp) :: vnn      ! Nucleus-nucleus potential energy
    real(kind=dp) :: vee      ! Electron-electron potential energy
    real(kind=dp) :: vtot     ! Total potential energy
    real(kind=dp) :: virial   ! Virial ratio (V/T)
    real(kind=dp) :: tkin     ! Kinetic energy
    real(kind=dp) :: eexc     ! Exchange-correlation energy for DFT
    real(kind=dp) :: totele   ! Total electron density for DFT
    real(kind=dp) :: totkin   ! Total kinetic energy for DFT
    real(kind=dp) :: e_pcm = 0.0_dp ! PCM solvent reaction-field energy (provisional; ddX path)
  contains
    procedure :: print_e => print_scf_energy
  end type scf_energy_t

contains
  function get_solver_name(solver_id) result(name)
    implicit none
    integer, intent(in) :: solver_id
    character(len=16)   :: name

    select case(solver_id)
    case (scf_diis)
       name = 'DIIS'
    case (scf_bfgs)
       name = 'BFGS/SOSCF'
    case (scf_trah)
       name = 'TRAH'
    case default
       name = 'UNKNOWN'
    end select
  end function get_solver_name

  pure function get_scf_name(code) result(name)
    integer, intent(in) :: code
    character(len=:), allocatable :: name
    select case (code)
    case (scf_rhf);  name = 'RHF'
    case (scf_uhf);  name = 'UHF'
    case (scf_rohf); name = 'ROHF'
    case default;    name = 'UNKNOWN'
    end select
  end function get_scf_name

  !> @brief Prints the final energy components of the SCF calculation.
  !> @detail Outputs a detailed breakdown of energy terms, including one-electron,
  !>         two-electron, nuclear repulsion, and total energies, as well as potential
  !>         (electron-electron, nucleus-electron, nucleus-nucleus, total) and
  !>         kinetic contributions, and the virial ratio.
  subroutine print_scf_energy(this)
     use precision, only: dp
     use io_constants, only: iw
     implicit none
     class(scf_energy_t), intent(in) :: this
     write(IW,"(/10X,17('=')/10X,'Energy components'/10X,17('=')/)")
     write(IW,"('         Wavefunction normalization =',F19.10)") this%psinrm
     write(IW,*)
     write(IW,"('                One electron energy =',F19.10)") this%ehf1
     write(IW,"('                Two electron energy =',F19.10)") this%vee
     write(IW,"('           Nuclear repulsion energy =',F19.10)") this%nenergy
     if (this%e_pcm /= 0.0_dp) then
        write(IW,"('           PCM solvent energy        =',F19.10)") this%e_pcm
     end if
     write(IW,"(38X,18('-'))")
     write(IW,"('                       TOTAL energy =',F19.10)") this%etot
     write(IW,*)
     write(IW,"(' Electron-electron potential energy =',F19.10)") this%vee
     write(IW,"('  Nucleus-electron potential energy =',F19.10)") this%vne
     write(IW,"('   Nucleus-nucleus potential energy =',F19.10)") this%vnn
     write(IW,"(38X,18('-'))")
     write(IW,"('             TOTAL potential energy =',F19.10)") this%vtot
     write(IW,"('               TOTAL kinetic energy =',F19.10)") this%tkin
     write(IW,"('                 Virial ratio (V/T) =',F19.10)") this%virial
     write(IW,*)
  end subroutine print_scf_energy

  !> @brief Applies the Maximum Overlap Method (MOM) to reorder orbitals.
  !> @detail Reorders the current iteration’s orbitals to maximize overlap
  !>         with the previous iteration’s orbitals,
  !>         ensuring consistent electronic state tracking during SCF convergence
  !>         (useful for avoiding state flipping).
  !> @param[in] infos System information.
  !> @param[in] v_prev Previous iteration’s MO coefficients.
  !> @param[in] e_prev Previous iteration’s orbital energies.
  !> @param[inout] v_curr Current iteration’s MO coefficients (reordered on output).
  !> @param[inout] e_curr Current iteration’s orbital energies (reordered on output).
  !> @param[in] s_ao Overlap matrix in AO basis.
  !> @param[in] n_occ Number of occupied orbitals.
  !> @param[in] spin_label Identifier for spin channel ("Alpha" or "Beta").
  !> @param[inout] work Work array for intermediate calculations (nbf x nbf).
  !> @param[inout] s_mo Work array for MO overlap matrix (nbf x nbf).
  subroutine apply_mom(infos, v_prev, e_prev, v_curr, e_curr, s_ao, n_occ, &
                       spin_label, work, s_mo)
    use precision, only: dp
    use io_constants, only: iw
    use types, only: information

    implicit none

    ! Input/output parameters
    type(information), intent(in) :: infos
    real(kind=dp), intent(in),    dimension(:,:) :: v_prev
    real(kind=dp), intent(in),    dimension(:)   :: e_prev
    real(kind=dp), intent(inout), dimension(:,:) :: v_curr
    real(kind=dp), intent(inout), dimension(:)   :: e_curr
    real(kind=dp), intent(in),    dimension(:,:) :: s_ao
    integer,       intent(in)                    :: n_occ
    character(*),  intent(in)                    :: spin_label
    real(kind=dp), intent(inout), dimension(:,:) :: work
    real(kind=dp), intent(inout), dimension(:,:) :: s_mo

    ! Local variables
    integer :: i, j, k, ip1, nbf
    integer :: max_idx
    real(kind=dp) :: max_overlap, overlap
    logical, allocatable :: reordered(:)

    nbf = size(v_curr, 1)

    if (infos%control%verbose>=1) then
      if (infos%control%rstctmo) then
        write(IW, fmt='(/,"Applying Reodering for ",A," spin channel")') trim(spin_label)
      else
        write(IW, fmt='(/,"Applying MOM for ",A," spin channel")') trim(spin_label)
      end if
    end if

    ! Allocate reordered flag array
    allocate(reordered(nbf), source=.false.)

    ! Calculate overlap between previous and current MOs: s_mo = v_prev^T * s_ao * v_curr
    call dgemm('t', 'n', nbf, nbf, nbf, 1.0_dp, v_prev, nbf, s_ao, nbf, 0.0_dp, work, nbf)
    call dgemm('n', 'n', nbf, nbf, nbf, 1.0_dp, work, nbf, v_curr, nbf, 0.0_dp, s_mo, nbf)

    ! Normalize columns to ensure proper comparison
    do i = 1, nbf
      s_mo(:,i) = s_mo(:,i) / max(norm2(s_mo(:,i)), 1.0e-10_dp)
    end do

    ! First, identify the best match for each orbital from the previous iteration
    ! Focus particularly on occupied orbitals and the HOMO-LUMO region
    ! Print information about important orbitals (HOMO, LUMO)
    if (infos%control%verbose>1) then
      write(IW,fmt='(1X,"MOM reordering for ",A," orbitals:")') trim(spin_label)
      write(IW,fmt='(1X,"Old Index → New Index   | Overlap |  Status")')
      write(IW,fmt='(1X,"--------------------------------------------")')
    end if

    ! First pass: check which orbitals need reordering
    do i = 1, nbf
      max_overlap = 0.0_dp
      max_idx = i  ! Default to no change

      ! Find the orbital with maximum overlap
      do j = 1, nbf
        if (.not. reordered(j)) then
          overlap = abs(s_mo(i,j))
          if (overlap > max_overlap) then
            max_overlap = overlap
            max_idx = j
          end if
        end if
      end do

      ! Mark the orbital as reordered and print info for occupied orbitals
      reordered(max_idx) = .true.
      if (infos%control%verbose>1) then
        ! Print info for important orbitals or those being reordered
        if (((i <= n_occ+1) .or. (i /= max_idx)).and. infos%control%verbose>=1) then
          write(IW, fmt='(3X,I3,5X,"→",5X,I3,5X,"| ",F7.5," |")', advance='no') &
            i, max_idx, max_overlap

          ! Add label for HOMO/LUMO
          if (i == n_occ)   write(IW, fmt='(1X,"HOMO")', advance='no')
          if (i == n_occ+1) write(IW, fmt='(1X,"LUMO")', advance='no')

          ! Add status message
          if (i /= max_idx .and. max_overlap < 0.9_dp) then
            write(IW, fmt='(1X,"Reordered (warning: low overlap)")')
          else if (i /= max_idx) then
            write(IW, fmt='(1X,"Reordered")')
          else if (max_overlap < 0.9_dp) then
            write(IW, fmt='(1X,"Unchanged (warning: low overlap)")')
          else
            write(IW, fmt='(1X,"Unchanged")')
          end if
        end if
      end if
    end do

    ! Check if all orbitals were successfully assigned
    if (.not. all(reordered)) then
      write(IW, fmt='(/,"WARNING: Some orbitals could not be properly reordered!")')
      write(IW, fmt='("This may indicate a significant change in electronic structure.")')
    end if

    ! Apply the reordering
    call reorder_orbitals(v_curr, e_curr, s_mo, nbf, &
                          start_mo=1, &
                          end_mo=n_occ+1)

    deallocate(reordered)
  end subroutine apply_mom

  !> @brief Reorders orbitals based on overlap with the previous iteration.
  !> @detail Internal helper routine for 'apply_mom' that swaps orbital
  !>         coefficients and energies to maximize overlap,
  !>         focusing on a specified range of molecular orbitals.
  !> @param[inout] v MO coefficients (reordered on output).
  !> @param[inout] e Orbital energies (reordered on output).
  !> @param[in] smo Overlap matrix between previous and current MOs.
  !> @param[in] nbf Number of basis functions.
  !> @param[in] start_mo First MO to reorder.
  !> @param[in] end_mo Last MO to reorder.
  subroutine reorder_orbitals(v, e, smo, nbf, start_mo, end_mo)
    use precision, only: dp

    implicit none

    real(kind=dp), intent(inout) :: v(nbf,*)
    real(kind=dp), intent(inout) :: e(*)
    real(kind=dp), intent(in) :: smo(nbf,*)
    integer, intent(in) :: nbf, start_mo, end_mo

    integer :: i, j, k, ip1
    integer, allocatable :: reorder_idx(:)
    real(kind=dp) :: smax, tmp_e

    ! Allocate array for reordering indices
    allocate(reorder_idx(nbf), source=0)

    ! Determine the reordering indices based on maximum overlap
    do i = 1, nbf
      smax = 0.0_dp
      reorder_idx(i) = 0

      ! Find maximum overlap
      do j = 1, nbf
        ! Skip already assigned orbitals
        if (any(reorder_idx(1:i-1) == j)) cycle

        if (abs(smo(i,j)) > smax) then
          smax = abs(smo(i,j))
          reorder_idx(i) = j
        end if
      end do

      ! Ensure sign consistency
      if (smo(i, reorder_idx(i)) < 0.0_dp) then
        v(:, reorder_idx(i)) = -v(:, reorder_idx(i))
      end if
    end do

    ! Apply reordering for the specified range
    do i = start_mo, end_mo
      j = reorder_idx(i)

      ! Swap orbital coefficients
      call dswap(nbf, v(1,i), 1, v(1,j), 1)

      ! Swap orbital energies
      tmp_e = e(i)
      e(i) = e(j)
      e(j) = tmp_e

      ! Update reordering indices for remaining swaps
      ip1 = i + 1
      do k = ip1, end_mo
        if (reorder_idx(k) == i) reorder_idx(k) = j
      end do
    end do

    deallocate(reorder_idx)
  end subroutine reorder_orbitals

  !> @brief Computes fractional occupation numbers using
  !>        the pseudo-Fractional Occupation Number (pFON) method.
  !> @detail Implements the pFON method to assign fractional occupations
  !>         via a Fermi-Dirac distribution, smoothing near-degenerate states.
  !>         Reference: https://doi.org/10.1063/1.478177
  !> @author Alireza Lashkaripour, January 2025
  !> @param[in] mo_energy Orbital energies.
  !> @param[in] nbf Number of basis functions.
  !> @param[in] nelec Total number of electrons.
  !> @param[inout] occ Occupation numbers (updated on output).
  !> @param[in] beta_pfon Inverse temperature parameter (1/(kB * T)).
  !> @param[in] scf_type SCF type (1=RHF, 2=UHF, 3=ROHF).
  !> @param[in] nsmear Number of orbitals to smear around the Fermi level.
  !> @param[in] is_beta Flag indicating beta spin calculation (optional).
  !> @param[in] nelec_a Number of alpha electrons (for UHF/ROHF).
  !> @param[in] nelec_b Number of beta electrons (for UHF/ROHF).
  subroutine pfon_occupations(mo_energy, nbf, nelec, occ, beta_pfon, &
                              scf_type, nsmear, is_beta, nelec_a, nelec_b)
    use precision, only: dp
    implicit none

    integer, intent(in) :: nbf
    integer, intent(in) :: nelec, nsmear
    real(kind=dp), intent(in) :: beta_pfon
    real(kind=dp), intent(in) :: mo_energy(nbf)
    real(kind=dp), intent(inout) :: occ(nbf)
    integer, intent(in) :: scf_type ! 1,2,3 RHF,UHF,ROHF
    logical, intent(in), optional :: is_beta
    integer, intent(in), optional :: nelec_a, nelec_b
    real(kind=dp) :: eF, sum_occ
    integer :: i, i_homo, i_lumo, i_low, i_high
    real(kind=dp) :: tmp
    logical :: is_beta_calc
    integer :: n_electrons, n_double, n_single

    is_beta_calc = .false.
    if (present(is_beta)) is_beta_calc = is_beta

    select case (scf_type)
    case(1) ! RHF
      i_homo = max(1, nelec/2)
      n_electrons = nelec

    case(2) ! UHF
      if (.not. present(nelec_a) .or. .not. present(nelec_b)) then
        stop 'UHF requires nelec_a and nelec_b'
      end if
      ! UHF: completely independent alpha and beta
      if (is_beta_calc) then
        i_homo = max(1, nelec_b)
        n_electrons = nelec_b
      else
        i_homo = max(1, nelec_a)
        n_electrons = nelec_a
      end if

    case(3) ! ROHF
      if (.not. present(nelec_a) .or. .not. present(nelec_b)) then
        stop 'ROHF requires nelec_a and nelec_b'
      end if
      ! ROHF: same spatial orbitals, different occupations
      n_double = nelec_b
      n_single = nelec_a - nelec_b
      if (is_beta_calc) then
        i_homo = n_double
        n_electrons = nelec_b
      else
        i_homo = n_double + n_single
        n_electrons = nelec_a
      end if
    end select

    i_lumo = i_homo + 1
    if (i_lumo > nbf) i_lumo = nbf

    ! Calculate Fermi level
    eF = 0.5_dp * (mo_energy(i_homo) + mo_energy(i_lumo))

    if (nsmear <= 0) then
      do i = 1, nbf
        tmp = beta_pfon * (mo_energy(i) - eF)
        if (scf_type == 1) then  ! RHF
          occ(i) = 2.0_dp / (1.0_dp + exp(tmp))
        else  ! UHF or ROHF
          occ(i) = 1.0_dp / (1.0_dp + exp(tmp))
        end if
      end do
    else
      i_low = max(1, i_homo - nsmear)
      i_high = min(nbf, i_lumo + nsmear)

      ! Special handling for ROHF
      if (scf_type == 3) then
        if (is_beta_calc) then
          do i = 1, n_double
            occ(i) = 1.0_dp
          end do
          do i = n_double + 1, nbf
            occ(i) = 0.0_dp
          end do
        else
          do i = 1, n_double
            occ(i) = 1.0_dp
          end do
          do i = n_double + 1, n_double + n_single
            occ(i) = 1.0_dp
          end do
          do i = n_double + n_single + 1, nbf
            occ(i) = 0.0_dp
          end do
        end if

        ! Apply smearing only around the Fermi level
        do i = i_low, i_high
          tmp = beta_pfon * (mo_energy(i) - eF)
          occ(i) = occ(i) / (1.0_dp + exp(tmp))
        end do
      else
        ! RHF/UHF handling
        do i = 1, i_low - 1
          if (scf_type == 1) then
            occ(i) = 2.0_dp
          else
            occ(i) = 1.0_dp
          end if
        end do

        do i = i_high + 1, nbf
          occ(i) = 0.0_dp
        end do

        do i = i_low, i_high
          tmp = beta_pfon * (mo_energy(i) - eF)
          if (scf_type == 1) then
            occ(i) = 2.0_dp / (1.0_dp + exp(tmp))
          else
            occ(i) = 1.0_dp / (1.0_dp + exp(tmp))
          end if
        end do
      end if
    end if

    ! Normalize occupations
    sum_occ = sum(occ(1:nbf))
    if (sum_occ < 1.0e-14_dp) then
      sum_occ = 1.0_dp
    end if
    occ(1:nbf) = occ(1:nbf) * (real(n_electrons,dp) / sum_occ)

  end subroutine pfon_occupations

  !> @brief Builds density matrices using fractional occupation numbers for the pFON method.
  !> @detail Constructs density matrices from molecular orbital coefficients
  !>         and fractional occupations.
  !> @param[inout] pdmat Density matrices (triangular format, updated on output).
  !> @param[in] mo_a Alpha MO coefficients.
  !> @param[in] mo_b Beta MO coefficients (UHF only).
  !> @param[in] occ_a Alpha occupation numbers.
  !> @param[in] occ_b Beta occupation numbers (UHF/ROHF).
  !> @param[in] scf_type SCF type (1=RHF, 2=UHF, 3=ROHF).
  !> @param[in] nbf Number of basis functions.
  !> @param[in] nelec_a Number of alpha electrons.
  !> @param[in] nelec_b Number of beta electrons.
  !> @param[inout] dtmp Work array for density matrix construction.
  !> @param[inout] work Additional work array.
  subroutine build_pfon_density(pdmat_a, mo_a, occ_a, scf_type, nbf, dtmp, work, &
                                pdmat_b, mo_b, occ_b, nelec_a, nelec_b)
    use precision, only: dp
    use mathlib, only: pack_matrix
    implicit none

    real(kind=dp), intent(inout) :: pdmat_a(:)
    real(kind=dp), intent(in) :: mo_a(:,:)
    real(kind=dp), intent(in) :: occ_a(:)
    integer, intent(in) :: nbf, scf_type
    real(kind=dp), intent(inout) :: dtmp(:,:), work(:,:)
    real(kind=dp), intent(inout), optional :: pdmat_b(:)
    real(kind=dp), intent(in), optional :: mo_b(:,:)
    real(kind=dp), intent(in), optional :: occ_b(:)
    integer, intent(in), optional :: nelec_a, nelec_b

    integer :: i, mu, nu
    integer :: n_double, n_single
    real(kind=dp) :: occ_factor


    select case(scf_type)
    case(1)  ! RHF
      ! Scale MO coefficients by square root of occupation numbers
      do i = 1, nbf
        if (occ_a(i) > 1.0e-14_dp) then
            call dger(nbf, nbf, occ_a(i), mo_a(:,i), 1, mo_a(:,i), 1, dtmp, nbf)
        end if
      end do
      pdmat_a = 0.0_dp
      call pack_matrix(dtmp, pdmat_a)

    case(2)  ! UHF
      do i = 1, nbf
        if (occ_a(i) > 1.0e-14_dp) then
          call dger(nbf, nbf, occ_a(i), mo_a(:,i), 1, mo_a(:,i), 1, dtmp, nbf)
        end if
      end do
      pdmat_a = 0.0_dp
      call pack_matrix(dtmp, pdmat_a)

      dtmp(:,:) = 0.0_dp
      do i = 1, nbf
        if (occ_b(i) > 1.0e-14_dp) then
          call dger(nbf, nbf, occ_b(i), mo_b(:,i), 1, mo_b(:,i), 1, dtmp, nbf)
        end if
      end do
      pdmat_b = 0.0_dp
      call pack_matrix(dtmp, pdmat_b)

    case(3)  ! ROHF
      n_double = nelec_b
      n_single = nelec_a - nelec_b

      dtmp(:,:) = 0.0_dp
      do i = 1, nbf
        if (occ_a(i) > 1.0e-14_dp) then
          if (i <= n_double) then
            occ_factor = occ_a(i)
          else if (i <= n_double + n_single) then
            occ_factor = 1.0_dp
          else
            occ_factor = occ_a(i)  ! Virtual orbitals
          end if

          ! dtmp += occ_factor * mo_a(:,i) * mo_a(:,i)^T
!         call dger(nbf, nbf, occ_factor, mo_a(:,i), 1, mo_a(:,i), 1, dtmp, nbf)
          do mu = 1, nbf
            do nu = 1, nbf
              dtmp(mu,nu) = dtmp(mu,nu) + occ_factor * mo_a(mu,i)*mo_a(nu,i)
            end do
          end do
        end if
      end do
      pdmat_a = 0.0_dp
      call pack_matrix(dtmp, pdmat_a)

      dtmp(:,:) = 0.0_dp
      do i = 1, nbf
        if (occ_b(i) > 1.0e-14_dp) then
          if (i <= n_double) then
            occ_factor = occ_b(i)
          else
            occ_factor = 0.0_dp
          end if

          ! dtmp += occ_factor * mo_a(:,i) * mo_a(:,i)^T
!         call dger(nbf, nbf, occ_factor, mo_a(:,i), 1, mo_a(:,i), 1, dtmp, nbf)
          do mu = 1, nbf
            do nu = 1, nbf
              dtmp(mu,nu) = dtmp(mu,nu) + occ_factor * mo_a(mu,i)*mo_a(nu,i)
            end do
          end do
        end if
      end do
      pdmat_b = 0.0_dp
      call pack_matrix(dtmp, pdmat_b)
    end select

  end subroutine build_pfon_density

  !> @brief Applies level shifting to the Fock matrix for improved SCF convergence.
  !> @detail Modifies the diagonal elements of the Fock matrix in the MO basis
  !>         for virtual orbitals by adding a shift parameter,
  !>         then transforms the result back to the AO basis.
  !> @param[inout] fock_ao Fock matrix in AO basis (triangular format, updated on output).
  !> @param[in] mo_coefs MO coefficients.
  !> @param[in] smat_full Full overlap matrix.
  !> @param[in] nocc Number of occupied orbitals.
  !> @param[in] nbf Number of basis functions.
  !> @param[in] vshift Level shift parameter value.
  subroutine level_shift_fock(fock_ao, mo_coefs, smat_full, nocc, nbf, vshift, &
                              work1, work2)
    use precision, only: dp
    use mathlib, only: orthogonal_transform_sym, &
                       orthogonal_transform2, &
                       unpack_matrix, &
                       pack_matrix

    implicit none

    integer, intent(in) :: nocc, nbf
    real(kind=dp), intent(inout) :: fock_ao(:)
    real(kind=dp), intent(in) :: mo_coefs(:,:)
    real(kind=dp), intent(in) :: smat_full(:,:)
    real(kind=dp), intent(in) :: vshift
    real(kind=dp), intent(inout) :: work1(:,:)
    real(kind=dp), intent(inout) :: work2(:,:)

    ! Local variables
    real(kind=dp), allocatable :: fock_mo_full(:,:), fock_mo(:), work_matrix(:,:)
    integer :: i, nbf_tri

    nbf_tri = nbf*(nbf+1)/2

    work1 = 0.0_dp
    work2 = 0.0_dp

    ! Allocate work arrays
    allocate(fock_mo_full(nbf, nbf), &
             fock_mo(nbf_tri), &
             work_matrix(nbf, nbf), &
             source=0.0_dp)

    ! Transform Fock from AO to MO basis: F_MO = C^T * F_AO * C
    call orthogonal_transform_sym(nbf, nbf, fock_ao, mo_coefs, nbf, fock_mo)

    ! Unpack triangular matrices to full format
    call unpack_matrix(fock_mo, fock_mo_full)

    ! Apply level shift to virtual orbitals in F_MO
    do i = nocc + 1, nbf
      fock_mo_full(i, i) = fock_mo_full(i, i) + vshift
    end do

    ! Back-transform ROHF Fock matrix to AO basis
    call dsymm('l', 'u', nbf, nbf, &
               1.0_dp, smat_full, nbf, &
                       mo_coefs, nbf, &
               0.0_dp, work1, nbf)
    call orthogonal_transform2('t', nbf, nbf, work1, nbf, fock_mo_full, nbf, &
                               work_matrix, nbf, work2)

    ! Pack the result back to triangular form
    call pack_matrix(work_matrix, fock_ao)

    deallocate(fock_mo_full, fock_mo, work_matrix)
  end subroutine level_shift_fock

  !> @brief Initialize pFON parameters
  !> @detail Sets up temperature, inverse temperature (beta),
  !>         and other pFON parameters based on input controls.
  !> @param[in] control Control structure containing pFON settings
  !> @param[in] nbf Number of basis functions
  !> @param[in] nelec Total number of electrons
  !> @param[in] nelec_a Number of alpha electrons
  !> @param[in] nelec_b Number of beta electrons
  !> @param[in] scf_type SCF type (1=RHF, 2=UHF, 3=ROHF)
  !> @param[inout] occ_a Pointer to alpha occupations array
  !> @param[inout] occ_b Pointer to beta occupations array (only for UHF/ROHF)
  subroutine pfon_init(this, control, nbf, nelec, nelec_a, nelec_b, scf_type, occ_a, occ_b)
    use types, only: control_parameters
    use constants, only: kB_HaK

    implicit none

    class(pfon_t), intent(inout) :: this
    type(control_parameters), intent(in) :: control
    integer, intent(in) :: nbf, nelec, nelec_a, nelec_b, scf_type
    real(dp), target, intent(inout) :: occ_a(:)
    real(dp), target, optional, intent(inout) :: occ_b(:)

    this%active = control%pfon
    if (.not. this%active) return

    this%nbf = nbf
    this%nelec = nelec
    this%nelec_a = nelec_a
    this%nelec_b = nelec_b
    this%scf_type = scf_type

    ! Set temperature parameters
    this%temp = control%pfon_start_temp
    if (this%temp <= 0.0_dp) this%temp = 2000.0_dp  ! Default temperature
    this%beta = 1.0_dp / (kB_HaK * this%temp)
    this%cooling_rate = control%pfon_cooling_rate
    if (this%cooling_rate <= 0.0_dp) this%cooling_rate = 50.0_dp

    ! Set number of orbitals to smear
    this%nsmear = int(control%pfon_nsmear)

    ! Set pointers to occupation arrays
    this%occ_a => occ_a
    if (present(occ_b)) this%occ_b => occ_b

  end subroutine pfon_init

  !> @brief Adjust pFON temperature based on convergence status
  !> @detail Dynamically modifies the temperature and beta parameters during SCF
  !>         iterations, reducing temperature as convergence improves.
  !> @param[in] iter Current SCF iteration
  !> @param[in] maxit Maximum number of SCF iterations
  !> @param[in] diis_error Current DIIS error
  !> @param[in] conv Convergence threshold
  subroutine pfon_adjust_temperature(this, iter, maxit, diis_error, conv, do_pfon ,do_final)
    use constants, only: kB_HaK
    use io_constants, only: iw
    class(pfon_t), intent(inout) :: this
    integer, intent(in) :: iter, maxit
    real(dp), intent(in) :: diis_error, conv
    logical, intent(in) :: do_final, do_pfon

    if (.not. do_pfon) return

    if (.not. this%active) return
    if (do_final) then
      this%temp = 1.0_dp
      this%beta = 1.0_dp / (kB_HaK * this%temp)
      write(IW, "(10x, 'Extra SCF iteration with Temp = 1K')")
    end if

    if (iter == maxit) then
      ! Final iteration: set temperature to zero for pure integer occupations
      this%temp = 0.0_dp
    else if (abs(diis_error) < 10.0_dp * conv) then
      ! Near convergence: set to minimum temperature (1K)
      if (this%temp > 1.0_dp) then
        this%last_cooled_temp = this%temp
      end if
      this%temp = 1.0_dp
    else
      ! Not converged yet: continue cooling temperature
      if (this%temp == 1.0_dp .and. this%last_cooled_temp > 1.0_dp) then
        this%temp = this%last_cooled_temp
      end if
      this%temp = this%temp - this%cooling_rate
      if (this%temp < 1.0_dp) then
        this%temp = 1.0_dp
      end if
      this%last_cooled_temp = this%temp
    end if

    ! Calculate beta = 1/(kB*T) for Fermi-Dirac distribution
    if (this%temp > 1.0e-12_dp) then
      this%beta = 1.0_dp / (kB_HaK * this%temp)
    else
      this%beta = 1.0e20_dp  ! Zero temperature
    end if
  end subroutine pfon_adjust_temperature

  !> @brief Compute fractional occupation numbers using pFON method
  !> @detail Uses current orbital energies to calculate occupations
  !>         via a Fermi-Dirac distribution.
  !> @param[in] mo_energy_a Alpha orbital energies
  !> @param[in] mo_energy_b Beta orbital energies (only for UHF)
  subroutine pfon_compute_occupations(this, mo_energy_a, do_pfon, mo_energy_b)
    class(pfon_t), intent(inout) :: this
    real(dp), intent(in) :: mo_energy_a(:)
    real(dp), intent(in), optional :: mo_energy_b(:)
    logical, intent(in) :: do_pfon

    if (.not. do_pfon) return

    if (.not. this%active) return

    ! Calculate alpha occupations
    call pfon_occupations(mo_energy_a, this%nbf, this%nelec, this%occ_a, &
                          this%beta, this%scf_type, this%nsmear, &
                          is_beta=.false., nelec_a=this%nelec_a, nelec_b=this%nelec_b)

    ! Calculate beta occupations if needed
    if (this%scf_type > 1 .and. associated(this%occ_b)) then
      if (this%scf_type == 2 .and. present(mo_energy_b)) then
        ! UHF case - use separate beta orbital energies
        call pfon_occupations(mo_energy_b, this%nbf, this%nelec, this%occ_b, &
                            this%beta, this%scf_type, this%nsmear, &
                            is_beta=.true., nelec_a=this%nelec_a, nelec_b=this%nelec_b)
      else
        ! ROHF case - use same orbital energies for alpha and beta
        call pfon_occupations(mo_energy_a, this%nbf, this%nelec, this%occ_b, &
                            this%beta, this%scf_type, this%nsmear, &
                            is_beta=.true., nelec_a=this%nelec_a, nelec_b=this%nelec_b)
      end if
    end if
  end subroutine pfon_compute_occupations

  !> @brief Build density matrices using fractional occupation numbers
  !> @detail Constructs density matrices for the current SCF iteration
  !>         using fractional occupations and MO coefficients.
  !> @param[inout] this pFON type instance.
  !> @param[out] pdmat Density matrices (triangular format).
  !> @param[in] mo_a Alpha MO coefficients.
  !> @param[inout] work1 Work array 1.
  !> @param[inout] work2 Work array 2.
  !> @param[in] mo_b Beta MO coefficients (optional for UHF).
  subroutine pfon_build_density(this, pdmat_a, mo_a, work1, work2, do_pfon, pdmat_b, mo_b)
    class(pfon_t), intent(inout) :: this
    real(kind=dp), intent(out) :: pdmat_a(:)
    real(kind=dp), intent(in) :: mo_a(:,:)
    real(kind=dp), intent(inout) :: work1(:,:)
    real(kind=dp), intent(inout) :: work2(:,:)
    real(kind=dp), intent(out), optional :: pdmat_b(:)
    real(kind=dp), intent(in), optional :: mo_b(:,:)
    logical , intent(in) :: do_pfon

    if (.not. do_pfon) return

    if (.not. this%active) return

    ! Nullify work arrays
    work1 = 0.0_dp
    work2 = 0.0_dp

    ! Call the existing build_pfon_density function with appropriate parameters
    select case (this%scf_type)
    case (1) ! RHF
      call build_pfon_density(pdmat_a, mo_a, this%occ_a, this%scf_type, this%nbf, &
                              work1, work2)
    case (2) ! UHF
      call build_pfon_density(pdmat_a, mo_a, this%occ_a, this%scf_type, this%nbf, &
                              work1, work2, pdmat_b, mo_b, this%occ_b)
    case (3) ! ROHF
      call build_pfon_density(pdmat_a, mo_a, this%occ_a, this%scf_type, this%nbf, &
                              work1, work2, pdmat_b, mo_b, this%occ_b, this%nelec_a, this%nelec_b)
    end select
  end subroutine pfon_build_density

  !> @brief Computes the two-electron part (Coulomb and exchange) of the Fock matrix.
  !> @detail Forms the Coulomb (J) and exchange (K) contributions to the Fock matrix
  !>         using two-electron integrals,
  !>         with optional scaling of the exchange term for hybrid DFT methods.
  !> @param[in] basis Basis set information.
  !> @param[in] d Density matrices (triangular format).
  !> @param[inout] f Fock matrices to be updated (triangular format).
  !> @param[in] scalefactor Optional scaling factor for exchange (default = 1.0).
  !> @param[inout] infos System information.
  subroutine fock_jk(basis, d, f, infos, scale_exch, nschwz, f_old, scale_coul, petite)
    use precision, only: dp
    use io_constants, only: iw
    use util, only: measure_time
    use basis_tools, only: basis_set
    use types, only: information
    use int2_compute, only: int2_compute_t, int2_fock_data_t, &
                            int2_rhf_data_t, int2_urohf_data_t

    implicit none

    type(basis_set), intent(in) :: basis
    type(information), intent(inout) :: infos
    real(kind=dp), optional, intent(in) :: scale_exch
    integer, optional, intent(inout) :: nschwz
    real(kind=dp), optional, intent(in) :: scale_coul
    real(kind=dp), target, intent(in) :: d(:,:)
    real(kind=dp), intent(inout) :: f(:,:)
    real(kind=dp), optional ,intent(inout) :: f_old(:,:)
    !> Opt into the symmetry petite-list reduction. Only valid for totally
    !> symmetric densities (SCF Fock); response/CPHF/Hessian callers with
    !> perturbed densities must not set this.
    logical, optional, intent(in) :: petite


    integer :: i, ii, nf
    real(kind=dp) :: scale_e, scale_c
    logical :: is_dft
    type(int2_compute_t) :: int2_driver
    class(int2_fock_data_t), allocatable :: int2_data

    ! Initial Settings
    scale_e = 1.0d0
    scale_c = 1.0d0
    if (present(scale_exch)) scale_e = scale_exch
    if (present(scale_coul)) scale_c = scale_coul
    is_dft = (infos%control%hamilton == 20)


    ! Initialize ERI calculations
    call int2_driver%init(basis, infos)
    if (present(petite)) then
      ! Petite-list reduction: only valid for totally symmetric densities
      ! (SCF Fock); the skeleton matrix is symmetrized below.
      if (petite) call int2_driver%enable_petite(infos)
    end if
    call int2_driver%set_screening()

    select case (infos%control%scftype)
    case (1)
      int2_data = int2_rhf_data_t(nfocks=1, d=d, scale_exchange=scale_e, scale_coulomb=scale_c)
    case (2)
      int2_data = int2_urohf_data_t(nfocks=2, d=d, scale_exchange=scale_e, scale_coulomb=scale_c)
    case (3)
      int2_data = int2_urohf_data_t(nfocks=2, d=d, scale_exchange=scale_e, scale_coulomb=scale_c)
    end select


    ! Constructing two electron Fock matrix
    call int2_driver%run(int2_data, &
                           cam=is_dft.and.infos%dft%cam_flag, &
                           alpha=infos%dft%cam_alpha, &
                           beta=infos%dft%cam_beta,&
                           mu=infos%dft%cam_mu)

    if (present(nschwz)) nschwz = int2_driver%skipped

    ! Scaling (everything except diagonal is halved)
    if (present(f_old)) then
      int2_data%f(:,:,1) = int2_data%f(:,:,1) + f_old
      f_old = int2_data%f(:,:,1)
    end if
    f =  0.5 * int2_data%f(:,:,1)
    do nf = 1, ubound(f,2)
      ii = 0
      do i = 1, basis%nbf
         ii = ii + i
         f(ii,nf) = 2*f(ii,nf)
      end do
    end do

    ! Petite-list runs produce a skeleton matrix; project onto the
    ! totally symmetric component: F <- (1/|G|) sum_op T_op F T_op^T.
    if (int2_driver%petite) call symmetrize_skeleton_fock(infos, basis, f)

    call int2_driver%clean()

  end subroutine fock_jk

!--------------------------------------------------------------------------------

!> @brief Symmetrize a packed-triangular skeleton Fock matrix.
!> @detail Applies F <- (1/|G|) sum_op T_op F T_op^T where T_op is the
!>   signed AO permutation of each abelian symmetry operation (standard
!>   orientation), using the maps written by pyoqp. No-op if the maps are
!>   missing.
  subroutine symmetrize_skeleton_fock(infos, basis, f)
    use precision, only: dp
    use types, only: information
    use basis_tools, only: basis_set
    use oqp_tagarray_driver
    use tagarray, only: TA_OK

    implicit none

    type(information), target, intent(inout) :: infos
    type(basis_set), intent(in) :: basis
    real(kind=dp), intent(inout) :: f(:,:)

    real(kind=dp), contiguous, pointer :: blocks(:)
    integer(4) :: status
    integer :: nbf

    nbf = basis%nbf
    call tagarray_get_data(infos%dat, OQP_sym_op_blocks, blocks, status=status)
    if (status == TA_OK) then
      call symmetrize_skeleton_blocked(infos, basis, f, blocks)
    else
      call symmetrize_skeleton_signed(infos, nbf, f)
    end if

  end subroutine symmetrize_skeleton_fock

!--------------------------------------------------------------------------------

!> @brief Full-group skeleton symmetrization with dense per-shell blocks.
!> @detail F <- (1/|G|) sum_op T_op F T_op^T where T_op permutes shells and
!>   mixes components within each shell (non-abelian operations such as the
!>   C6 rotations of D6h). Blocks staged by pyoqp, column-major per shell,
!>   concatenated shell-by-shell then op-by-op.
  subroutine symmetrize_skeleton_blocked(infos, basis, f, blocks)
    use precision, only: dp
    use types, only: information
    use basis_tools, only: basis_set
    use oqp_tagarray_driver
    use tagarray, only: TA_OK

    implicit none

    type(information), target, intent(inout) :: infos
    type(basis_set), intent(in) :: basis
    real(kind=dp), intent(inout) :: f(:,:)
    real(kind=dp), contiguous, intent(in) :: blocks(:)

    integer(8), contiguous, pointer :: shell_map(:)
    integer(4) :: status
    integer :: nbf, nshell, nops, iop, nf, k, j, s, off_k, off_j
    integer :: mu, nu, idx, blk0, blk_per_op
    real(kind=dp), allocatable :: fsq(:,:), y(:,:), acc(:,:)

    call tagarray_get_data(infos%dat, OQP_sym_shell_map, shell_map, status=status)
    if (status /= TA_OK) return

    nbf = basis%nbf
    nshell = basis%nshell
    if (mod(size(shell_map), nshell) /= 0) return
    nops = int(size(shell_map)/nshell)
    if (nops < 2) return

    blk_per_op = 0
    do k = 1, nshell
      s = shell_size(basis, k, nbf)
      blk_per_op = blk_per_op + s*s
    end do
    if (size(blocks) /= nops*blk_per_op) return

    allocate(fsq(nbf, nbf), y(nbf, nbf), acc(nbf, nbf))

    do nf = 1, ubound(f, 2)
      ! unpack the packed lower triangle
      idx = 0
      do mu = 1, nbf
        do nu = 1, mu
          idx = idx + 1
          fsq(mu, nu) = f(idx, nf)
          fsq(nu, mu) = f(idx, nf)
        end do
      end do

      acc = 0.0_dp
      do iop = 1, nops
        ! Operator transform: F <- T^T F T (T maps shell k to shell j with
        ! block B_k; T is metric-orthogonal, not orthogonal, so the
        ! transpose side matters once d shells mix under rotations).
        ! Y = T^T F : rows of source shell k get B_k^T @ rows of shell j.
        blk0 = (iop-1)*blk_per_op
        do k = 1, nshell
          s = shell_size(basis, k, nbf)
          j = int(shell_map((iop-1)*nshell + k))
          off_k = basis%ao_offset(k) - 1
          off_j = basis%ao_offset(j) - 1
          associate(b => reshape(blocks(blk0+1:blk0+s*s), [s, s]))
            y(off_k+1:off_k+s, :) = matmul(transpose(b), fsq(off_j+1:off_j+s, :))
          end associate
          blk0 = blk0 + s*s
        end do
        ! acc += Y T : columns of source shell k get Y cols j @ B_k.
        blk0 = (iop-1)*blk_per_op
        do k = 1, nshell
          s = shell_size(basis, k, nbf)
          j = int(shell_map((iop-1)*nshell + k))
          off_k = basis%ao_offset(k) - 1
          off_j = basis%ao_offset(j) - 1
          associate(b => reshape(blocks(blk0+1:blk0+s*s), [s, s]))
            acc(:, off_k+1:off_k+s) = acc(:, off_k+1:off_k+s) &
                + matmul(y(:, off_j+1:off_j+s), b)
          end associate
          blk0 = blk0 + s*s
        end do
      end do

      acc = acc/real(nops, dp)

      idx = 0
      do mu = 1, nbf
        do nu = 1, mu
          idx = idx + 1
          f(idx, nf) = acc(mu, nu)
        end do
      end do
    end do

  contains

    integer function shell_size(basis, k, nbf) result(s)
      type(basis_set), intent(in) :: basis
      integer, intent(in) :: k, nbf
      if (k < basis%nshell) then
        s = basis%ao_offset(k+1) - basis%ao_offset(k)
      else
        s = nbf - basis%ao_offset(k) + 1
      end if
    end function shell_size

  end subroutine symmetrize_skeleton_blocked

!--------------------------------------------------------------------------------

!> @brief Abelian (signed-permutation) skeleton symmetrization.
  subroutine symmetrize_skeleton_signed(infos, nbf, f)
    use precision, only: dp
    use types, only: information
    use oqp_tagarray_driver
    use tagarray, only: TA_OK

    implicit none

    type(information), target, intent(inout) :: infos
    integer, intent(in) :: nbf
    real(kind=dp), intent(inout) :: f(:,:)

    integer(8), contiguous, pointer :: target_map(:)
    real(kind=dp), contiguous, pointer :: sign_map(:)
    real(kind=dp), allocatable :: acc(:)
    integer(4) :: status
    integer :: nops, iop, nf, mu, nu, tm, tn, a, b, idx, tidx, base

    call tagarray_get_data(infos%dat, OQP_sym_ao_target, target_map, status=status)
    if (status /= TA_OK) return
    call tagarray_get_data(infos%dat, OQP_sym_ao_sign, sign_map, status=status)
    if (status /= TA_OK) return
    if (size(sign_map) /= size(target_map)) return
    if (mod(size(target_map), nbf) /= 0) return

    ! Flat layout, AO index fastest: target(mu, op) = target_map((op-1)*nbf+mu).
    nops = int(size(target_map)/nbf)
    if (nops < 2) return

    allocate(acc(nbf*(nbf+1)/2))

    do nf = 1, ubound(f, 2)
      acc = 0.0_dp
      do iop = 1, nops
        base = (iop-1)*nbf
        idx = 0
        do mu = 1, nbf
          tm = int(target_map(base + mu))
          do nu = 1, mu
            idx = idx + 1
            tn = int(target_map(base + nu))
            a = max(tm, tn)
            b = min(tm, tn)
            tidx = a*(a-1)/2 + b
            acc(tidx) = acc(tidx) &
                    + sign_map(base + mu)*sign_map(base + nu)*f(idx, nf)
          end do
        end do
      end do
      f(:, nf) = acc/real(nops, dp)
    end do

  end subroutine symmetrize_skeleton_signed
  !> @brief Builds AO-space linear response vector(s) in packed (triangular) form.
  !> @detail Forms the Coulomb/exchange response and, when using DFT, the
  !>         exchange–correlation kernel contribution in AO space.
  !>         For RHF: v1 = J/K(dm1) + f_xc(dm1).
  !>         For UHF/ROHF: spin-separated v1α, v1β using dm1α, dm1β.
  !>         Uses `fock_jk` for J/K and `tddft_fxc` / `utddft_fxc` for the XC kernel.
  !>         All AO matrices are in packed (upper-triangular) storage unless noted.
  !> @param[in]  basis     Basis set information.
  !> @param[inout] infos   System/control information (used to detect SCF type and DFT flags).
  !> @param[in]  molGrid   DFT molecular grid (required if DFT/XC kernel is used).
  !> @param[inout] mo_a    AO→MO coefficients for α (nbf×nbf). May be updated by XC routines.
  !> @param[in]  dm1_tri   First-order AO density in packed form:
  !>                       RHF: (nbf*(nbf+1)/2, 1)
  !>                       U/R: (nbf*(nbf+1)/2, 2) for α,β.
  !> @param[out] v1_tri    Packed AO response vector(s), same shape as dm1_tri.
  !> @param[inout,opt] mo_b AO→MO coefficients for β (nbf×nbf). Required for UHF;
  !>                        for ROHF it may be absent, in which case α is reused.
  !> @author Mohsen Mazaherifar
  !> @date August 2025
  subroutine get_response_packed(basis, infos, molGrid, mo_a, dm1_tri, v1_tri, mo_b)
      use precision,           only: dp
      use basis_tools,         only: basis_set
      use types,               only: information
      use mathlib,             only: unpack_matrix, pack_matrix,symmetrize_matrix
      use mod_dft_molgrid,     only: dft_grid_t
      use mod_dft_gridint_fxc, only: tddft_fxc, utddft_fxc
      implicit none

      type(basis_set),   intent(in)    :: basis
      type(information), intent(inout) :: infos
      type(dft_grid_t),  intent(in)    :: molGrid
      real(dp),          intent(inout)    :: mo_a(:,:)      ! nbf x nbf  (AO->MO)
      real(dp),          intent(in)    :: dm1_tri(:,:)     ! nbf*(nbf+1)/2  (packed)
      real(dp),          intent(out)   :: v1_tri(:,:)      ! packed AO response
      real(dp), optional, intent(inout) :: mo_b(:,:)

      integer :: nbf, nbf2, ok
      logical :: is_dft

      ! Packed work for int2
      real(dp), allocatable :: d_pack(:,:)

      ! Full AO work for XC response
      real(dp), allocatable :: dm1_full(:,:), fx_full(:,:), fx_pack(:)
      real(dp), allocatable :: dx3(:,:,:), fx3(:,:,:)  ! rank-3 wrappers for tddft_fxc
      real(kind=dp), allocatable :: dxa(:,:,:), dxb(:,:,:)
      real(kind=dp), allocatable :: fxa(:,:,:), fxb(:,:,:)
      real(dp) :: scalefactor

      nbf   = basis%nbf
      nbf2  = nbf*(nbf+1)/2
      is_dft = (infos%control%hamilton == 20)

      if (is_dft) then
        scalefactor = infos%dft%HFscale
      else
        scalefactor = 1.0_dp
      end if

      ! --- (2) XC-kernel part (DFT only): v_xc^(1) ---
      select case (infos%control%scftype)
      case (scf_rhf)
        call fock_jk(basis, d=dm1_tri, f=v1_tri, scale_exch=scalefactor, infos=infos)
        if (is_dft) then
          allocate(dm1_full(nbf,nbf), fx_full(nbf,nbf), fx_pack(nbf2), stat=ok)
          if (ok/=0) stop "alloc fail full"
          call unpack_matrix(dm1_tri(:,1), dm1_full)
          allocate(dx3(nbf,nbf,1), fx3(nbf,nbf,1), stat=ok); if (ok/=0) stop "alloc fail dx3/fx3"
          dx3(:,:,1) = dm1_full
          fx3(:,:,1) = 0.0_dp
          call tddft_fxc( basis=basis, molGrid=molGrid, isVecs=.true., wf=mo_a, &
                          fx=fx3, dx=dx3, nmtx=1, threshold=0.0_dp, infos=infos )
          fx_full = fx3(:,:,1)*0.5
          call pack_matrix(fx_full, fx_pack)
          v1_tri(:,1) = v1_tri(:,1) + fx_pack
        end if
      case (scf_rohf, scf_uhf)
        call fock_jk(basis, d=dm1_tri, f=v1_tri, scale_exch=scalefactor, infos=infos)
        if (is_dft) then
          allocate(dxa(nbf,nbf,1), dxb(nbf,nbf,1), fxa(nbf,nbf,1), fxb(nbf,nbf,1), fx_pack(nbf2), fx_full(nbf,nbf), stat=ok)
          if (ok/=0) stop "alloc fail full"
          call unpack_matrix(dm1_tri(:,1), dxa(:,:,1))
          call unpack_matrix(dm1_tri(:,2), dxb(:,:,1))
          fxa = 0
          fxb = 0
          call utddft_fxc(basis=basis, molGrid=molGrid, isVecs=.true., &
                     wfa=mo_a, wfb=mo_b, &
                     fxa=fxa, fxb=fxb, &
                     dxa=dxa, dxb=dxb, &
                     nMtx=1, threshold=0.0_dp, infos=infos)
          fx_full = fxa(:,:,1)
          fx_pack = 0.0_dp
          call pack_matrix(fx_full, fx_pack)
          v1_tri(:,1) = v1_tri(:,1) + fx_pack
          fx_full = fxb(:,:,1)
          fx_pack = 0.0_dp
          call pack_matrix(fx_full, fx_pack)
          v1_tri(:,2) = v1_tri(:,2) + fx_pack
          deallocate(dxa,dxb,fxa,fxb,fx_pack,fx_full)
        end if
      end select

  end subroutine get_response_packed

  !> @brief Computes DFT exchange–correlation contributions (matrix and energies).
  !> @detail Calls `dftexcor` to form the packed AO XC matrix pfxc and the
  !>         XC/total electron/kinetic energies. Handles RHF, UHF, and ROHF:
  !>         - RHF: single packed matrix used for both spins.
  !>         - UHF: separate α/β packed matrices.
  !>         - ROHF: β MOs are taken equal to α (mo_b := mo_a) for the call.
  !> @param[inout] infos   System/control information (reads SCF type).
  !> @param[in]    basis   Basis set information.
  !> @param[in]    molgrid DFT molecular grid and quadrature weights.
  !> @param[out]   pfxc    Packed AO XC matrix/matrices:
  !>                       RHF: (nbf_tri,1)
  !>                       U/R: (nbf_tri,2) for α,β.
  !> @param[out]   eexc    Exchange–correlation energy.
  !> @param[out]   totele  Total electron energy on the grid (xc driver report).
  !> @param[out]   totkin  Kinetic energy on the grid (xc driver report).
  !> @param[inout] mo_a    AO→MO coefficients for α (nbf×nbf).
  !> @param[inout] mo_b    AO→MO coefficients for β (nbf×nbf). For ROHF, set to mo_a.
  !> @author Mohsen Mazaherifar
  !> @date August 2025
  subroutine calc_dft_xc(infos, basis, molgrid, pfxc, eexc, totele, totkin, mo_a, mo_b)
    use precision, only: dp
    use types, only: information
    use dft, only: dftexcor
    use mod_dft_molgrid, only: dft_grid_t
    use basis_tools, only: basis_set
    use oqp_tagarray_driver
    use tagarray, only: TA_OK
    implicit none

    type(basis_set), intent(in) :: basis
    type(dft_grid_t), intent(in) :: molgrid
    real(kind=dp), intent(inout) :: mo_a(:,:)
    real(kind=dp), intent(inout) :: mo_b(:,:)
    real(kind=dp), intent(out) :: pfxc(:,:)
    real(kind=dp), contiguous, pointer :: sym_atom_weight(:)
    integer(8), contiguous, pointer :: sym_petite_flag(:)
    integer(4) :: sym_status
    logical :: sym_active
    real(kind=dp), intent(out) :: eexc
    real(kind=dp), intent(out) :: totele
    real(kind=dp), intent(out) :: totkin
    type(information), intent(inout) :: infos
    integer :: scf_type, nbf, nbf_tri
    ! Local parameters for SCF type
    integer, parameter :: scf_rhf = 1, scf_uhf = 2, scf_rohf = 3

    ! Initialize exchange-correlation contribution
    pfxc = 0.0_dp
    eexc = 0.0_dp
    totele = 0.0_dp
    totkin = 0.0_dp
    scf_type = infos%control%scftype
    nbf = basis%nbf
    nbf_tri = nbf*(nbf+1)/2

    ! Symmetry XC reduction: integrate only unique atoms' grid slices
    ! (orbit-weighted) and symmetrize the resulting skeleton XC matrix.
    ! Gated by the same petite flag as the two-electron reduction, so the
    ! stability-stage fail-safe applies here as well.
    sym_active = .false.
    sym_atom_weight => null()
    call tagarray_get_data(infos%dat, OQP_sym_petite, sym_petite_flag, status=sym_status)
    if (sym_status == TA_OK) then
      if (sym_petite_flag(1) /= 0) then
        call tagarray_get_data(infos%dat, OQP_sym_atom_weight, sym_atom_weight, &
                               status=sym_status)
        sym_active = sym_status == TA_OK
        if (sym_active) sym_active = size(sym_atom_weight) == infos%mol_prop%natom
      end if
    end if

    ! Calculate exchange-correlation based on SCF type
    if (sym_active) then
      if (scf_type == scf_rhf) then
        call dftexcor(basis, molgrid, 1, pfxc, pfxc, mo_a, mo_a, &
                      nbf, nbf_tri, eexc, totele, totkin, infos, sym_atom_weight)
      else if (scf_type == scf_uhf) then
        call dftexcor(basis, molgrid, 2, pfxc(:,1), pfxc(:,2), mo_a, mo_b, &
                      nbf, nbf_tri, eexc, totele, totkin, infos, sym_atom_weight)
      else if (scf_type == scf_rohf) then
        mo_b = mo_a
        call dftexcor(basis, molgrid, 2, pfxc(:,1), pfxc(:,2), mo_a, mo_b, &
                      nbf, nbf_tri, eexc, totele, totkin, infos, sym_atom_weight)
      end if
      ! The reduced-grid XC matrix is a skeleton: project onto the totally
      ! symmetric component (the XC energy/electron count are already exact).
      call symmetrize_skeleton_fock(infos, basis, pfxc)
    else if (scf_type == scf_rhf) then
      ! Restricted calculation - same matrix for alpha and beta
      call dftexcor(basis, molgrid, 1, pfxc, pfxc, mo_a, mo_a, &
                    nbf, nbf_tri, eexc, totele, totkin, infos)
    else if (scf_type == scf_uhf) then
      ! Unrestricted calculation - separate matrices for alpha and beta
      call dftexcor(basis, molgrid, 2, pfxc(:,1), pfxc(:,2), mo_a, mo_b, &
                    nbf, nbf_tri, eexc, totele, totkin, infos)
    else if (scf_type == scf_rohf) then
      ! Restricted open-shell calculation
      ! ROHF does not have MO_B, so we copy MO_A to MO_B
      mo_b = mo_a
      call dftexcor(basis, molgrid, 2, pfxc(:,1), pfxc(:,2), mo_a, mo_b, &
                    nbf, nbf_tri, eexc, totele, totkin, infos)
    end if

  end subroutine calc_dft_xc

  !> @brief Computes DFT exchange-correlation contributions from explicit AO density matrices.
  subroutine calc_dft_xc_density(infos, basis, molgrid, dmat, pfxc, eexc, totele, totkin)
    use precision, only: dp
    use types, only: information
    use mod_dft_molgrid, only: dft_grid_t
    use basis_tools, only: basis_set
    use mod_dft_gridint_energy, only: dmatd_density_blk
    use mathlib, only: unpack_matrix
    implicit none

    type(information), intent(inout) :: infos
    type(basis_set), intent(in) :: basis
    type(dft_grid_t), intent(in) :: molgrid
    real(kind=dp), intent(in) :: dmat(:,:)
    real(kind=dp), intent(out) :: pfxc(:,:)
    real(kind=dp), intent(out) :: eexc, totele, totkin

    integer :: scf_type, nbf, nbf_tri, nang
    logical :: urohf
    real(kind=dp), allocatable :: da(:,:), db(:,:)

    scf_type = infos%control%scftype
    urohf = scf_type /= scf_rhf
    nbf = basis%nbf
    nbf_tri = nbf*(nbf+1)/2
    nang = maxval(basis%am)+1+1
    allocate(da(nbf,nbf), source=0.0_dp)
    call unpack_matrix(dmat(:,1), da, nbf, "U")
    allocate(db(nbf,nbf), source=0.0_dp)
    if (urohf .and. size(dmat,2) > 1) then
      call unpack_matrix(dmat(:,2), db, nbf, "U")
    else
      db = da
    end if

    pfxc = 0.0_dp
    call dmatd_density_blk(basis, molgrid, da, db, pfxc(:,1), pfxc(:,min(2,size(pfxc,2))), &
                          eexc, totele, totkin, nang, nbf, infos%dft%grid_density_cutoff, &
                          urohf, infos)

    deallocate(da, db)
  end subroutine calc_dft_xc_density

  !> @brief Builds J/K (and optional DFT XC) Fock contribution(s) and energies.
  !> @detail Forms two-electron Fock using `fock_jk`, adds the one-electron core
  !>         Hamiltonian, and accumulates SCF energy components:
  !>         E_hf1 = Tr[D·Hcore], E_hf = ½·Σ_i Tr[D_i·F_i] + ½·E_hf1, E_tot = E_hf + E_nuc.
  !>         If DFT (infos%control%hamilton ≥ 20), adds packed XC matrix (via `calc_dft_xc`)
  !>         and XC energy to F and E.
  !>         Supports incremental updates when both d_old and f_old are provided:
  !>         builds F for ΔD = D − D_old, accumulates into F (and updates D_old,F_old).
  !> @param[in]     basis    Basis set information.
  !> @param[inout]  infos    System/control information and runtime data.
  !> @param[inout]  d        Packed AO density(ies), shape (nbf_tri,nfocks).
  !> @param[in]     hcore    Packed one-electron core Hamiltonian (nbf_tri).
  !> @param[in]     nfocks   Number of spin blocks: 1 (RHF) or 2 (UHF/ROHF).
  !> @param[inout]  f        Packed Fock matrix(ces) to fill (nbf_tri,nfocks).
  !> @param[inout]  E        SCF energy accumulator (ehf1, ehf, etot, … are updated).
  !> @param[in,opt] molgrid  DFT grid (required if DFT is active).
  !> @param[inout,opt] mo_a  AO→MO α (nbf×nbf); required if DFT is active.
  !> @param[inout,opt] mo_b  AO→MO β (nbf×nbf); required for UHF. For ROHF α is reused.
  !> @param[inout]  nschwz   (Output) number of Schwarz-screened quartets (from ERI driver).
  !> @param[inout,opt] f_old Previously accumulated packed Fock(ces) for incremental build.
  !> @param[inout,opt] d_old Previous packed density(ies) for incremental build.
  !> @note For DFT hybrids, exchange scaling is taken from infos%dft%HFscale.
  !> @note Continuum solvent (PCM) is the single canonical runtime path: gated on
  !>       infos%control%pcm_enabled, applied via add_pcm_reaction_field, and
  !>       reported in E%e_pcm. There is no second reaction-field hook.
  !> @throws error stop if DFT is requested but molgrid/mo_a are not provided.
  !> @author Mohsen Mazaherifar
  !> @date August 2025
  subroutine calc_jk_xc(basis, infos, d, hcore, nfocks, f, E, &
                               molgrid, mo_a, mo_b, nschwz, f_old, d_old, density_xc, xc_reuse)
    use precision,       only : dp
    use basis_tools,     only : basis_set
    use types,           only : information
    use mod_dft_molgrid, only : dft_grid_t
    use mathlib,          only : traceprod_sym_packed
    use solvent_pcm,      only : add_pcm_reaction_field
    use mod_dft_incdft,  only : g_xc_ref, incdft_store
    implicit none

    type(basis_set),   intent(in)    :: basis
    type(information), intent(inout) :: infos
    real(dp),          intent(inout)    :: d(:,:)              ! (nbf_tri, nfocks)
    real(dp),          intent(in)    :: hcore(:)
    type(scf_energy_t), intent(inout)        :: E
    type(dft_grid_t),  intent(in),   optional :: molgrid
    real(dp),          intent(inout),optional :: mo_a(:,:)  ! (nbf, nbf)
    real(dp),          intent(inout),optional :: mo_b(:,:)  ! (nbf, nbf)
    real(dp),          intent(inout) :: f(:,:)              ! (nbf_tri, nfocks)
    real(dp), intent(inout), optional        :: d_old(:,:), f_old(:,:)
    logical, intent(in), optional :: density_xc
    integer,  intent(inout)    :: nschwz
    integer, intent(in) :: nfocks
    !> Opt 2 (IncDFT): when present and .true., reuse the stored reference XC
    !> matrix/energy instead of rebuilding from the density this iteration.
    logical, intent(in), optional :: xc_reuse

    real(dp) :: scale_factor

    integer  :: scf_type, nbf
    integer :: ii
    real(dp), allocatable :: pfxc(:,:)
    logical :: is_dft = .false., use_density_xc
    logical :: xc_reused
    ! Env-gated (OQP_XC_TIMING) per-iteration wall split: J/K vs XC build.
    logical :: do_t
    character(len=8) :: tenv
    integer :: tst, tln
    integer(8) :: clk0, clk1, clkr
    real(dp) :: wall_jk, wall_xc

    call get_environment_variable('OQP_XC_TIMING', tenv, length=tln, status=tst)
    do_t = (tst == 0 .and. tln > 0 .and. &
        (tenv(1:1) == '1' .or. tenv(1:1) == 't' .or. tenv(1:1) == 'T' .or. &
         tenv(1:1) == 'y' .or. tenv(1:1) == 'Y' .or. tenv(1:1) == 'o' .or. tenv(1:1) == 'O'))
    wall_jk = 0.0_dp; wall_xc = 0.0_dp
    call system_clock(count_rate=clkr)

    is_dft = infos%control%hamilton >= 20
    use_density_xc = .false.
    if (present(density_xc)) use_density_xc = density_xc
    if (is_dft) then
      scale_factor = infos%dft%HFscale
    else
      scale_factor = 1.0_dp
    end if

    nbf      = basis%nbf
    if (do_t) call system_clock(count=clk0)
    if(present(d_old) .and. present(f_old)) then
      d = d - d_old
      call fock_jk(basis, d, f, infos, scale_factor, nschwz , f_old, petite=.true.)
      d = d + d_old
      d_old = d
    else
      call fock_jk(basis, d, f, infos, scale_factor, nschwz, petite=.true.)
    end if
    if (do_t) then
      call system_clock(count=clk1)
      wall_jk = real(clk1-clk0, dp)/real(clkr, dp)
    end if
    ii = 0
    do ii = 1, nfocks
      f(:,ii) =  f(:,ii) + hcore
    end do
    !----------------------------------------------------------------------------
    ! Compute HF Energy Components
    !----------------------------------------------------------------------------
    E%ehf = 0.0_dp
    E%ehf1 = 0.0_dp

    ! compute one and two-electron energies
    do ii = 1, nfocks
      E%ehf1 = E%ehf1 + traceprod_sym_packed(d(:,ii), hcore, nbf)
      E%ehf = E%ehf + traceprod_sym_packed(d(:,ii), f(:,ii), nbf)
    end do

    E%ehf = 0.5_dp * (E%ehf + E%ehf1)
    E%etot = E%ehf + E%nenergy

    ! PCM solvent reaction field (provisional energy-only path; ddX backend).
    ! Mirrors the XC pattern below: the V_pcm operator is added to the Fock
    ! blocks used for the next density update, and a distinct E_pcm term is
    ! added to the total energy. It is applied AFTER the vacuum HF energy is
    ! formed so the reaction field is not double-counted in E%ehf. Gated on
    ! pcm_enabled; aborts at runtime if built without ddX (OQP_ENABLE_DDX).
    E%e_pcm = 0.0_dp
    if (infos%control%pcm_enabled) then
      call add_pcm_reaction_field(basis, infos, d, nfocks, f, E%e_pcm)
      E%etot = E%etot + E%e_pcm
    end if

    if (.not. is_dft) return

    if (.not.present(molgrid) .or. .not.present(mo_a)) then
      error stop 'calc_jk_xc: DFT requested but molgrid/mo_a/mo_b not provided.'
    end if


    allocate(pfxc(nbf*(nbf+1)/2, nfocks))
    pfxc = 0.0_dp

    if (do_t) call system_clock(count=clk0)

    ! Opt 2 (IncDFT): reuse the reference XC matrix/energy when the caller signals
    ! the density has effectively stopped changing (controlled, late-SCF window).
    xc_reused = .false.
    if (present(xc_reuse)) xc_reused = xc_reuse .and. g_xc_ref%valid &
                            .and. g_xc_ref%ntri == nbf*(nbf+1)/2 .and. g_xc_ref%nf == nfocks
    if (xc_reused) then
      pfxc      = g_xc_ref%vxc
      E%eexc    = g_xc_ref%eexc
      E%totele  = g_xc_ref%totele
      E%totkin  = g_xc_ref%totkin
      g_xc_ref%reuse_run = g_xc_ref%reuse_run + 1
      g_xc_ref%n_reuse   = g_xc_ref%n_reuse + 1
    else
      if (use_density_xc) then
        call calc_dft_xc_density(infos, basis, molgrid, d, pfxc, E%eexc, E%totele, E%totkin)
      else
        call calc_dft_xc(infos, basis, molgrid, pfxc, E%eexc, E%totele, E%totkin, mo_a, mo_b)
      end if
      ! Refresh the IncDFT reference from this full build (only when IncDFT is on).
      if (infos%control%xc_incdft /= 0) call incdft_store(pfxc, E%eexc, E%totele, E%totkin)
    end if

    if (do_t) then
      call system_clock(count=clk1)
      wall_xc = real(clk1-clk0, dp)/real(clkr, dp)
      write(*,'(1x,a,f9.4,a,f9.4,a,f6.1,a,l2)') '[SCFTIME] wall_JK=', wall_jk, &
        's  wall_XCbuild=', wall_xc, 's  XC_frac=', &
        100.0_dp*wall_xc/max(1.0d-12, wall_jk+wall_xc), '%  xc_reused=', xc_reused
    end if

    f = f + pfxc
    E%etot=E%etot + E%eexc

    deallocate(pfxc)

  end subroutine calc_jk_xc

  !> @brief High-level AO-Fock builder and energy evaluation.
  !> @detail Retrieves required matrices/vectors from the internal tagarray
  !>         (Hcore, T, S, D_α[,_β], MO_α[,_β]) and constructs the packed AO
  !>         Fock matrix(ces) via `calc_jk_xc`. Computes one- and two-electron
  !>         energy components, nuclear repulsion, virial ratio, and stores
  !>         the packed Fock back to tagarray (FOCK_A[, FOCK_B]).
  !>         Supports optional overrides for MO/Density and incremental updates.
  !> @param[in]     basis     Basis set information.
  !> @param[inout]  infos     System/control information and tagarray store.
  !> @param[in]     molgrid   DFT molecular grid (used when DFT is active).
  !> @param[inout]  fock_ao   Packed AO Fock output: (nbf_tri, nfocks).
  !> @param[inout]  E         SCF energy structure (fields updated).
  !> @param[inout,opt] mo_a_in Override AO→MO α (nbf×nbf).
  !> @param[inout,opt] mo_b_in Override AO→MO β (nbf×nbf).
  !> @param[inout,opt] dens_in Override packed AO density(ies) (nbf_tri, nfocks).
  !> @param[inout,opt] dens_old Previous packed density(ies) for incremental build.
  !> @param[inout,opt] f_old    Previous packed Fock(ces) for incremental build.
  !> @param[inout,opt] nschwz   (Output) count of Schwarz-screened quartets.
  !> @note Continuum solvent (PCM) is driven entirely inside calc_jk_xc via the
  !>       single infos%control%pcm_enabled gate; calc_fock takes no PCM argument.
  !> @author Mohsen Mazaherifar
  !> @date August 2025
  subroutine calc_fock(basis, infos, molgrid, fock_ao, E, mo_a_in, dens_in, mo_b_in, nschwz, f_old, dens_old, xc_reuse)
    use precision,       only : dp
    use oqp_tagarray_driver
    use types,           only : information
    use mod_dft_molgrid, only : dft_grid_t
    use basis_tools,     only : basis_set
    use util,            only : e_charge_repulsion
    use mathlib,         only : traceprod_sym_packed, unpack_matrix
    implicit none

    type(basis_set), intent(in)              :: basis
    type(information), target, intent(inout) :: infos
    type(dft_grid_t), intent(in)             :: molgrid
    real(dp), intent(inout), target          :: fock_ao(:,:)
    type(scf_energy_t), intent(inout)        :: E

    ! optionals
    real(dp), intent(inout), optional        :: mo_a_in(:,:)
    real(dp), intent(inout), optional        :: mo_b_in(:,:)
    real(dp), intent(inout), optional        :: dens_in(:,:)
    real(dp), intent(inout), optional        :: dens_old(:,:)
    real(dp), intent(inout), optional        :: f_old(:,:)
    integer,  intent(inout), optional        :: nschwz
    !> Opt 2 (IncDFT): reuse the reference XC matrix this iteration (caller-decided)
    logical,  intent(in),    optional        :: xc_reuse

    ! locals
    integer :: nbf, nbf_tri, nfocks, nelec, scf_type, ii
    logical :: is_dft
    real(dp), allocatable :: pdmat(:,:), pfock(:,:)
    real(dp), contiguous, pointer :: hcore(:), tmat(:), smat(:)
    real(dp), contiguous, pointer :: dmat_a(:), dmat_b(:), fock_a(:), fock_b(:)
    real(dp), contiguous, pointer :: mo_a(:,:), mo_b(:,:)

    ! SCF type & sizes
    select case (infos%control%scftype)
    case (1); scf_type = 1; nfocks = 1
    case (2,3); scf_type = 2; nfocks = 2
    end select
    nelec   = infos%mol_prop%nelec
    nbf     = basis%nbf
    nbf_tri = nbf*(nbf+1)/2
    is_dft  = infos%control%hamilton >= 20
    ! tag arrays
    call tagarray_get_data(infos%dat, OQP_Hcore, hcore)
    call tagarray_get_data(infos%dat, OQP_TM,    tmat)
    call tagarray_get_data(infos%dat, OQP_SM, smat)
    call tagarray_get_data(infos%dat, OQP_DM_A,  dmat_a)
    call tagarray_get_data(infos%dat, OQP_FOCK_A,fock_a)
    call tagarray_get_data(infos%dat, OQP_VEC_MO_A, mo_a)
    if (nfocks > 1) then
      call tagarray_get_data(infos%dat, OQP_DM_B,  dmat_b)
      call tagarray_get_data(infos%dat, OQP_FOCK_B,fock_b)
      call tagarray_get_data(infos%dat, OQP_VEC_MO_B, mo_b)
    end if

    if (present(mo_a_in)) mo_a = mo_a_in
    if (present(mo_b_in) .and. nfocks>1) mo_b = mo_b_in

    allocate(pdmat(nbf_tri,nfocks)); pdmat(:,1) = dmat_a
    if (nfocks > 1) pdmat(:,2) = dmat_b
    if (present(dens_in)) pdmat = dens_in

    E%nenergy = e_charge_repulsion(infos%atoms%xyz, infos%atoms%zn - infos%basis%ecp_zn_num)

    fock_ao = 0.0_dp
    if (present(dens_old)) then
      call calc_jk_xc(basis, infos, pdmat, hcore, nfocks, &
                    fock_ao, E, molgrid, mo_a, mo_b, nschwz, f_old, dens_old, density_xc=present(dens_in), &
                    xc_reuse=xc_reuse)
    else
      call calc_jk_xc(basis, infos, pdmat, hcore, nfocks, &
                    fock_ao, E, molgrid, mo_a, mo_b, nschwz, density_xc=present(dens_in), &
                    xc_reuse=xc_reuse)
    end if

    E%psinrm    = 0.0_dp
    E%tkin = 0.0_dp
    do ii = 1, nfocks
      E%psinrm   = E%psinrm    + traceprod_sym_packed(pdmat(:,ii), smat, nbf)/nelec
      E%tkin = E%tkin + traceprod_sym_packed(pdmat(:,ii), tmat,  nbf)
    end do
    E%vne    = E%ehf1 - E%tkin
    E%vee    = E%etot - E%ehf1 - E%nenergy
    E%vnn    = E%nenergy
    E%vtot   = E%vne + E%vnn + E%vee
    E%virial = - E%vtot / E%tkin

    ! store Fock back
    fock_a = fock_ao(:,1)
    if (nfocks > 1) then
      fock_b = fock_ao(:,2)
    end if

    infos%mol_energy%energy = E%etot

    deallocate(pdmat)

  end subroutine calc_fock

  function compute_energy(energy) result(etot)
    implicit none
    type(scf_energy_t), pointer :: energy
    real(dp)  :: etot
    etot = energy%etot
  end function

end module scf_addons
