program export_trexio
  implicit none
  read_wf = .True.
  SOFT_TOUCH read_wf
  call run
end

subroutine run
  use trexio
  implicit none
  BEGIN_DOC
  !     Exports the wave function in TREXIO format
  END_DOC

  integer(8)                     :: f ! TREXIO file handle
  integer                        :: rc
  double precision, allocatable  :: factor(:)

  print *, 'TREXIO file : '//trim(trexio_filename)
  print *, ''

  call system('rm -rf '//trim(trexio_filename)//'.bak')
  call system('mv '//trim(trexio_filename)//' '//trim(trexio_filename)//'.bak')
  if (backend == 0) then
    f = trexio_open(trexio_filename, 'w', TREXIO_HDF5, rc)
  else if (backend == 1) then
    f = trexio_open(trexio_filename, 'w', TREXIO_TEXT, rc)
  endif
  if (f == 0_8) then
    print *, 'Unable to open TREXIO file for writing'
    print *, 'rc = ', rc
    stop -1
  endif

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

! Electrons
! ---------

  print *, 'Electrons'

  rc = trexio_write_electron_up_num(f, elec_alpha_num)
  call check_success(rc)

  rc = trexio_write_electron_dn_num(f, elec_beta_num)
  call check_success(rc)


! Nuclei
! ------

  print *, 'Nuclei'

  rc = trexio_write_nucleus_num(f, nucl_num)
  call check_success(rc)

  rc = trexio_write_nucleus_charge(f, nucl_charge)
  call check_success(rc)

  rc = trexio_write_nucleus_coord(f, nucl_coord_transp)
  call check_success(rc)

  rc = trexio_write_nucleus_label(f, nucl_label, 32)
  call check_success(rc)

  rc = trexio_write_nucleus_repulsion(f, nuclear_repulsion)
  call check_success(rc)


! Pseudo-potentials
! -----------------

  if (do_pseudo) then

    print *, 'ECP'
    integer                       :: num

    num = 0
    do k=1,pseudo_klocmax
      do i=1,nucl_num
        if (pseudo_dz_k(i,k) /= 0.d0) then
          num = num+1
        end if
      end do
    end do

    do l=0,pseudo_lmax
      do k=1,pseudo_kmax
        do i=1,nucl_num
          if (pseudo_dz_kl(i,k,l) /= 0.d0) then
            num = num+1
          end if
        end do
      end do
    end do

    integer, allocatable          :: ang_mom(:), nucleus_index(:), power(:), lmax(:)
    double precision, allocatable :: exponent(:), coefficient(:)

    allocate(ang_mom(num), nucleus_index(num), exponent(num), coefficient(num), power(num), &
            lmax(nucl_num) )

    do i=1,nucl_num
      lmax(i) = -1
      do l=0,pseudo_lmax
        do k=1,pseudo_kmax
          if (pseudo_dz_kl_transp(k,l,i) /= 0.d0) then
            lmax(i) = max(lmax(i), l)
          end if
        end do
      end do
    end do

    j = 0
    do i=1,nucl_num
      do k=1,pseudo_klocmax
        if (pseudo_dz_k_transp(k,i) /= 0.d0) then
          j = j+1
          ang_mom(j) = lmax(i)+1
          nucleus_index(j) = i
          exponent(j) = pseudo_dz_k_transp(k,i)
          coefficient(j) = pseudo_v_k_transp(k,i)
          power(j) = pseudo_n_k_transp(k,i)
        end if
      end do

      do l=0,lmax(i)
        do k=1,pseudo_kmax
          if (pseudo_dz_kl_transp(k,l,i) /= 0.d0) then
            j = j+1
            ang_mom(j) = l
            nucleus_index(j) = i
            exponent(j) = pseudo_dz_kl_transp(k,l,i)
            coefficient(j) = pseudo_v_kl_transp(k,l,i)
            power(j) = pseudo_n_kl_transp(k,l,i)
          end if
        end do
      end do
    end do


    lmax(:) = lmax(:)+1
    rc = trexio_write_ecp_max_ang_mom_plus_1(f, lmax)
    call check_success(rc)

    rc = trexio_write_ecp_z_core(f, int(nucl_charge_remove))
    call check_success(rc)

    rc = trexio_write_ecp_num(f, num)
    call check_success(rc)

    rc = trexio_write_ecp_ang_mom(f, ang_mom)
    call check_success(rc)

    rc = trexio_write_ecp_nucleus_index(f, nucleus_index)
    call check_success(rc)

    rc = trexio_write_ecp_exponent(f, exponent)
    call check_success(rc)

    rc = trexio_write_ecp_coefficient(f, coefficient)
    call check_success(rc)

    rc = trexio_write_ecp_power(f, power)
    call check_success(rc)

  endif


! Basis
! -----

  print *, 'Basis'


  rc = trexio_write_basis_type(f, 'Gaussian', len('Gaussian'))
  call check_success(rc)

  rc = trexio_write_basis_prim_num(f, prim_num)
  call check_success(rc)

   rc = trexio_write_basis_shell_num(f, shell_num)
   call check_success(rc)

   rc = trexio_write_basis_nucleus_index(f, basis_nucleus_index)
   call check_success(rc)

   rc = trexio_write_basis_shell_ang_mom(f, shell_ang_mom)
   call check_success(rc)

   allocate(factor(shell_num))
   if (ao_normalized) then
     factor(1:shell_num) = shell_normalization_factor(1:shell_num)
   else
     factor(1:shell_num) = 1.d0
   endif
   rc = trexio_write_basis_shell_factor(f, factor)
   call check_success(rc)

   deallocate(factor)

  rc = trexio_write_basis_shell_index(f, shell_index)
  call check_success(rc)

  rc = trexio_write_basis_exponent(f, prim_expo)
  call check_success(rc)

  rc = trexio_write_basis_coefficient(f, prim_coef)
  call check_success(rc)

  allocate(factor(prim_num))
  if (primitives_normalized) then
    factor(1:prim_num) = prim_normalization_factor(1:prim_num)
  else
    factor(1:prim_num) = 1.d0
  endif
  rc = trexio_write_basis_prim_factor(f, factor)
  call check_success(rc)
  deallocate(factor)


! Atomic orbitals
! ---------------

  print *, 'AOs'

  rc = trexio_write_ao_num(f, ao_num)
  call check_success(rc)

  rc = trexio_write_ao_cartesian(f, 1)
  call check_success(rc)

  rc = trexio_write_ao_shell(f, ao_shell)
  call check_success(rc)

  integer :: i, pow0(3), powA(3), j, k, l, nz
  double precision :: normA, norm0, C_A(3), overlap_x, overlap_z, overlap_y, c
  nz=100

  C_A(1) = 0.d0
  C_A(2) = 0.d0
  C_A(3) = 0.d0

  allocate(factor(ao_num))
  if (ao_normalized) then
    do i=1,ao_num
      l = ao_first_of_shell(ao_shell(i))
      factor(i) = (ao_coef_normalized(i,1)+tiny(1.d0))/(ao_coef_normalized(l,1)+tiny(1.d0))
    enddo
  else
    factor(:) = 1.d0
  endif
  rc = trexio_write_ao_normalization(f, factor)
  call check_success(rc)
  deallocate(factor)

! One-e AO integrals
! ------------------

  print *, 'AO integrals'

  rc = trexio_write_ao_1e_int_overlap(f,ao_overlap)
  call check_success(rc)

  rc = trexio_write_ao_1e_int_kinetic(f,ao_kinetic_integrals)
  call check_success(rc)

  rc = trexio_write_ao_1e_int_potential_n_e(f,ao_integrals_n_e)
  call check_success(rc)

  if (do_pseudo) then
    rc = trexio_write_ao_1e_int_ecp(f, ao_pseudo_integrals_local + ao_pseudo_integrals_non_local)
    call check_success(rc)
  endif

  rc = trexio_write_ao_1e_int_core_hamiltonian(f,ao_one_e_integrals)
  call check_success(rc)

! Two-e AO integrals
! ------------------

  if (export_ao_ints) then
    PROVIDE ao_two_e_integrals_in_map

    integer(8), parameter :: BUFSIZE=10000_8
    double precision :: eri_buffer(BUFSIZE), integral
    integer(4) :: eri_index(4,BUFSIZE)
    integer(8) :: icount, offset

    double precision, external :: get_ao_two_e_integral


    icount = 0_8
    offset = 0_8
    do l=1,ao_num
      do k=1,ao_num
        do j=l,ao_num
          do i=k,ao_num
            if (i==j .and. k<l) cycle
            if (i<j) cycle
            integral = get_ao_two_e_integral(i,j,k,l,ao_integrals_map)
            if (integral == 0.d0) cycle
            icount += 1_8
            eri_buffer(icount) = integral
            eri_index(1,icount) = i
            eri_index(2,icount) = j
            eri_index(3,icount) = k
            eri_index(4,icount) = l
            if (icount == BUFSIZE) then
              rc = trexio_write_ao_2e_int_eri(f, offset, icount, eri_index, eri_buffer)
              call check_success(rc)
              offset += icount
              icount = 0_8
            end if
          end do
        end do
      end do
    end do

    if (icount >= 0_8) then
      rc = trexio_write_ao_2e_int_eri(f, offset, icount, eri_index, eri_buffer)
      call check_success(rc)
    end if
  end if


! Molecular orbitals
! ------------------

  print *, 'MOs'

  rc = trexio_write_mo_type(f, mo_label, len(trim(mo_label)))
  call check_success(rc)

  rc = trexio_write_mo_num(f, mo_num)
  call check_success(rc)

  rc = trexio_write_mo_coefficient(f, mo_coef)
  call check_success(rc)


! One-e MO integrals
! ------------------

  print *, 'MO integrals'

  rc = trexio_write_mo_1e_int_kinetic(f,mo_kinetic_integrals)
  call check_success(rc)

  rc = trexio_write_mo_1e_int_potential_n_e(f,mo_integrals_n_e)
  call check_success(rc)

  if (do_pseudo) then
    rc = trexio_write_mo_1e_int_ecp(f,mo_pseudo_integrals_local)
    call check_success(rc)
  endif

  rc = trexio_write_mo_1e_int_core_hamiltonian(f,mo_one_e_integrals)
  call check_success(rc)

! Two-e MO integrals
! ------------------

  if (export_mo_ints) then
    PROVIDE mo_two_e_integrals_in_map

    double precision, external :: mo_two_e_integral


    icount = 0_8
    offset = 0_8
    do l=1,mo_num
      do k=1,mo_num
        do j=l,mo_num
          do i=k,mo_num
            if (i==j .and. k<l) cycle
            if (i<j) cycle
            integral = mo_two_e_integral(i,j,k,l)
            if (integral == 0.d0) cycle
            icount += 1_8
            eri_buffer(icount) = integral
            eri_index(1,icount) = i
            eri_index(2,icount) = j
            eri_index(3,icount) = k
            eri_index(4,icount) = l
            if (icount == BUFSIZE) then
              rc = trexio_write_mo_2e_int_eri(f, offset, icount, eri_index, eri_buffer)
              call check_success(rc)
              offset += icount
              icount = 0_8
            end if
          end do
        end do
      end do
    end do

    if (icount >= 0_8) then
      rc = trexio_write_mo_2e_int_eri(f, offset, icount, eri_index, eri_buffer)
      call check_success(rc)
    end if
  end if



! One-e RDM
! ---------

  rc = trexio_write_rdm_1e(f,one_e_dm_mo)
  call check_success(rc)

  rc = trexio_write_rdm_1e_up(f,one_e_dm_mo_alpha_average)
  call check_success(rc)

  rc = trexio_write_rdm_1e_dn(f,one_e_dm_mo_beta_average)
  call check_success(rc)


! Two-e RDM
! ---------

  if (export_rdm) then
    PROVIDE two_e_dm_mo

    icount = 0_8
    offset = 0_8
    do l=1,mo_num
      do k=1,mo_num
        do j=1,mo_num
          do i=1,mo_num
            integral = two_e_dm_mo(i,j,k,l)
            if (integral == 0.d0) cycle
            icount += 1_8
            eri_buffer(icount) = integral
            eri_index(1,icount) = i
            eri_index(2,icount) = j
            eri_index(3,icount) = k
            eri_index(4,icount) = l
            if (icount == BUFSIZE) then
              rc = trexio_write_rdm_2e(f, offset, icount, eri_index, eri_buffer)
              call check_success(rc)
              offset += icount
              icount = 0_8
            end if
          end do
        end do
      end do
    end do

    if (icount >= 0_8) then
      rc = trexio_write_rdm_2e(f, offset, icount, eri_index, eri_buffer)
      call check_success(rc)
    end if
  end if


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

  rc = trexio_close(f)
  call check_success(rc)

end


subroutine check_success(rc)
  use trexio
  implicit none
  integer, intent(in)            :: rc
  character*(128)                :: str
  if (rc /= TREXIO_SUCCESS) then
    call trexio_string_of_error(rc,str)
    print *, 'TREXIO Error: ' //trim(str)
    stop -1
  endif
end

! -*- mode: f90 -*-