module QUIP_model_driver_mod

use kim_model_driver_headers_module

use, intrinsic :: iso_c_binding
use kim_model_driver_headers_module, only : &
    kim_log_entry, kim_log_verbosity_error, &
    kim_model_compute_handle_type, kim_model_compute_arguments_handle_type, &
    kim_model_destroy_handle_type, &
    kim_get_argument_pointer, &
    ARG_number_of_particles => kim_compute_argument_name_number_of_particles,  &
    ARG_particle_species_codes => kim_compute_argument_name_particle_species_codes, &
    ARG_particle_contributing => kim_compute_argument_name_particle_contributing, &
    ARG_coordinates => kim_compute_argument_name_coordinates, &
    ARG_partial_energy => kim_compute_argument_name_partial_energy, &
    ARG_partial_particle_energy => kim_compute_argument_name_partial_particle_energy, &
    ARG_partial_forces => kim_compute_argument_name_partial_forces, &
    ARG_partial_virial => kim_compute_argument_name_partial_virial, &
    ARG_partial_particle_virial => kim_compute_argument_name_partial_particle_virial, &
    kim_set_argument_support_status, kim_support_status_optional, &
    kim_length_unit_a, kim_energy_unit_ev, kim_time_unit_unused, kim_charge_unit_unused, &
    kim_temperature_unit_unused, kim_numbering_one_based, &
    kim_get_number_of_parameter_files, &
    kim_get_parameter_file_name, &
    kim_set_influence_distance_pointer, kim_set_neighbor_list_pointers, &
    kim_set_model_buffer_pointer, kim_get_neighbor_list, &
    kim_model_routine_name_compute, &
    kim_model_routine_name_compute_arguments_create, &
    kim_model_routine_name_compute_arguments_destroy, &
    kim_model_routine_name_destroy, &
    kim_from_string

use system_module, only : dp
use periodictable_module, only : atomic_number_from_symbol
use potential_module, only : Potential, potential_filename_initialise, finalise, cutoff, calc
use atoms_types_module, only : Atoms, add_property, assign_property_pointer, &
    add_property_from_pointer
use atoms_module, only : initialise, finalise, print, add_atoms, remove_atoms, set_atoms, &
    get_param_value, remove_property
use connection_module, only : initialise, wipe, add_bond, print
use extendable_str_module, only : extendable_str, string, read, index, substr_replace

implicit none
private

integer, save :: n_initialized = 0

type, bind(C) :: model_buffer_type
    type(c_ptr) :: pot_p
    real(c_double)  :: cutoff(1), influence_distance
    integer(c_int) :: no_need_padding_neighbors(1)
    integer(c_int) :: calc_args_len
    character(kind=c_char) :: calc_args(1024)
end type model_buffer_type

type, bind(C) :: args_buffer_type
    type(c_ptr) :: at_p
end type args_buffer_type

public :: mod_create

contains

recursive subroutine compute(model_compute_handle, &
                              model_compute_arguments_handle, ierr) bind(c)
implicit none
    type(kim_model_compute_handle_type), intent(in) :: model_compute_handle
    type(kim_model_compute_arguments_handle_type), intent(in) :: model_compute_arguments_handle
    integer(c_int), intent(out) :: ierr

    type(model_buffer_type), pointer :: model_buffer; type(c_ptr) pmodel_buffer
    type(args_buffer_type), pointer :: args_buffer; type(c_ptr) pargs_buffer
    type(Atoms), pointer :: at_p
    type(Potential), pointer :: pot_p

    integer(c_int), pointer :: n_at
    real(c_double), pointer :: pos(:,:)

    integer(c_int), pointer :: particleSpecies(:), particleContrib(:)

    real(c_double), pointer :: energy_out, local_energy_out(:), forces_out(:,:)
    real(c_double), pointer :: virial_out(:), local_virial_out(:,:)
    real(dp) :: my_virial_out(3,3)
    real(dp), allocatable :: my_local_virial_out(:,:)

    integer(c_int) :: ierrt
    integer, allocatable :: added_Z(:)
    real(dp), allocatable :: added_pos(:,:)
    integer, allocatable :: remove_indices(:)
    integer :: i_at, jj, j_at, i_remove
    logical, pointer :: local_p(:)

    integer(c_int) :: nn_num
    integer(c_int), pointer :: nn_list(:)
    real(dp) :: lattice(3,3), nn_dist

    integer :: i

    character(len=2048, kind=c_char) :: calc_args

    ! get buffers
    call kim_get_model_buffer_pointer(model_compute_handle, pmodel_buffer)
    call c_f_pointer(pmodel_buffer, model_buffer)
    call kim_get_model_buffer_pointer(model_compute_arguments_handle, &
        pargs_buffer)
    call c_f_pointer(pargs_buffer, args_buffer)
    call c_f_pointer(args_buffer%at_p, at_p)
    call c_f_pointer(model_buffer%pot_p, pot_p)

    ! get numner of atoms
    ierr = 0_c_int
    call kim_get_argument_pointer(model_compute_arguments_handle, &
                     ARG_number_of_particles, n_at, ierrt); ierr = ierr + ierrt
    if (ierr /= 0_c_int) then
        call kim_log_entry(model_compute_arguments_handle, &
            kim_log_verbosity_error, "get_argument_pointer number_of_particles")
        return
    endif
    ! e.g. LAMMPS vacuum region
    !NB intialize quantities that aren't n_at dependent
    if (n_at == 0) then
        return
    endif

    ! get other arguments
    call kim_get_argument_pointer(model_compute_arguments_handle, &
                     ARG_particle_species_codes, n_at, particleSpecies, ierrt); ierr = ierr + ierrt
    call kim_get_argument_pointer(model_compute_arguments_handle, &
                     ARG_particle_contributing, n_at, particleContrib, ierrt); ierr = ierr + ierrt
    call kim_get_argument_pointer(model_compute_arguments_handle, &
                     ARG_coordinates, 3, n_at, pos, ierrt); ierr = ierr + ierrt
    call kim_get_argument_pointer(model_compute_arguments_handle, &
                     ARG_partial_energy, energy_out, ierrt); ierr = ierr + ierrt
    call kim_get_argument_pointer(model_compute_arguments_handle, &
                     ARG_partial_particle_energy, n_at, local_energy_out, ierrt); ierr = ierr + ierrt
    call kim_get_argument_pointer(model_compute_arguments_handle, &
                     ARG_partial_virial, 6, virial_out, ierrt); ierr = ierr + ierrt
    call kim_get_argument_pointer(model_compute_arguments_handle, &
                     ARG_partial_particle_virial, 6, n_at, local_virial_out, ierrt); ierr = ierr + ierrt
    call kim_get_argument_pointer(model_compute_arguments_handle, &
                     ARG_partial_forces, 3, n_at, forces_out, ierrt); ierr = ierr + ierrt
    if (ierr /= 0_c_int) then
        call kim_log_entry(model_compute_arguments_handle, &
            kim_log_verbosity_error, "get_argument_pointer most quantities")
        return
    endif

    ! resize atoms object
    if (n_at > at_p%n) then
        allocate(added_Z(n_at-at_p%n))
        added_Z = 1
        allocate(added_pos(3,n_at-at_p%n))
        call add_atoms(at_p, added_pos, added_Z)
        deallocate(added_Z, added_pos)
        call initialise(at_p%connect, n_at, Nbuffer=n_at*15)
    else if (n_at < at_p%n) then
        allocate(remove_indices(at_p%n - n_at))
        do i_remove=n_at+1, at_p%n
            remove_indices(i_remove-n_at) = i_remove
        end do
        call remove_atoms(at_p, remove_indices)
        deallocate(remove_indices)
        call initialise(at_p%connect, n_at, Nbuffer=n_at*15)
    else
        call wipe(at_p%connect)
    endif

    ! set types
    call set_atoms(at_p, particleSpecies(1:n_at))
    ! set positions
    at_p%pos(1:3,1:n_at) = pos(1:3,1:n_at)
    ! set local status
    call assign_property_pointer(at_p, "local", local_p)
    local_p = (particleContrib /= 0)

    ! set Atoms object cutoff in case potential checks explicitly
    at_p%cutoff = model_buffer%cutoff(1)

    ! create neighbor list
    !!NB let's try this the clean way first (add_bond), without directly messing with internal data structures
    lattice = 0.0
!!DEBUG print *, "quip make nn list"
    do i_at=1, n_at
        if (particleContrib(i_at) /= 0 .or. model_buffer%no_need_padding_neighbors(1) == 0) then
            call kim_get_neighbor_list( model_compute_arguments_handle, &
                1, i_at, nn_num, nn_list, ierr)
            if (ierr /= 0_c_int) then
                call kim_log_entry(model_compute_arguments_handle, &
                    kim_log_verbosity_error, "get_argument_pointer neighbor_list")
                return
            endif
!!DEBUG print *, "quip make nn list i", i_at, nn_num, "contrib", particleContrib(i_at), "pos i", at_p%pos(:,i_at)
            do jj = 1, nn_num
                j_at = nn_list(jj)
                if (j_at < i_at .or. (model_buffer%no_need_padding_neighbors(1) == 1 .and. &
                    particleContrib(i_at) /= particleContrib(j_at))) then
!!DEBUG print *, "    j_at ", j_at, "pos j", at_p%pos(:,j_at)
                    nn_dist = sqrt(sum((at_p%pos(:,i_at)-at_p%pos(:,j_at))**2))
                    call add_bond(at_p%connect, at_p%pos, &
                        lattice, i_at, j_at, (/ 0,0,0 /), nn_dist)
!!DEBUG else
!!DEBUG print *, "    NO j_at ", j_at, "pos j", at_p%pos(:,j_at)
                endif
            end do
        endif
    end do
    lattice(1,1) = maxval(pos(1,1:n_at)) - minval(pos(1,1:n_at)) + 2.0*model_buffer%cutoff(1)
    lattice(2,2) = maxval(pos(2,1:n_at)) - minval(pos(2,1:n_at)) + 2.0*model_buffer%cutoff(1)
    lattice(3,3) = maxval(pos(3,1:n_at)) - minval(pos(3,1:n_at)) + 2.0*model_buffer%cutoff(1)

    calc_args = ""
    do i=1, model_buffer%calc_args_len
        calc_args(i:i) = model_buffer%calc_args(i)
    end do
    calc_args = trim(calc_args) // " atom_mask_name=local do_calc_connect=F"
    ! things that will be returned as parameters
    if (associated(energy_out)) then
        calc_args = trim(calc_args)//" energy"
    endif
    if (associated(virial_out)) then
        calc_args = trim(calc_args)//" virial"
    endif
    ! things that wil be return as properties, but have matching shape
    ! so can use KIM's allocation
    if (associated(local_energy_out)) then
        calc_args = trim(calc_args)//" local_energy"
        call add_property_from_pointer(at_p, "local_energy", local_energy_out)
    endif
    if (associated(forces_out)) then
        calc_args = trim(calc_args)//" force"
        call add_property_from_pointer(at_p, "force", forces_out)
    endif
    ! things that wil be return as properties, but need to be allocated locally
    if (associated(local_virial_out)) then
        allocate(my_local_virial_out(9,n_at))
        calc_args = trim(calc_args)//" local_virial"
        call add_property_from_pointer(at_p, "local_virial", my_local_virial_out)
    endif

    call calc(pot_p, at_p, args_str=trim(calc_args))

    ! read back things that are returned as parameters
    if (associated(energy_out)) then
        call get_param_value(at_p, "energy", energy_out)
    endif
    if (associated(virial_out)) then
        call get_param_value(at_p, "virial", my_virial_out)
        virial_out(1) = my_virial_out(1,1)
        virial_out(2) = my_virial_out(2,2)
        virial_out(3) = my_virial_out(3,3)
        virial_out(4) = my_virial_out(2,3)
        virial_out(5) = my_virial_out(1,3)
        virial_out(6) = my_virial_out(1,2)
    endif
    ! things where storage was externally allocated don't need to be copied
    ! but association should be removed
    if (associated(local_energy_out)) then
        call remove_property(at_p, "local_energy")
    endif
    if (associated(forces_out)) then
        call remove_property(at_p, "force")
    endif
    ! read back things that are properties but were locally allocated
    ! then deallocate local copy (maybe buffer?)
    if (associated(local_virial_out)) then
        local_virial_out(1,:) = my_local_virial_out(1,:)
        local_virial_out(2,:) = my_local_virial_out(5,:)
        local_virial_out(3,:) = my_local_virial_out(9,:)
        local_virial_out(4,:) = my_local_virial_out(6,:)
        local_virial_out(5,:) = my_local_virial_out(3,:)
        local_virial_out(6,:) = my_local_virial_out(2,:)
        call remove_property(at_p, "local_virial")
        deallocate(my_local_virial_out)
    endif

end subroutine compute

recursive subroutine destroy(model_destroy_handle, ierr) bind(c)
use system_module, only : system_finalise
    type(kim_model_destroy_handle_type), intent(inout) :: model_destroy_handle
    integer(c_int), intent(out) :: ierr

    type(model_buffer_type), pointer :: model_buffer; type(c_ptr) :: pmodel_buffer
    type(Potential), pointer :: pot_p

    call kim_get_model_buffer_pointer(model_destroy_handle, pmodel_buffer)
    call c_f_pointer(pmodel_buffer, model_buffer)
    call c_f_pointer(model_buffer%pot_p, pot_p)

    call finalise(pot_p)
    deallocate(pot_p)
    deallocate(model_buffer)

    n_initialized = n_initialized - 1
    if (n_initialized == 0) call system_finalise()

    ierr = 0_c_int
end subroutine destroy

recursive subroutine compute_arguments_create(model_compute_handle, &
    model_compute_arguments_create_handle, ierr) bind(c)
implicit none
    type(kim_model_compute_handle_type), intent(in) :: model_compute_handle
    type(kim_model_compute_arguments_create_handle_type), intent(inout) :: &
      model_compute_arguments_create_handle
    integer(c_int), intent(out) :: ierr

    type(args_buffer_type), pointer :: args_buffer
    type(Atoms), pointer :: at_p
    integer(c_int) :: ierrt
    real(dp) :: lattice(3,3)

    ! avoid unused dummy argument warnings
    if (model_compute_handle .eq. KIM_MODEL_COMPUTE_NULL_HANDLE) continue

    ! register arguments
    ierr = 0_c_int
    call kim_set_argument_support_status( model_compute_arguments_create_handle, &
        arg_partial_energy, kim_support_status_optional, ierrt); ierr = ierr + ierrt
    call kim_set_argument_support_status( model_compute_arguments_create_handle, &
        arg_partial_particle_energy, kim_support_status_optional, ierrt); ierr = ierr + ierrt
    call kim_set_argument_support_status( model_compute_arguments_create_handle, &
        arg_partial_virial, kim_support_status_optional, ierrt); ierr = ierr + ierrt
    call kim_set_argument_support_status( model_compute_arguments_create_handle, &
        arg_partial_particle_virial, kim_support_status_optional, ierrt); ierr = ierr + ierrt
    call kim_set_argument_support_status( model_compute_arguments_create_handle, &
        arg_partial_forces, kim_support_status_optional, ierrt); ierr = ierr + ierrt
    if (ierr /= 0_c_int) then
      call kim_log_entry(model_compute_arguments_create_handle, &
          kim_log_verbosity_error, "Unable to register arguments support_status")
      return
    end if

    allocate(args_buffer)
    allocate(at_p)
    args_buffer%at_p = c_loc(at_p)
    lattice = 0.0_dp
    lattice(1,1) = 1.0_dp
    lattice(2,2) = 1.0_dp
    lattice(3,3) = 1.0_dp
    call initialise(at_p, 1, lattice)
    call add_property(at_p,'local',.false.)

    call initialise(at_p%connect, 1, Nbuffer=1*15)

    call kim_set_model_buffer_pointer( &
        model_compute_arguments_create_handle, c_loc(args_buffer))

end subroutine compute_arguments_create

recursive subroutine compute_arguments_destroy(model_compute_handle, &
  model_compute_arguments_destroy_handle, ierr) bind(c)
implicit none
    type(kim_model_compute_handle_type), intent(in) :: model_compute_handle
    type(kim_model_compute_arguments_destroy_handle_type), intent(inout) :: &
      model_compute_arguments_destroy_handle
    integer(c_int), intent(out) :: ierr

    type(args_buffer_type), pointer :: args_buffer; type(c_ptr) pargs_buffer
    type(Atoms), pointer :: at_p

    ! avoid unused dummy argument warnings
    if (model_compute_handle .eq. KIM_MODEL_COMPUTE_NULL_HANDLE) continue

    call kim_get_model_buffer_pointer( &
        model_compute_arguments_destroy_handle, pargs_buffer)
    call c_f_pointer(pargs_buffer, args_buffer)
    call c_f_pointer(args_buffer%at_p, at_p)

    call finalise(at_p)
    deallocate(at_p)
    deallocate(args_buffer)

    ierr = 0_c_int

end subroutine compute_arguments_destroy

recursive subroutine mod_create(model_driver_create_handle, &
  requested_length_unit, requested_energy_unit, requested_charge_unit, &
  requested_temperature_unit, requested_time_unit, ierr) bind(c)
use, intrinsic :: iso_c_binding
use system_module, only : dp, system_initialise, PRINT_SILENT, system_abort, operator(//)
use kim_model_driver_headers_module
implicit none
    type(kim_model_driver_create_handle_type), intent(inout) &
      :: model_driver_create_handle
    type(kim_length_unit_type), intent(in), value :: requested_length_unit
    type(kim_energy_unit_type), intent(in), value :: requested_energy_unit
    type(kim_charge_unit_type), intent(in), value :: requested_charge_unit
    type(kim_temperature_unit_type), intent(in), value :: requested_temperature_unit
    type(kim_time_unit_type), intent(in), value :: requested_time_unit
    integer(c_int), intent(out) :: ierr

    type(Potential), pointer :: pot_p

    integer(c_int) :: ierrt
    integer(c_int) :: number_of_parameter_files
    character(len=1024, kind=c_char) :: parameter_file_name, xml_file_name
    character(len=1024, kind=c_char) :: pot_init_args, pot_calc_args
    real(c_double) :: r_cut_influence_multiplier
    integer(c_int) :: padding_neighbor_hint
    integer(c_int) :: n_species, i_species, i_file
    character(len=2, kind=c_char), allocatable :: species_list(:)
    type(kim_species_name_type) species_name

    type(model_buffer_type), pointer :: model_buffer
    type(extendable_str) :: params_es

    integer :: i

    if (n_initialized == 0) call system_initialise(verbosity=PRINT_SILENT)
    n_initialized = n_initialized + 1

    call kim_set_units( &
        model_driver_create_handle, &
        kim_length_unit_a, &
        kim_energy_unit_ev, &
        kim_charge_unit_unused, &
        kim_temperature_unit_unused, &
        kim_time_unit_unused, ierr)
    if (ierr /= 0_c_int) then
      call kim_log_entry(model_driver_create_handle, &
          kim_log_verbosity_error, "Unable to set units")
        return
    end if

    ! register numbering
    call kim_set_model_numbering( &
        model_driver_create_handle, kim_numbering_one_based, ierr)
    if (ierr /= 0_c_int) then
      call kim_log_entry(model_driver_create_handle, &
          kim_log_verbosity_error, "Unable to set numbering")
        return
    end if

    ! store callback pointers in KIM object
    ierr = 0_c_int
    call kim_set_routine_pointer( &
        model_driver_create_handle, kim_model_routine_name_compute, &
        kim_language_name_fortran, 1_c_int, &
        c_funloc(compute), ierrt); ierr = ierr + ierrt
    call kim_set_routine_pointer( &
        model_driver_create_handle, kim_model_routine_name_compute_arguments_create, &
        kim_language_name_fortran, 1_c_int, &
        c_funloc(compute_arguments_create), ierrt); ierr = ierr + ierrt
    call kim_set_routine_pointer( &
        model_driver_create_handle, kim_model_routine_name_compute_arguments_destroy, &
        kim_language_name_fortran, 1_c_int, &
        c_funloc(compute_arguments_destroy), ierrt); ierr = ierr + ierrt
    call kim_set_routine_pointer( &
        model_driver_create_handle, kim_model_routine_name_destroy, &
        kim_language_name_fortran, 1_c_int, &
        c_funloc(destroy), ierrt); ierr = ierr + ierrt
    if (ierr /= 0_c_int) then
      call kim_log_entry(model_driver_create_handle, &
          kim_log_verbosity_error, "Unable to store callback pointers")
        return
    end if

    ! process parameter files
    call kim_get_number_of_parameter_files( &
        model_driver_create_handle, number_of_parameter_files)
    if (number_of_parameter_files < 2) then
        call kim_log_entry(model_driver_create_handle, &
          kim_log_verbosity_error, "wrong number of parameter files")
        return
    end if

    ! get master filename
    call kim_get_parameter_file_name( model_driver_create_handle, &
        1, parameter_file_name, ierr)

    ! read info from master filename
    open(10,file=parameter_file_name,status="old")
    read(10,'(A)',iostat=ierr,err=100) pot_init_args
    read(10,*,iostat=ierr,err=100) n_species
    allocate(species_list(n_species))
    read(10,*,iostat=ierr,err=100) species_list
    read(10,'(A)',iostat=ierr,err=100) pot_calc_args
    read(10,*,iostat=ierr,err=100) r_cut_influence_multiplier
    read(10,*,iostat=ierr,err=100) padding_neighbor_hint
    close(10)

    ! register species
    do i_species=1, n_species
        ! register species
        !!NB check for "uknown" return value?
        call kim_from_string(species_list(i_species), species_name)
        call kim_set_species_code( &
            model_driver_create_handle, species_name, &
            atomic_number_from_symbol(species_list(i_species)), ierr)
        if (ierr /= 0_c_int) then
            call kim_log_entry(model_driver_create_handle, &
              kim_log_verbosity_error, "Unable to set species code "//trim(species_list(i_species)))
            return
        end if
    end do
    deallocate(species_list)

    ! get name of xml file
    call kim_get_parameter_file_name( model_driver_create_handle, &
        2, xml_file_name, ierr)
    ! read xml file and mangle sub file names
    allocate(model_buffer)
    ! actually initialise
    allocate(pot_p)
    model_buffer%pot_p = c_loc(pot_p)
    call potential_filename_initialise(pot_p, args_str=pot_init_args, param_filename=xml_file_name)

    if (pot_calc_args == '_NONE_') then
        model_buffer%calc_args_len = 0
    else
        if (len_trim(pot_calc_args) >= 1024) then
            call kim_log_entry(model_driver_create_handle, &
              kim_log_verbosity_error, "pot_calc_args too long")
            return
        endif
        model_buffer%calc_args_len = len_trim(pot_calc_args)
        do i=1, len_trim(pot_calc_args)
            model_buffer%calc_args(i) = pot_calc_args(i:i)
        end do
    endif
    model_buffer%cutoff(1) = cutoff(pot_p)
!!DEBUG print *, "setting cutoff ", model_buffer%cutoff(1)
    model_buffer%influence_distance = model_buffer%cutoff(1) * r_cut_influence_multiplier
    model_buffer%no_need_padding_neighbors(1) = padding_neighbor_hint

    call kim_set_influence_distance_pointer( model_driver_create_handle, &
        model_buffer%influence_distance)
    call kim_set_neighbor_list_pointers( model_driver_create_handle, &
        1, model_buffer%cutoff, model_buffer%no_need_padding_neighbors)
    call kim_set_model_buffer_pointer( model_driver_create_handle, c_loc(model_buffer))

    return

100 continue
    call kim_log_entry(model_driver_create_handle, &
      kim_log_verbosity_error, "Unable to read potential initialization parameters")

end subroutine mod_create

end module QUIP_model_driver_mod

recursive subroutine create(model_driver_create_handle, &
  requested_length_unit, requested_energy_unit, requested_charge_unit, &
  requested_temperature_unit, requested_time_unit, ierr) bind(c)

    use kim_model_driver_headers_module
    use QUIP_model_driver_mod, only : mod_create

    type(kim_model_driver_create_handle_type), intent(inout) &
      :: model_driver_create_handle
    type(kim_length_unit_type), intent(in), value :: requested_length_unit
    type(kim_energy_unit_type), intent(in), value :: requested_energy_unit
    type(kim_charge_unit_type), intent(in), value :: requested_charge_unit
    type(kim_temperature_unit_type), intent(in), value :: requested_temperature_unit
    type(kim_time_unit_type), intent(in), value :: requested_time_unit
    integer(c_int), intent(out) :: ierr

    call mod_create(model_driver_create_handle, &
      requested_length_unit, requested_energy_unit, requested_charge_unit, &
      requested_temperature_unit, requested_time_unit, ierr)
end subroutine