(* ./run_Localisation.native -b basis -x xyz -c charge -o name.mos -i EZFIOdirectory > name.out *)
(* 
#.mos contains the localised orbitales 
#.out contains the localization convergence and the analysis of the spatial extent of the orbitales
*)

open Lacaml.D

let read_qp_mo dirname =
  let ic = Unix.open_process_in ("zcat "^dirname^"/mo_basis/mo_coef.gz") in
  let check = String.trim (input_line ic) in
  assert (check = "2");
  let int_list = 
    input_line ic
    |> String.split_on_char ' ' 
    |> List.filter (fun x -> x <> "")
    |> List.map int_of_string
  in
  let n_ao, n_mo = 
    match int_list with
    | [ x ; y ] -> x, y
    | _ -> assert false
  in
  let result = 
          Mat.init_cols n_ao n_mo (fun i j ->
            let s = input_line ic in
            Scanf.sscanf s " %f " (fun x -> x)
          )
  in
  let exit_code = Unix.close_process_in ic in
  assert (exit_code = Unix.WEXITED 0);
  result



let () =
  let open Command_line in
  begin
    set_header_doc (Sys.argv.(0) ^ " - QCaml command");
    set_description_doc "Localizes MOs";
    set_specs
      [ { short='b' ; long="basis" ; opt=Mandatory;
          arg=With_arg "<string>";
          doc="Name of the file containing the basis set"; } ;

        { short='x' ; long="xyz" ; opt=Mandatory;
          arg=With_arg "<string>";
          doc="Name of the file containing the nuclear coordinates in xyz format"; } ;

        { short='m' ; long="multiplicity" ; opt=Optional;
          arg=With_arg "<int>";
          doc="Spin multiplicity (2S+1). Default is singlet"; } ;

        { short='c' ; long="charge" ; opt=Optional;
          arg=With_arg "<int>";
          doc="Total charge of the molecule. Default is 0"; } ;
        
        { short='o' ; long="output" ; opt=Mandatory;
          arg=With_arg "<string>";
          doc="Name of the file containing the localized MOs"; } ;

        { short='i' ; long="import" ; opt=Optional;
          arg=With_arg "<string>";
          doc="Name of the EZFIO directory containing MOs"; } ;

        { short='s' ; long="separation" ; opt=Optional;
          arg=With_arg "<string>";
          doc="Type of separation: [ov: occ/vir | sp: sigma/pi | custom] "; } ;

      ]
  end;


  let basis_file  = Util.of_some @@ Command_line.get "basis" in

  let nuclei_file = Util.of_some @@ Command_line.get "xyz" in

  let ezfio_file = Command_line.get "import" in

  let charge = 
    match Command_line.get "charge" with
    | Some x -> int_of_string x
    | None   -> 0
  in

  let multiplicity =
    match Command_line.get "multiplicity" with
    | Some x -> int_of_string x
    | None -> 1
  in       

  let output_filename = Util.of_some @@ Command_line.get "output" in

  let separation =
    match Command_line.get "separation" with
      | None
      | Some "ov" -> Localisation.Occ_vir
      | Some "sp" -> Localisation.Sigma_pi
      | Some "custom" ->
        let list_of_lists =
          print_endline "Enter ranges, one per line:";
          let rec loop accu =
            try
              let r =
                input_line stdin
                |> Range.of_string
                |> Range.to_int_list
              in
              loop (r::accu)
            with End_of_file -> List.rev accu
          in loop []
        in
        Localisation.Custom list_of_lists
      | _ -> invalid_arg "separation"
  in
    

  (* II : Hartree-Fock *)

  (* 1. Def pour HF *)

  let simulation = 
    Simulation.of_filenames
      ~charge ~multiplicity ~nuclei:nuclei_file basis_file 
  in 

  (* 2. Calcul de Hartree-Fock*)
  let mo_basis = 
     match ezfio_file with
     | Some ezfio_file ->
           let mo_coef =  read_qp_mo  ezfio_file in
           let mo_type = MOBasis.Localized "Boys" in
           let elec = Simulation.electrons simulation in
           let n_mo = Mat.dim2 mo_coef in
           let mo_occupation = Vec.init n_mo (fun i ->
                   if i <= Electrons.n_beta elec then 2.
                   else if i <= Electrons.n_alfa elec then 1.
                   else 0.) in
           MOBasis.make ~simulation ~mo_type ~mo_occupation ~mo_coef ()
     | None -> HartreeFock.make simulation 
               |> MOBasis.of_hartree_fock
  in

  let local_mos = Localisation.localize mo_basis separation in

  let oc = open_out output_filename in
  Printf.fprintf oc "[\n";
  Mat.as_vec local_mos
  |> Vec.iter (fun x -> Printf.fprintf oc "%20.15e,\n" x);
  Printf.fprintf oc "]\n";
  close_out oc