!--------------------------------------------------------------------------------------------------!
!   CP2K-RS Extended Interface: Stress Tensor Wrapper                                             !
!   Copyright 2025 CP2K-RS developers                                                             !
!                                                                                                  !
!   SPDX-License-Identifier: GPL-2.0-or-later                                                      !
!--------------------------------------------------------------------------------------------------!

! **************************************************************************************************
!> \brief Extended CP2K C/C++ interface for stress tensor calculations
!> \par History
!>       01.2025 created as part of cp2k-rs extensions
!> \author CP2K-RS Developers
!> \note This module extends the standard libcp2k interface with additional functionality
!>       for accessing stress tensors, which are crucial for NPT molecular dynamics and
!>       materials science applications.
! **************************************************************************************************
MODULE cp2k_stress_extended
   USE ISO_C_BINDING,                   ONLY: C_CHAR,&
                                              C_DOUBLE,&
                                              C_INT,&
                                              C_NULL_CHAR
   USE f77_interface,                   ONLY: f_env_add_defaults,&
                                              f_env_rm_defaults,&
                                              f_env_type
   USE force_env_types,                 ONLY: force_env_get
   USE virial_types,                    ONLY: virial_type
   USE cell_types,                      ONLY: cell_type

   IMPLICIT NONE

   PRIVATE

   PUBLIC :: cp2k_get_stress_tensor
   PUBLIC :: cp2k_get_virial_tensor

CONTAINS

! **************************************************************************************************
!> \brief Get the stress tensor from a force environment
!> \param env_id The force environment ID
!> \param stress 3x3 stress tensor in GPa (output)
!> \par History
!>      01.2025 created
!> \author CP2K-RS Developers
!> \note The stress tensor is returned in GPa. It is computed from the virial tensor and cell volume.
!>       For a valid stress tensor, an energy/force calculation must have been performed first.
! **************************************************************************************************
   SUBROUTINE cp2k_get_stress_tensor(env_id, stress) BIND(C, name='cp2k_get_stress_tensor')
      INTEGER(C_INT), VALUE, INTENT(IN)                  :: env_id
      REAL(C_DOUBLE), DIMENSION(3, 3), INTENT(OUT)       :: stress

      INTEGER                                            :: ierr
      REAL(KIND=C_DOUBLE)                                :: volume
      TYPE(cell_type), POINTER                           :: cell
      TYPE(f_env_type), POINTER                          :: f_env
      TYPE(virial_type), POINTER                         :: virial

      ! Initialize output
      stress = 0.0_C_DOUBLE

      NULLIFY (f_env, virial, cell)

      ! Get the force environment
      CALL f_env_add_defaults(env_id, f_env)

      ! Try to get virial and cell from force environment
      try: BLOCK
         CALL force_env_get(f_env%force_env, cell=cell, virial=virial)

         ! Check if virial is available
         IF (.NOT. ASSOCIATED(virial)) THEN
            ! No virial available - stress will remain zero
            EXIT try
         END IF

         ! Check if cell is available
         IF (.NOT. ASSOCIATED(cell)) THEN
            EXIT try
         END IF

         ! Get cell volume
         volume = cell%deth

         ! Avoid division by zero
         IF (ABS(volume) < 1.0E-10_C_DOUBLE) THEN
            EXIT try
         END IF

         ! Calculate stress tensor: stress = -virial / volume
         ! Convert to GPa (atomic units to GPa conversion factor)
         ! 1 a.u. of pressure = 2.9421912e13 Pa = 2.9421912e4 GPa
         stress(1:3, 1:3) = -virial%pv_virial(1:3, 1:3) / volume * 2.9421912e4_C_DOUBLE

      END BLOCK try

      ! Remove defaults (cleanup)
      CALL f_env_rm_defaults(f_env, ierr)
      CPASSERT(ierr == 0)

   END SUBROUTINE cp2k_get_stress_tensor

! **************************************************************************************************
!> \brief Get the virial tensor from a force environment
!> \param env_id The force environment ID
!> \param virial_tensor 3x3 virial tensor in atomic units (output)
!> \par History
!>      01.2025 created
!> \author CP2K-RS Developers
!> \note The virial tensor is returned in atomic units (Hartree).
!>       For a valid virial, an energy/force calculation must have been performed first.
! **************************************************************************************************
   SUBROUTINE cp2k_get_virial_tensor(env_id, virial_tensor) BIND(C, name='cp2k_get_virial_tensor')
      INTEGER(C_INT), VALUE, INTENT(IN)                  :: env_id
      REAL(C_DOUBLE), DIMENSION(3, 3), INTENT(OUT)       :: virial_tensor

      INTEGER                                            :: ierr
      TYPE(f_env_type), POINTER                          :: f_env
      TYPE(virial_type), POINTER                         :: virial

      ! Initialize output
      virial_tensor = 0.0_C_DOUBLE

      NULLIFY (f_env, virial)

      ! Get the force environment
      CALL f_env_add_defaults(env_id, f_env)

      ! Try to get virial from force environment
      try: BLOCK
         CALL force_env_get(f_env%force_env, virial=virial)

         ! Check if virial is available
         IF (.NOT. ASSOCIATED(virial)) THEN
            ! No virial available - tensor will remain zero
            EXIT try
         END IF

         ! Copy virial tensor
         virial_tensor(1:3, 1:3) = virial%pv_virial(1:3, 1:3)

      END BLOCK try

      ! Remove defaults (cleanup)
      CALL f_env_rm_defaults(f_env, ierr)
      CPASSERT(ierr == 0)

   END SUBROUTINE cp2k_get_virial_tensor

END MODULE cp2k_stress_extended
