(*
  vim::syntax=ocaml
 *)

open Qputils
open Qptypes
open Sexplib.Std

(** Interactive editing of the input.

WARNING
This file is automatically generated by
`${{QP_ROOT}}/scripts/ezfio_interface/ei_handler.py`
*)


(** Keywords used to define input sections *)
type keyword =
| Ao_basis
| Ao_extra_basis
| Determinants_by_hand
| Electrons
| Mo_basis
| Nuclei_by_hand
{keywords}



let keyword_to_string = function
| Ao_basis              -> "AO basis"
| Ao_extra_basis              -> "AO extra_basis"
| Determinants_by_hand  -> "Determinants_by_hand"
| Electrons             -> "Electrons"
| Mo_basis              -> "MO basis"
| Nuclei_by_hand                -> "Molecule"
{keywords_to_string}




(** Create the header of the temporary file *)
let file_header filename =
  Printf.sprintf  "
==================================================================
                       Quantum Package
==================================================================

Editing file `%s`

" filename



(** Creates the header of a section *)
let make_header kw =
  let s = keyword_to_string kw in
  let l = String.length s in
  "\n\n"^s^"\n"^(String.init l (fun _ -> '='))^"\n\n"



(** Returns the rst string of section [s] *)
let get s =
  let header = (make_header s) in
  let f (read,to_rst) =
    match read () with
    | Some text -> header ^ (Rst_string.to_string (to_rst text))
    | None      -> ""
  in
  let rst =
    try
      begin
         let open Input in
         match s with
         | Mo_basis ->
           f Mo_basis.(read, to_rst)
         | Electrons ->
           f Electrons.(read, to_rst)
         | Nuclei_by_hand ->
           f Nuclei_by_hand.(read, to_rst)
         | Ao_extra_basis ->
           f Ao_extra_basis.(read, to_rst)
         | Ao_basis ->
           f Ao_basis.(read, to_rst)
         | Determinants_by_hand ->
           f Determinants_by_hand.(read ~full:false, to_rst)
{section_to_rst}
      end
    with
    | Sys_error msg -> (Printf.eprintf "Info: %s\n%!" msg ; "")
  in
  rst



(** Applies the changes from the string [str] corresponding to section [s] *)
let set str s =
  let header = (make_header s) in
  match String_ext.substr_index ~pos:0 ~pattern:header str with
  | None -> ()
  | Some idx ->
    begin
      let index_begin = idx + (String.length header) in
      let index_end   =
        match ( String_ext.substr_index ~pos:(index_begin+(String.length header)+1)
          ~pattern:"==" str) with
          | Some i -> i
          | None -> String.length str
      in
      let l = index_end - index_begin in
      let str = String.sub str index_begin l
      |> Rst_string.of_string
      in
      let write (of_rst,w) s =
        try
        match of_rst str with
        | Some data -> w data
        | None -> ()
        with
        | _ -> (Printf.eprintf "Info: Read error in %s\n%!"
               (keyword_to_string s); ignore (of_rst str) )
      in
      let open Input in
        match s with
{write}
        | Electrons        -> write Electrons.(of_rst, write) s
        | Determinants_by_hand     -> write Determinants_by_hand.(of_rst, write ~force:false) s
        | Nuclei_by_hand           -> write Nuclei_by_hand.(of_rst, write) s
        | Ao_basis         -> () (* TODO *)
        | Ao_extra_basis         -> () (* TODO *)
        | Mo_basis         -> () (* TODO *)
    end



(** Creates the temporary file for interactive editing *)
let create_temp_file ?filename ezfio_filename fields =
  let temp_filename  =
    match filename with
    | None -> Filename.temp_file "qp_edit_" ".rst"
    | Some f -> f
  in
  let () =
    match filename with
    | None -> at_exit (fun () -> Sys.remove temp_filename)
    | _ -> ()
  in
  begin
    let oc = open_out temp_filename in
    (file_header ezfio_filename) :: (List.map get fields)
    |> String.concat "\n"
    |> Printf.fprintf oc "%s";
    close_out oc;
    temp_filename
  end





let run check_only ?ndet ?state ?read ?write ezfio_filename =

  (* Set check_only if the arguments are not empty *)
  let open_editor =
    match ndet, state, read, write with
    | None, None, None, None -> not check_only
    | _ -> false
  in

  (* Open EZFIO *)
  if (not (Sys.file_exists ezfio_filename)) then
    failwith (ezfio_filename^" does not exists");

  Ezfio.set_file ezfio_filename;

  (* Clean qp_stop status *)
  [ "qpstop" ; "qpkill" ]
  |> List.iter (fun f ->
    let stopfile =
      Filename.concat (Qpackage.ezfio_work ezfio_filename) f
    in
    if Sys.file_exists stopfile then
      Sys.remove stopfile
  );

  (* Reorder basis set *)
  begin
    match Input.Ao_basis.read() with
    | Some aos ->
        let ordering = Input.Ao_basis.ordering aos in
        let test = Array.copy ordering in
        Array.sort compare test ;
        if test <> ordering then
        begin
          Printf.eprintf "Warning: Basis set is not properly ordered. Redordering.\n";
          let new_aos = Input.Ao_basis.reorder aos in
          Input.Ao_basis.write new_aos;
          match Input.Mo_basis.read() with
          | None -> ()
          | Some mos ->
              let new_mos = Input.Mo_basis.reorder mos ordering in
              Input.Mo_basis.write new_mos
        end
    | _ -> ()
  end;

  (* Reorder extra basis set *)
  begin
    match Input.Ao_extra_basis.read() with
    | Some aos ->
        let ordering = Input.Ao_extra_basis.ordering aos in
        let test = Array.copy ordering in
        Array.sort compare test ;
        if test <> ordering then
        begin
          Printf.eprintf "Warning: Basis set is not properly ordered. Redordering.\n";
          let new_aos = Input.Ao_extra_basis.reorder aos in
          Input.Ao_extra_basis.write new_aos;
        end
    | _ -> ()
  end;

  begin
    match ndet with
    | None -> ()
    | Some n -> Input.Determinants_by_hand.update_ndet (Det_number.of_int n)
  end;

  begin
    match state with
    | None -> ()
    | Some range ->
       begin
         Input.Determinants_by_hand.extract_states range
       end
  end;


  (*
  let output = (file_header ezfio_filename) :: (
    List.map get [
      Ao_basis ;
      Ao_extra_basis ;
      Mo_basis ;
    ])
   in
  String.concat output
  |> print_string
  *)

  let tasks = [
      Nuclei_by_hand ;
      Ao_basis;
      Ao_extra_basis;
      Electrons ;
{tasks}
      Mo_basis;
      Determinants_by_hand ;
  ]
  in

  (* Create the temp file *)
  let temp_filename =
    match read, write with
    | None, None -> create_temp_file ezfio_filename tasks
    | Some filename, None -> filename
    | None, filename -> create_temp_file ?filename ezfio_filename tasks
    | x, y -> failwith "read and write options are incompatible"
  in


  if open_editor then
      begin
        (* Open the temp file with external editor *)
        let editor =
          try Sys.getenv "EDITOR"
          with Not_found -> "vi"
        in
        Printf.sprintf "%s %s" editor temp_filename
        |> Sys.command |> ignore
      end;

  if write = None then
    (* Re-read the temp file *)
    let temp_string  =
      let ic = open_in temp_filename in
      let result =
        input_lines ic
        |> String.concat "\n"
      in
      close_in ic;
      result
    in
    List.iter (fun x -> set temp_string x) tasks





(** Remove the backup file *)
let remove_backup ezfio_filename =
  let backup_filename =
    Printf.sprintf "%s/work/backup.tar" ezfio_filename
  in
  try Sys.remove backup_filename
  with _ -> ()

(** Create a backup file in case of an exception *)
let create_backup ezfio_filename =
  remove_backup ezfio_filename;
  let backup_filename =
    Printf.sprintf "%s/work/backup.tar" ezfio_filename
  in
  try
    Printf.sprintf "
      mkdir -p %s/work ;
      tar -cf .backup.tar --exclude=\"work/*\" %s && (mv .backup.tar %s || rm .backup.tar)
    "
      ezfio_filename ezfio_filename backup_filename
    |> Sys.command
    |> ignore
  with _ -> ()


(** Restore the backup file when an exception occuprs *)
let restore_backup ezfio_filename =
  let filename =
    Printf.sprintf "%s/work/backup.tar" ezfio_filename
  in
  if Sys.file_exists filename then
    begin
      Printf.sprintf "tar -xf %s" filename
      |> Sys.command |> ignore;
      remove_backup ezfio_filename
    end



let () =
  let open Command_line in
  begin
    set_header_doc (Sys.argv.(0) ^ " - quantum_package command");
    set_description_doc
      "Opens a text editor to edit the parameters of an EZFIO directory.";

    [ {{ short='c'; long="check"; opt=Optional;
        doc="Checks the input data";
        arg=Without_arg; }};

      {{
        short='w'; long="write"; opt=Optional;
        doc="Writes the qp_edit file to a file\"";
        arg=With_arg "<string>"; }};

      {{
        short='r'; long="read"; opt=Optional;
        doc="Reads the file and applies it to the EZFIO\"";
        arg=With_arg "<string>"; }};

      {{ short='n'; long="ndet"; opt=Optional;
        doc="Truncates the wavefunction to the target number of determinants";
        arg=With_arg "<int>"; }};

      {{
        short='s'; long="state"; opt=Optional;
        doc="Extracts selected states, for example \"[1,3-5]\"";
        arg=With_arg "<range>"; }};

      anonymous "EZFIO_DIR" Mandatory "EZFIO directory";
    ]
    |> set_specs ;

  end;

  (*  Handle options *)
  let write =
    Command_line.get "write"
  in
  let read =
    Command_line.get "read"
  in
  let ndet =
    match Command_line.get "ndet" with
    | None -> None
    | Some s -> (try Some (int_of_string s)
            with _ -> failwith "[-n|--ndet] expects an integer")
  in
  let state =
    match Command_line.get "state" with
    | None -> None
    | Some s -> (try Some (Range.of_string s)
            with _ -> failwith "[-s|--state] expects a range")
  in

  let c = Command_line.get_bool "check" in

  let ezfio_filename =
    match Command_line.anon_args () with
    | [x] -> x
    | _ -> (Command_line.help () ; failwith "EZFIO_DIR is missing")
  in

  at_exit (fun () -> remove_backup ezfio_filename);
  (* Run the program *)
  try
    if (not c) then create_backup ezfio_filename;
    run c ?ndet ?state ?read ?write ezfio_filename
  with
  | Failure exc
  | Invalid_argument exc ->
    begin
      Printf.eprintf "=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-\n\n";
      Printf.eprintf "%s\n\n" exc;
      Printf.eprintf "=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-\n\n";
      restore_backup ezfio_filename;
      ignore @@ exit 1
    end
  | Assert_failure (file, line, ch) ->
    begin
      Printf.eprintf "=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-\n\n";
      Printf.eprintf "Assert error in file $QP_ROOT/ocaml/%s, line %d, character %d\n\n" file line ch;
      Printf.eprintf "=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-\n\n";
      restore_backup ezfio_filename;
      ignore @@ exit 2
    end;

    exit 0