open Util
open Xyz_ast

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

let of_xyz_lexbuf lexbuf =
  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_xyz_string buffer =
  Zmatrix.of_string buffer
  |> Zmatrix.to_xyz
  |> Array.map (fun (e,x,y,z) ->
        (e, Coordinate.(angstrom_to_bohr @@ make_angstrom { x ; y ; z} ))
      )


let of_xyz_string input_string = 
  Lexing.from_string input_string
  |> of_xyz_lexbuf


let of_xyz_file filename =
  let ic = open_in filename in
  let lexbuf =
    Lexing.from_channel ic
  in
  let result = 
    of_xyz_lexbuf lexbuf
  in
  close_in ic;
  result


let of_zmt_string buffer =
  Zmatrix.of_string buffer
  |> Zmatrix.to_xyz
  |> Array.map (fun (e,x,y,z) ->
        (e, Coordinate.(angstrom_to_bohr @@ make_angstrom { x ; y ; z} ))
      )
    

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 []
  |> of_zmt_string
    

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

-----------------------------------------------------------------------
   Center   Atomic   Element          Coordinates (Angstroms)
            Number                 X             Y             Z
-----------------------------------------------------------------------
" ^
    (Array.mapi (fun i (e, coord) ->
         let open Coordinate in
         let coord =
           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.x coord.y coord.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


let charge nuclei = 
  Array.fold_left (fun accu (e, _) -> accu + Charge.to_int (Element.to_charge e) )
    0 nuclei 
  |> Charge.of_int


let to_xyz_string t =
  [ string_of_int (Array.length t) ; "" ] @
  ( Array.mapi (fun i (e, coord) ->
         let open Coordinate in
         let coord =
           bohr_to_angstrom coord
         in
         Printf.sprintf " %5s   %12.6f  %12.6f  %12.6f"
            (Element.to_string e) coord.x coord.y coord.z
       ) t 
    |> Array.to_list )
  |> String.concat "\n"


let to_t2_string t =
  [ "# nAt nEl nCore nRyd" ;
    Printf.sprintf "  %d   %d  %d     0" (Array.length t)
    (Array.fold_left (+) 0 (Array.map (fun (e,_) -> Element.to_int e) t) )
    (2 * Array.length t); 
    "# Znuc   x            y           z" ]
  @ (Array.mapi (fun i (e, coord) ->
         let open Coordinate in
         Printf.sprintf " %5f    %12.6f  %12.6f  %12.6f"
            (Element.to_int e |> float_of_int)  coord.x coord.y coord.z
       ) t 
    |> Array.to_list )
  |> String.concat "\n"


let small_core a = 
  Array.fold_left (fun accu (e,_) -> accu + (Element.small_core e)) 0 a