module restraints_constraints_xml_module
use system_module, only : dp, system_abort, operator(//)
use constraints_module
use dynamicalsystem_module
use quip_common_module
implicit none
private

type(DynamicalSystem), pointer, private :: parse_ds
logical, private :: parse_in_restraints = .false., parse_in_constraints = .false.

public :: init_restraints_constraints

contains

   subroutine restraints_constraints_startElement_handler(URI, localname, name, attributes)
      character(len=*), intent(in)   :: URI
      character(len=*), intent(in)   :: localname
      character(len=*), intent(in)   :: name
      type(dictionary_t), intent(in) :: attributes

      character(len=1024) :: value
      integer :: status, di_status, tau_status, t0_status
      integer :: n
      integer :: atom_1, atom_2, atom_3, Z
      real(dp) :: k
      character(len=1024) :: bound_str
      integer :: bound
      real(dp) :: c
      real(dp) :: d, plane_n(3), di, p, q(3), SF
      real(dp) :: t0, tau
      real(dp) :: egap
      real(dp) :: tol
      character(len=20) :: type_str
      logical :: print_summary

      if (parse_in_restraints) type_str = "restraint"
      if (parse_in_constraints) type_str = "constraint"

      if (name == 'restraints' .or. name == 'constraints') then ! new restraints stanze
         if (name == 'restraints') then
            parse_in_restraints = .true.
            parse_in_constraints = .false.
            call QUIP_FoX_get_value(attributes, "N", value, status)
            if (status /= 0) call system_abort("restraint_startElement_handler failed to read N in restraints stanza")
            read (value, *) n
            call print("allocating array for "//n//" restraints")
            if (allocated(parse_ds%restraint)) deallocate(parse_ds%restraint)
            allocate(parse_ds%restraint(n))
         else ! constraints
            parse_in_constraints = .true.
            parse_in_restraints = .false.
            call QUIP_FoX_get_value(attributes, "N", value, status)
            if (status /= 0) call system_abort("restraint_startElement_handler failed to read N in constraints stanza")
            read (value, *) n
            call print("allocating array for "//n//" constraints")
            if (allocated(parse_ds%constraint)) deallocate(parse_ds%constraint)
            allocate(parse_ds%constraint(n))
         endif

      else if (parse_in_restraints .or. parse_in_constraints) then
         call QUIP_FoX_get_value(attributes, "print_summary", value, status)
         if (status == 0) then
            read (value, *) print_summary
         else
            print_summary=.true.
         endif

         if (parse_in_restraints) then
            call QUIP_FoX_get_value(attributes, "k", value, status)
            if (status /= 0) call system_abort("restraint_startElement_handler failed to read k in bond_length "//trim(type_str))
            read (value, *) k
            call QUIP_FoX_get_value(attributes, "bound", value, status)
            if (status == 0) then
               read (value, *) bound_str
               call interpret_bound_string(trim(bound_str),bound)
            else
               bound = BOTH_UPPER_AND_LOWER_BOUNDS
            endif
         else
            call QUIP_FoX_get_value(attributes, "tol", value, status) !constraint tolerance
            if (status == 0) then
               read (value, *) tol
            else
               tol=-1._dp !default
            endif
         endif

         if (name == 'bond_length') then
            call QUIP_FoX_get_value(attributes, "atom_1", value, status)
            if (status /= 0) call system_abort("restraint_startElement_handler failed to read atom_1 in bond_length "// &
              trim(type_str))
            read (value, *) atom_1
            call QUIP_FoX_get_value(attributes, "atom_2", value, status)
            if (status /= 0) call system_abort("restraint_startElement_handler failed to read atom_2 in bond_length "// &
              trim(type_str))
            read (value, *) atom_2
            call QUIP_FoX_get_value(attributes, "d", value, status)
            if (status /= 0) call system_abort("restraint_startElement_handler failed to read d in bond_length "//trim(type_str))
            read (value, *) d

            tau = -1.0_dp
            di = -1.0_dp
            call QUIP_FoX_get_value(attributes, "tau", value, status)
            if (status == 0) then
               read (value, *) tau
            endif
            call QUIP_FoX_get_value(attributes, "di", value, status)
            if (status == 0) then
               read (value, *) di
            endif
            call QUIP_FoX_get_value(attributes, "t0", value, status)
            if (status == 0) then
               read (value, *) t0
            endif
            if ((status == 0 .and. (tau <= 0.0_dp .or. di < 0.0_dp)) .or. &
                (status /= 0 .and. (tau > 0.0_dp .or. di >= 0.0_dp))) then
              call system_abort("got t0, but tau or di is missing or invalid")
            endif

            if (tau > 0.0_dp) then
               if (parse_in_restraints) then
                  call constrain_bondlength(parse_ds, atom_1, atom_2, d, di=di, t0=t0, tau=tau, restraint_k=k, bound=bound, print_summary=print_summary)
               else
                  call constrain_bondlength(parse_ds, atom_1, atom_2, d, di=di, t0=t0, tau=tau, tol=tol, print_summary=print_summary)
               endif
            else
               if (parse_in_restraints) then
                  call constrain_bondlength(parse_ds, atom_1, atom_2, d, restraint_k=k, bound=bound, print_summary=print_summary)
               else
                  call constrain_bondlength(parse_ds, atom_1, atom_2, d, tol=tol, print_summary=print_summary)
               endif
            endif

         else if (name == 'bond_length_dev_pow') then
            call QUIP_FoX_get_value(attributes, "atom_1", value, status)
            if (status /= 0) call system_abort("restraint_startElement_handler failed to read atom_1 in bond_length_dev_pow "// &
              trim(type_str))
            read (value, *) atom_1
            call QUIP_FoX_get_value(attributes, "atom_2", value, status)
            if (status /= 0) call system_abort("restraint_startElement_handler failed to read atom_2 in bond_length_dev_pow "// &
              trim(type_str))
            read (value, *) atom_2
            call QUIP_FoX_get_value(attributes, "p", value, status)
            if (status /= 0) call system_abort("restraint_startElement_handler failed to read p in bond_length_dev_pow "//trim(type_str))
            read (value, *) p
            call QUIP_FoX_get_value(attributes, "d", value, status)
            if (status /= 0) call system_abort("restraint_startElement_handler failed to read d in bond_length_dev_pow "//trim(type_str))
            read (value, *) d

            tau = -1.0_dp
            di = -1.0_dp
            call QUIP_FoX_get_value(attributes, "tau", value, status)
            if (status == 0) then
               read (value, *) tau
            endif
            call QUIP_FoX_get_value(attributes, "di", value, status)
            if (status == 0) then
               read (value, *) di
            endif
            call QUIP_FoX_get_value(attributes, "t0", value, status)
            if (status == 0) then
               read (value, *) t0
            endif
            if ((status == 0 .and. (tau <= 0.0_dp .or. di < 0.0_dp)) .or. &
                (status /= 0 .and. (tau > 0.0_dp .or. di >= 0.0_dp))) then
              call system_abort("got t0, but tau or di is missing or invalid")
            endif

            if (tau > 0.0_dp) then
               if (parse_in_restraints) then
                  call constrain_bondlength_dev_pow(parse_ds, atom_1, atom_2, p, d, di=di, t0=t0, tau=tau, restraint_k=k, bound=bound, print_summary=print_summary)
               else
                  call constrain_bondlength_dev_pow(parse_ds, atom_1, atom_2, p, d, di=di, t0=t0, tau=tau, tol=tol, print_summary=print_summary)
               endif
            else
               if (parse_in_restraints) then
                  call constrain_bondlength_dev_pow(parse_ds, atom_1, atom_2, p, d, restraint_k=k, bound=bound, print_summary=print_summary)
               else
                  call constrain_bondlength_dev_pow(parse_ds, atom_1, atom_2, p, d, tol=tol, print_summary=print_summary)
               endif
            endif

         else if (name == 'bond_length_sq') then

            call QUIP_FoX_get_value(attributes, "atom_1", value, status)
            if (status /= 0) call system_abort("restraint_startElement_handler failed to read atom_1 in bond_length_sq "//trim(type_str))
            read (value, *) atom_1
            call QUIP_FoX_get_value(attributes, "atom_2", value, status)
            if (status /= 0) call system_abort("restraint_startElement_handler failed to read atom_2 in bond_length_sq "//trim(type_str))
            read (value, *) atom_2
            call QUIP_FoX_get_value(attributes, "d", value, status)
            if (status == 0) then
               read (value, *) d
               if (parse_in_restraints) then
                  call constrain_bondlength_sq(parse_ds, atom_1, atom_2, d, restraint_k=k, bound=bound, print_summary=print_summary)
               else
                  call constrain_bondlength_sq(parse_ds, atom_1, atom_2, d, tol=tol, print_summary=print_summary)
               endif
            else
               if (parse_in_restraints) then
                  call constrain_bondlength_sq(parse_ds, atom_1, atom_2, restraint_k=k, bound=bound, print_summary=print_summary)
               else
                  call constrain_bondlength_sq(parse_ds, atom_1, atom_2, tol=tol, print_summary=print_summary)
               endif
            endif

         else if (name == 'bond_angle_cos') then

            call QUIP_FoX_get_value(attributes, "atom_1", value, status)
            if (status /= 0) call system_abort("restraint_startElement_handler failed to read atom_1 in bond_angle_cos "//trim(type_str))
            read (value, *) atom_1
            call QUIP_FoX_get_value(attributes, "atom_2", value, status)
            if (status /= 0) call system_abort("restraint_startElement_handler failed to read atom_2 in bond_angle_cos "//trim(type_str))
            read (value, *) atom_2
            call QUIP_FoX_get_value(attributes, "atom_3", value, status)
            if (status /= 0) call system_abort("restraint_startElement_handler failed to read atom_3 in bond_angle_cos "//trim(type_str))
            read (value, *) atom_3
            call QUIP_FoX_get_value(attributes, "c", value, status)
            if (status == 0) then
               read (value, *) c
               if (parse_in_restraints) then
                  call constrain_bondanglecos(parse_ds, atom_1, atom_2, atom_3, c, restraint_k=k, bound=bound, print_summary=print_summary)
               else
                  call constrain_bondanglecos(parse_ds, atom_1, atom_2, atom_3, c, tol=tol, print_summary=print_summary)
               endif
            else
               if (parse_in_restraints) then
                  call constrain_bondanglecos(parse_ds, atom_1, atom_2, atom_3, restraint_k=k, bound=bound, print_summary=print_summary)
               else
                  call constrain_bondanglecos(parse_ds, atom_1, atom_2, atom_3, tol=tol, print_summary=print_summary)
               endif
            endif

         else if (name == 'bond_length_diff') then

            call QUIP_FoX_get_value(attributes, "atom_1", value, status)
            if (status /= 0) call system_abort("restraint_startElement_handler failed to read atom_1 in bond_length_diff "//trim(type_str))
            read (value, *) atom_1
            call QUIP_FoX_get_value(attributes, "atom_2", value, status)
            if (status /= 0) call system_abort("restraint_startElement_handler failed to read atom_2 in bond_length_diff "//trim(type_str))
            read (value, *) atom_2
            call QUIP_FoX_get_value(attributes, "atom_3", value, status)
            if (status /= 0) call system_abort("restraint_startElement_handler failed to read atom_3 in bond_length_diff "//trim(type_str))
            read (value, *) atom_3
            call QUIP_FoX_get_value(attributes, "d", value, status)
            if (status /= 0) call system_abort("restraint_startElement_handler failed to read d in bond_length_diff "//trim(type_str))
            read (value, *) d

            call QUIP_FoX_get_value(attributes, "di", value, di_status)
            if (di_status == 0) read (value, *) di
            call QUIP_FoX_get_value(attributes, "tau", value, tau_status)
            if (tau_status == 0) read (value, *) tau
            call QUIP_FoX_get_value(attributes, "t0", value, t0_status)
            if (t0_status == 0) read (value, *) t0
            if (count( (/ di_status, tau_status, t0_status /) == 0) /= 0 .and. &
                count( (/ di_status, tau_status, t0_status /) == 0) /= 3) then
               call system_abort("restraint_startElement_handler needs either all or none of di, tau, t0")
            endif

            if (di_status == 0) then ! relax target value with time
               if (parse_in_restraints) then
                  call constrain_bondlength_diff(parse_ds, atom_1, atom_2, atom_3, d, di=di, tau=tau, t0=t0, restraint_k=k, bound=bound, print_summary=print_summary)
               else
                  call constrain_bondlength_diff(parse_ds, atom_1, atom_2, atom_3, d, di=di, tau=tau, t0=t0, tol=tol, print_summary=print_summary)
               endif
            else ! fixed target value
               if (parse_in_restraints) then
                  call constrain_bondlength_diff(parse_ds, atom_1, atom_2, atom_3, d, restraint_k=k, bound=bound, print_summary=print_summary)
               else
                  call constrain_bondlength_diff(parse_ds, atom_1, atom_2, atom_3, d, tol=tol, print_summary=print_summary)
               endif
            endif

         else if (name == 'gap_energy') then

            call QUIP_FoX_get_value(attributes, "egap", value, status)
            if (status == 0) then
               read (value, *) egap
               if (parse_in_restraints) then
                  call constrain_gap_energy(parse_ds, egap, restraint_k=k, bound=bound, print_summary=print_summary)
               else
                  call constrain_gap_energy(parse_ds, egap, tol=tol, print_summary=print_summary)
               endif
            else
               call system_abort("restraint_startElement_handler failed to read egap in gap_energy "//trim(type_str))
               !if (parse_in_restraints) then
               !  call constrain_gap_energy(parse_ds, restraint_k=k, print_summary=print_summary)
               !else
               !  call constrain_gap_energy(parse_ds, tol=tol, print_summary=print_summary)
               !endif
            endif

         else if (name == 'atom_plane') then

            call QUIP_FoX_get_value(attributes, "atom", value, status)
            if (status /= 0) call system_abort("restraint_startElement_handler failed to read atom in atom_plane "//trim(type_str))
            read (value, *) atom_1
            call QUIP_FoX_get_value(attributes, "normal", value, status)
            if (status /= 0) call system_abort("restraint_startElement_handler failed to read normal in atom_plane "//trim(type_str))
            read (value, *) plane_n
            call QUIP_FoX_get_value(attributes, "d", value, status)
            if (status == 0) then
               read (value, *) d
               if (parse_in_restraints) then
                  call constrain_atom_plane(parse_ds, atom_1, plane_n, d, restraint_k=k, bound=bound, print_summary=print_summary)
               else
                  call constrain_atom_plane(parse_ds, atom_1, plane_n, d, tol=tol, print_summary=print_summary)
               endif
            else
               if (parse_in_restraints) then
                  call constrain_atom_plane(parse_ds, atom_1, plane_n, restraint_k=k, bound=bound, print_summary=print_summary)
               else
                  call constrain_atom_plane(parse_ds, atom_1, plane_n, tol=tol, print_summary=print_summary)
               endif
            endif

         else if (name == 'struct_factor_like_mag') then

            call QUIP_FoX_get_value(attributes, "Z", value, status)
            if (status /= 0) call system_abort("restraint_startElement_handler failed to read Z in struct_factor_like_mag "//trim(type_str))
            read (value, *) Z
            call QUIP_FoX_get_value(attributes, "q", value, status)
            if (status /= 0) call system_abort("restraint_startElement_handler failed to read q in struct_factor_like_mag "//trim(type_str))
            read (value, *) q
            call QUIP_FoX_get_value(attributes, "SF", value, status)
            if (status /= 0) call system_abort("restraint_startElement_handler failed to read SF in struct_factor_like_mag "//trim(type_str))
            read (value, *) SF
            if (parse_in_restraints) then
               call constrain_struct_factor_like_mag(parse_ds, Z, q, SF, restraint_k=k, bound=bound, print_summary=print_summary)
            else
               call constrain_struct_factor_like_mag(parse_ds, Z, q, SF, tol=tol, print_summary=print_summary)
            endif

         else if (name == 'struct_factor_like_r') then

            call QUIP_FoX_get_value(attributes, "Z", value, status)
            if (status /= 0) call system_abort("restraint_startElement_handler failed to read Z in struct_factor_like_r "//trim(type_str))
            read (value, *) Z
            call QUIP_FoX_get_value(attributes, "q", value, status)
            if (status /= 0) call system_abort("restraint_startElement_handler failed to read q in struct_factor_like_r "//trim(type_str))
            read (value, *) q
            call QUIP_FoX_get_value(attributes, "SF", value, status)
            if (status /= 0) call system_abort("restraint_startElement_handler failed to read SF in struct_factor_like_r "//trim(type_str))
            read (value, *) SF
            if (parse_in_restraints) then
               call constrain_struct_factor_like_r(parse_ds, Z, q, SF, restraint_k=k, bound=bound, print_summary=print_summary)
            else
               call constrain_struct_factor_like_r(parse_ds, Z, q, SF, tol=tol, print_summary=print_summary)
            endif

         else if (name == 'struct_factor_like_i') then

            call QUIP_FoX_get_value(attributes, "Z", value, status)
            if (status /= 0) call system_abort("restraint_startElement_handler failed to read Z in struct_factor_like_i "//trim(type_str))
            read (value, *) Z
            call QUIP_FoX_get_value(attributes, "q", value, status)
            if (status /= 0) call system_abort("restraint_startElement_handler failed to read q in struct_factor_like_i "//trim(type_str))
            read (value, *) q
            call QUIP_FoX_get_value(attributes, "SF", value, status)
            if (status /= 0) call system_abort("restraint_startElement_handler failed to read SF in struct_factor_like_i "//trim(type_str))
            read (value, *) SF
            if (parse_in_restraints) then
               call constrain_struct_factor_like_i(parse_ds, Z, q, SF, restraint_k=k, bound=bound, print_summary=print_summary)
            else
               call constrain_struct_factor_like_i(parse_ds, Z, q, SF, tol=tol, print_summary=print_summary)
            endif

         else
            call system_abort("unknown " // trim(type_str) //" type '"//name//"' in " // trim(type_str) //"s xml stanza")
         endif
      endif

   end subroutine

   subroutine restraints_constraints_endElement_handler(URI, localname, name)
      character(len=*), intent(in)   :: URI
      character(len=*), intent(in)   :: localname
      character(len=*), intent(in)   :: name

      if (name == 'restraints') then ! end of restraints stanza
         parse_in_restraints = .false.
      else if (name == 'constraints') then ! end of constraints stanza
         parse_in_constraints = .false.
      endif

   end subroutine

   !% Initialise restraints and constraints for a dynamical system from an XML containing string.
   !%
   !% XML format for constraints
   !%
   !% \begin{verbatim}
   !% <constraints N="n">
   !%   <bond_length atom_1="i" atom_2="j" d="r" />
   !%   <bond_length_sq atom_1="i" atom_2="j" d="r" />
   !%   <bond_length_diff atom_1="i" atom_2="j" atom_3="k" d="r" />
   !%   <bond_angle_cos atom_1="i" atom_2="j" atom_3="k" c="r" />
   !% </constraints N="n">
   !% \end{verbatim}
   !%
   !% In all types, constraint value (c or d) is optional.  If not
   !% specified, defaults to value at initialization time, just like
   !% underlying constraint adding routines.
   !%
   !% Format for restraints is similar, except surrounding stanza name
   !% is `restraints', and the additional {\tt k} attribute
   !% specifying spring constant for each restraint
   !
   subroutine init_restraints_constraints(this, param_str)
      type(DynamicalSystem), intent(inout), target :: this
      character(len=*), intent(in) :: param_str

      type(xml_t) :: fxml

      if (len_trim(param_str) == 0) return

      parse_ds => this

      call open_xml_string(fxml, param_str)

      call parse(fxml, startElement_handler = restraints_constraints_startElement_handler, &
         endElement_handler = restraints_constraints_endElement_handler)

      call close_xml_t(fxml)

   end subroutine init_restraints_constraints

   subroutine interpret_bound_string(bound_str,bound)
     character(len=*), intent(in) :: bound_str
     integer, intent(out) :: bound

     if ( trim(bound_str) == trim(BOUND_STRING(UPPER_BOUND)) ) then
        bound = UPPER_BOUND
     elseif ( trim(bound_str) == trim(BOUND_STRING(LOWER_BOUND)) ) then
        bound = LOWER_BOUND
     elseif ( trim(bound_str) == trim(BOUND_STRING(BOTH_UPPER_AND_LOWER_BOUNDS)) ) then
        bound = BOTH_UPPER_AND_LOWER_BOUNDS
     else
        call system_abort("bound must take a value of "//trim(BOUND_STRING(UPPER_BOUND))//", "// &
           trim(BOUND_STRING(LOWER_BOUND))//" and "//trim(BOUND_STRING(BOTH_UPPER_AND_LOWER_BOUNDS)))
     endif

   end subroutine interpret_bound_string

end module restraints_constraints_xml_module
