module sirius_mode
  use iso_c_binding, only: c_ptr
  use f_refcnts, only: f_reference_counter
  use sirius, only: sirius_context_handler, sirius_kpoint_set_handler

  implicit none

  private

  logical, parameter :: sirius_false = .false.

  type, public :: SIRIUS_restart_objects
     type(f_reference_counter) :: refcnt
     type(sirius_context_handler) :: ctx
     type(sirius_kpoint_set_handler) :: kpts
     integer :: niter
  end type SIRIUS_restart_objects

  public :: nullify_SIRIUS_restart_objects, free_SIRIUS_restart_objects, init_SIRIUS_restart_objects
  public :: sirius_ground_state

contains

  pure subroutine nullify_SIRIUS_restart_objects(sirius_rst)
    use iso_c_binding, only: c_null_ptr
    use f_refcnts
    implicit none
    type(SIRIUS_restart_objects), intent(out) :: sirius_rst
    call nullify_f_ref(sirius_rst%refcnt)
    sirius_rst%ctx%handler_ptr_ = c_null_ptr
    sirius_rst%kpts%handler_ptr_ = c_null_ptr
  end subroutine nullify_SIRIUS_restart_objects

  subroutine free_SIRIUS_restart_objects(sirius_rst)
    use sirius
    use f_refcnts
    use iso_c_binding, only: c_associated
    implicit none
    type(SIRIUS_restart_objects), intent(inout) :: sirius_rst
    !check if the object can be freed
    if (f_ref_count(sirius_rst%refcnt) < 0) return
    call f_ref_free(sirius_rst%refcnt)
    if (c_associated(sirius_rst%kpts%handler_ptr_)) then
       call sirius_free_handler(sirius_rst%kpts)
    end if
    if (c_associated(sirius_rst%ctx%handler_ptr_)) then
       call sirius_free_handler(sirius_rst%ctx)
       call sirius_finalize(sirius_false)
    end if
  end subroutine free_SIRIUS_restart_objects

  subroutine init_SIRIUS_restart_objects(sirius_rst, inputs, atoms, mpi_comm)
    use f_refcnts
    use sirius
    use module_input_keys
    use module_atoms
    use f_precisions
    use module_bigdft_arrays
    use abi_defs_basis
    use yaml_output, only: yaml_map, yaml_dict_dump
    implicit none
    type(SIRIUS_restart_objects), intent(inout) :: sirius_rst
    type(input_variables), intent(in) :: inputs
    type(atoms_data), intent(in) :: atoms
    integer, intent(in) :: mpi_comm
    real(f_double), dimension(:,:), pointer :: kpts   !< K points coordinates
    character(len=4096) :: filename
    character(:), allocatable :: dict_content

    real(f_double), dimension(3) :: a1, a2, a3
    integer :: iatype, iat, i, n

    sirius_rst%refcnt = f_ref_new('sirius_rst')

    call sirius_initialize(sirius_false)

    call sirius_create_context(mpi_comm, sirius_rst%ctx)

    call sirius_import_parameters(sirius_rst%ctx, inputs%sirius_config)
    call sirius_set_parameters(sirius_rst%ctx, & !iter_solver_tol = inputs%gnrm_cv, &
         verbosity = inputs%verbosity)
    a1 = (/ atoms%astruct%cell_dim(1), 0._f_double, 0._f_double /)
    a2 = (/ 0._f_double, atoms%astruct%cell_dim(2), 0._f_double /)
    a3 = (/ 0._f_double, 0._f_double, atoms%astruct%cell_dim(3) /)
    call sirius_set_lattice_vectors(sirius_rst%ctx, a1(1), a2(1), a3(1))
    do iatype = 1, atoms%astruct%ntypes
!!$       call sirius_add_atom_type(sirius_rst%ctx, atoms%astruct%atomnames(iatype), &
!!$            zn = atoms%nzatom(iatype), symbol = atoms%astruct%atomnames(iatype), &
!!$            mass = atoms%amu(iatype))
       call sirius_add_atom_type(sirius_rst%ctx, atoms%astruct%atomnames(iatype), &
            fname = trim(atoms%astruct%atomnames(iatype)) // ".json")
    end do
    do iat = 1, atoms%astruct%nat
       iatype = atoms%astruct%iatype(iat)
       call sirius_add_atom(sirius_rst%ctx, atoms%astruct%atomnames(iatype), &
                atoms%astruct%rxyz(1:3, iat)/atoms%astruct%cell_dim(1:3))
    end do
    call sirius_initialize_context(sirius_rst%ctx)

    kpts = f_malloc_ptr(src_ptr=inputs%gen_kpt,id='kpts')
    do i = 1, inputs%gen_nkpt, 1
       kpts(:, i) = inputs%gen_kpt(:, i) * atoms%astruct%cell_dim / two_pi
    end do

    call sirius_create_kset(sirius_rst%ctx, inputs%gen_nkpt, &
         kpts(1:, 1:), inputs%gen_wkpt(1:), sirius_false, sirius_rst%kpts)
    call sirius_initialize_kset(sirius_rst%kpts)
    call f_free_ptr(kpts)
    sirius_rst%niter = inputs%itermax
  end subroutine init_SIRIUS_restart_objects

  subroutine sirius_ground_state(sirius_rst, energy, forces, infocode)
    use sirius
    use module_precisions
    use f_precisions, only: UNINITIALIZED
    use yaml_output
    implicit none
    type(SIRIUS_restart_objects), intent(inout) :: sirius_rst
    real(gp), intent(out) :: energy
    real(gp), dimension(:,:), intent(out) :: forces
    integer, intent(inout) :: infocode
    !local variables
    logical :: stat

    type(sirius_ground_state_handler) :: gs

    energy = UNINITIALIZED(energy)

    call sirius_context_initialized(sirius_rst%ctx, stat)
    if (.not. stat) then
       call yaml_warning("SIRIUS context not initialzed.")
       infocode = 1
       return
    end if

    call sirius_create_ground_state(sirius_rst%kpts, gs)
    call sirius_find_ground_state(gs, niter = sirius_rst%niter)

    infocode = 0
    call sirius_get_energy(gs, "total", energy)
    ! Forces require Scalapack in Sirius, leave it like that for the moment.
    call sirius_get_forces(gs, "total", forces)

    call sirius_free_handler(gs)
  END SUBROUTINE sirius_ground_state

end module sirius_mode
