open Util
open Xyz_ast

type t = (Element.t * Coordinate.t) array

let of_xyz_file ~filename =
  let lexbuf =
    let ic = open_in filename in
    Lexing.from_channel ic
  in
  let data = 
    Xyz_parser.input Nuclei_lexer.read_all lexbuf
  in

  let len = List.length data.nuclei in
  if len <> data.number_of_atoms then
    Printf.sprintf "Error: expected %d atoms but %d read" 
      data.number_of_atoms len
    |> failwith;

  List.map (fun nucleus ->
    nucleus.element, Coordinate.angstrom_to_bohr nucleus.coord
  ) data.nuclei
  |> Array.of_list



let of_zmt_file ~filename =
  let ic = open_in filename in
  let rec aux accu =
    try
      let line = input_line ic in
      aux (line::accu)
    with End_of_file ->
      close_in ic;
      List.rev accu
      |> String.concat "\n"
  in aux []
  |> Zmatrix.of_string
  |> Zmatrix.to_xyz
  |> Array.map (fun (e,x,y,z) -> (e, (Angstrom.make {Point.x ; y ; z} )))
    

let to_string atoms =
    "
                Nuclear Coordinates (Angstrom)
                ------------------------------

-----------------------------------------------------------------------
   Center   Atomic   Element          Coordinates (Angstroms)
            Number                 X             Y             Z
-----------------------------------------------------------------------
" ^
    (Array.mapi (fun i (e, coord) ->
         let coord =
           Coordinate.bohr_to_angstrom coord
         in
         Printf.sprintf " %5d    %5d    %5s   %12.6f  %12.6f  %12.6f"
            (i+1)  (Element.to_int e)  (Element.to_string e) 
            coord.Angstrom.x coord.Angstrom.y coord.Angstrom.z
       ) atoms
     |> Array.to_list
     |> String.concat "\n" ) ^
"
-----------------------------------------------------------------------

"


let of_filename ~filename =
    of_xyz_file filename
  

let repulsion ~nuclei =
  let get_charge e = 
    Element.to_charge e
    |> Charge.to_float
  in
  Array.fold_left ( fun accu (e1, coord1) -> 
    accu +. 
    Array.fold_left (fun accu (e2, coord2) ->
      let r = Coordinate.(norm (coord1 |- coord2)) in
      if r > 0. then
          accu +. 0.5 *. (get_charge e2) *. (get_charge e1) /. r
      else accu
      ) 0. nuclei
    ) 0.  nuclei