open Qptypes;;
open Qputils;;
open Sexplib.Std;;

module Determinants_by_hand : sig
  type t =
    { n_int                  : N_int_number.t;
      bit_kind               : Bit_kind.t;
      n_det                  : Det_number.t;
      n_det_qp_edit          : Det_number.t;
      n_states               : States_number.t;
      expected_s2            : Positive_float.t;
      psi_coef               : Det_coef.t array;
      psi_det                : Determinant.t array;
      state_average_weight   : Positive_float.t array;
      mo_label               : MO_label.t;
    } [@@deriving sexp]
  val read : ?full:bool -> unit -> t option
  val write : ?force:bool -> t -> unit
  val to_string : t -> string
  val to_rst : t -> Rst_string.t
  val of_rst : Rst_string.t -> t option
  val read_n_int : unit -> N_int_number.t
  val update_ndet : Det_number.t -> unit
  val extract_state : States_number.t -> unit
  val extract_states : Range.t -> unit
end = struct
  type t =
    { n_int                  : N_int_number.t;
      bit_kind               : Bit_kind.t;
      n_det                  : Det_number.t;
      n_det_qp_edit          : Det_number.t;
      n_states               : States_number.t;
      expected_s2            : Positive_float.t;
      psi_coef               : Det_coef.t array;
      psi_det                : Determinant.t array;
      state_average_weight   : Positive_float.t array;
      mo_label               : MO_label.t;
    } [@@deriving sexp]
  ;;

  let get_default = Qpackage.get_ezfio_default "determinants";;

  let read_mo_label () =
    if not (Ezfio.has_determinants_mo_label ()) then
      if Ezfio.has_mo_basis_mo_label () then (
        let label = Ezfio.get_mo_basis_mo_label () in
        Ezfio.set_determinants_mo_label label) ;
    Ezfio.get_determinants_mo_label ()
    |> MO_label.of_string
  ;;

  let read_n_int () =
    if not (Ezfio.has_determinants_n_int()) then
       Ezfio.get_mo_basis_mo_num ()
       |> Bitlist.n_int_of_mo_num
       |> N_int_number.to_int
       |> Ezfio.set_determinants_n_int
    ;
    Ezfio.get_determinants_n_int ()
    |> N_int_number.of_int
  ;;

  let write_n_int n =
    N_int_number.to_int n
    |> Ezfio.set_determinants_n_int
  ;;


  let read_bit_kind () =
    if not (Ezfio.has_determinants_bit_kind ()) then
      Lazy.force Qpackage.bit_kind
      |> Bit_kind.to_int
      |> Ezfio.set_determinants_bit_kind
    ;
    Ezfio.get_determinants_bit_kind ()
    |> Bit_kind.of_int
  ;;

  let write_bit_kind b =
    Bit_kind.to_int b
    |> Ezfio.set_determinants_bit_kind
  ;;

  let read_n_det () =
    if not (Ezfio.has_determinants_n_det ()) then
      Ezfio.set_determinants_n_det 1
    ;
    Ezfio.get_determinants_n_det ()
    |> Det_number.of_int
  ;;

  let read_n_det_qp_edit () =
    if not (Ezfio.has_determinants_n_det_qp_edit ()) then
      begin
        let n_det = read_n_det () |> Det_number.to_int in
        Ezfio.set_determinants_n_det_qp_edit n_det
      end;
    Ezfio.get_determinants_n_det_qp_edit ()
    |> Det_number.of_int
  ;;

  let write_n_det n =
    Det_number.to_int n
    |> Ezfio.set_determinants_n_det
  ;;

  let write_n_det_qp_edit n =
    let n_det = read_n_det () |> Det_number.to_int in
    min n_det (Det_number.to_int n)
    |> Ezfio.set_determinants_n_det_qp_edit
  ;;

  let read_n_states () =
    if not (Ezfio.has_determinants_n_states ()) then
      Ezfio.set_determinants_n_states 1
    ;
    Ezfio.get_determinants_n_states ()
    |> States_number.of_int
  ;;

  let write_n_states n =
    let n_states =
      States_number.to_int n
    in
    (*
    let old_nstates, read_wf =
      Ezfio.get_determinants_n_states (),
      Ezfio.get_determinants_read_wf  ()
    in
    if read_wf && old_nstates <> n_states then
      Printf.eprintf "Warning : n_states could not be changed because read_wf is true\n%!"
    else
    *)
      begin
        Ezfio.set_determinants_n_states n_states;
        let data =
          Array.make n_states 1.
          |> Array.to_list
        in
        Ezfio.ezfio_array_of_list ~rank:1 ~dim:[| n_states |] ~data
        |> Ezfio.set_determinants_state_average_weight
      end
  ;;

  let write_state_average_weight data =
      let n_states =
        read_n_states ()
        |> States_number.to_int
      in
      let data =
        Array.map Positive_float.to_float data
        |> Array.to_list
      in
      Ezfio.ezfio_array_of_list ~rank:1 ~dim:[| n_states |] ~data
      |> Ezfio.set_determinants_state_average_weight
  ;;

  let read_state_average_weight () =
    let n_states =
      read_n_states ()
      |> States_number.to_int
    in
    if not (Ezfio.has_determinants_state_average_weight ()) then
     begin
        let data =
          Array.init n_states (fun _ -> 1./.(float_of_int n_states))
          |> Array.map Positive_float.of_float
        in
        write_state_average_weight data
     end;
    let result =
      Ezfio.get_determinants_state_average_weight ()
      |> Ezfio.flattened_ezfio
      |> Array.map Positive_float.of_float
    in
    if Array.length result = n_states then
      result
    else
      let data =
        Array.init n_states (fun _ -> 1./.(float_of_int n_states))
        |> Array.map Positive_float.of_float
      in
      (write_state_average_weight data; data)
  ;;

  let read_expected_s2 () =
    if not (Ezfio.has_determinants_expected_s2 ()) then
      begin
        let na = Ezfio.get_electrons_elec_alpha_num ()
        and nb = Ezfio.get_electrons_elec_beta_num  ()
        in
        let s = 0.5 *. (Float.of_int (na - nb))
        in
        Ezfio.set_determinants_expected_s2 ( s *. (s +. 1.) )
      end
    ;
    Ezfio.get_determinants_expected_s2 ()
    |> Positive_float.of_float
  ;;

  let write_expected_s2 s2 =
    Positive_float.to_float s2
    |> Ezfio.set_determinants_expected_s2
  ;;

  let read_psi_coef ~read_only () =
    if not (Ezfio.has_determinants_psi_coef ()) then
      begin
        let n_states =
          read_n_states ()
          |> States_number.to_int
        in
        Ezfio.ezfio_array_of_list ~rank:2 ~dim:[| 1 ; n_states |]
          ~data:(List.init n_states (fun i -> if (i=0) then 1. else 0. ))
          |> Ezfio.set_determinants_psi_coef
      end;
    begin
      if read_only then
        Ezfio.get_determinants_psi_coef_qp_edit ()
      else
        Ezfio.get_determinants_psi_coef ()
    end
    |> Ezfio.flattened_ezfio
    |> Array.map Det_coef.of_float
  ;;

  let write_psi_coef ~n_det ~n_states c =
    let n_det = Det_number.to_int n_det
    and c =
      Array.map Det_coef.to_float c
      |> Array.to_list
    and n_states =
      States_number.to_int n_states
    in
    let r =
      Ezfio.ezfio_array_of_list ~rank:2 ~dim:[| n_det ; n_states |] ~data:c
    in
    Ezfio.set_determinants_psi_coef r;
    Ezfio.set_determinants_psi_coef_qp_edit r
  ;;


  let read_psi_det ~read_only () =
    let n_int = read_n_int ()
    and n_alpha = Ezfio.get_electrons_elec_alpha_num ()
        |> Elec_alpha_number.of_int
    and n_beta = Ezfio.get_electrons_elec_beta_num ()
        |> Elec_beta_number.of_int
    in
    if not (Ezfio.has_determinants_psi_det ()) then
      begin
        let mo_num = MO_number.get_max () in
        let rec build_data accu =  function
          | 0 -> accu
          | n -> build_data ((MO_number.of_int ~max:mo_num n)::accu) (n-1)
        in
        let det_a = build_data [] (Elec_alpha_number.to_int n_alpha)
          |> Bitlist.of_mo_number_list n_int
        and det_b = build_data [] (Elec_beta_number.to_int n_beta)
          |> Bitlist.of_mo_number_list n_int
        in
        let data = ( (Bitlist.to_int64_list det_a) @
          (Bitlist.to_int64_list det_b) )
        in
        Ezfio.ezfio_array_of_list ~rank:3 ~dim:[| N_int_number.to_int n_int ; 2 ; 1 |] ~data:data
          |> Ezfio.set_determinants_psi_det ;
      end  ;
    let n_int = N_int_number.to_int n_int in
    let psi_det_array =
      if read_only then
        Ezfio.get_determinants_psi_det_qp_edit ()
      else
        Ezfio.get_determinants_psi_det ()
    in
    let dim = psi_det_array.Ezfio.dim
    and data =  Ezfio.flattened_ezfio psi_det_array
    in
    assert (n_int = dim.(0));
    assert (dim.(1) = 2);
    if read_only then
      assert (dim.(2) = (Det_number.to_int (read_n_det_qp_edit ())))
    else
      assert (dim.(2) = (Det_number.to_int (read_n_det ())));
    Array.init dim.(2) (fun i ->
      Array.sub data (2*n_int*i) (2*n_int) )
    |> Array.map (Determinant.of_int64_array
      ~n_int:(N_int_number.of_int n_int)
      ~alpha:n_alpha ~beta:n_beta )
  ;;

  let write_psi_det ~n_int ~n_det d =
    let data = Array.to_list d
      |> Array.concat
      |> Array.to_list
    in
    let r =
      Ezfio.ezfio_array_of_list ~rank:3 ~dim:[| N_int_number.to_int n_int ; 2 ; Det_number.to_int n_det |] ~data:data
    in
    Ezfio.set_determinants_psi_det r;
    Ezfio.set_determinants_psi_det_qp_edit r
  ;;

  let write_mo_label a =
    MO_label.to_string a
    |> Ezfio.set_determinants_mo_label


  let read ?(full=true) () =

    let n_det_qp_edit = read_n_det_qp_edit () in
    let n_det         = read_n_det         () in
    let read_only =
      if full then false else n_det_qp_edit <> n_det
    in

    if (Ezfio.has_mo_basis_mo_num ()) then
      try
        Some
        { n_int                  = read_n_int ()                ;
          bit_kind               = read_bit_kind ()             ;
          n_det                  = read_n_det ()                ;
          n_det_qp_edit          = read_n_det_qp_edit ()        ;
          expected_s2            = read_expected_s2 ()          ;
          psi_coef               = read_psi_coef ~read_only ()  ;
          psi_det                = read_psi_det ~read_only  ()  ;
          n_states               = read_n_states ()             ;
          state_average_weight   = read_state_average_weight () ;
          mo_label               = read_mo_label () ;
        }
      with _ -> None
    else
      (* No molecular orbitals, so no determinants *)
      None
  ;;

  let write ?(force=false)
    { n_int                ;
      bit_kind             ;
      n_det                ;
      n_det_qp_edit        ;
      expected_s2          ;
      psi_coef             ;
      psi_det              ;
      n_states             ;
      state_average_weight ;
      mo_label             ;
    } =
     write_n_int n_int ;
     write_bit_kind bit_kind;
     write_n_det n_det;
     write_n_states n_states;
     write_expected_s2 expected_s2;
     if force || (n_det <= n_det_qp_edit) then
        begin
          write_n_det_qp_edit n_det;
          write_psi_coef ~n_det:n_det ~n_states:n_states psi_coef ;
          write_psi_det ~n_int:n_int ~n_det:n_det psi_det
        end;
     write_state_average_weight state_average_weight ;
     write_mo_label mo_label ;
     ()
  ;;


  let to_rst b =
    let max =
      Ezfio.get_mo_basis_mo_num ()
    in
    let mo_num =
      MO_number.of_int ~max max
    in
    let det_text =
      let nstates =
        read_n_states ()
        |> States_number.to_int
      and ndet_qp_edit =
        Det_number.to_int b.n_det_qp_edit
      in
      let coefs_string i =
        Array.init nstates (fun j ->
          let ishift =
            j*ndet_qp_edit
          in
          if (ishift < Array.length b.psi_coef) then
            b.psi_coef.(i+ishift)
            |> Det_coef.to_float
            |> Float.to_string
          else
            "0."
        )
        |> Array.to_list |> String.concat "\t"
      in
      Array.init ndet_qp_edit (fun i ->
        Printf.sprintf "  %s\n%s\n"
          (coefs_string i)
          (Determinant.to_string ~mo_num:mo_num b.psi_det.(i)
           |> String_ext.split ~on:'\n'
           |> list_map (fun x -> "  "^x)
           |> String.concat "\n"
          )
      )
      |> Array.to_list
      |> String.concat "\n"
    in
    Printf.sprintf "
Force the selected wave function to be an eigenfunction of S^2.
If true, input the expected value of S^2 ::

  expected_s2 = %s

Number of determinants ::

  n_det = %s

State average weights ::

  state_average_weight = (%s)

Determinants ::

%s
"
     (b.expected_s2   |> Positive_float.to_string)
     (b.n_det         |> Det_number.to_string)
     (b.state_average_weight |> Array.map Positive_float.to_string |> Array.to_list |> String.concat "\t")
     det_text
     |> Rst_string.of_string
  ;;

  let to_string b =
    let mo_num = Ezfio.get_mo_basis_mo_num () in
    let mo_num = MO_number.of_int mo_num ~max:mo_num in
    Printf.sprintf "
n_int                  = %s
bit_kind               = %s
n_det                  = %s
n_states               = %s
expected_s2            = %s
state_average_weight   = %s
psi_coef               = %s
psi_det                = %s
"
     (b.n_int         |> N_int_number.to_string)
     (b.bit_kind      |> Bit_kind.to_string)
     (b.n_det         |> Det_number.to_string)
     (b.n_states      |> States_number.to_string)
     (b.expected_s2   |> Positive_float.to_string)
     (b.state_average_weight |> Array.to_list |> list_map Positive_float.to_string |> String.concat ",")
     (b.psi_coef  |> Array.map Det_coef.to_string |> Array.to_list
      |> String.concat ", ")
     (b.psi_det   |> Array.map (Determinant.to_string ~mo_num) |> Array.to_list
      |> String.concat "\n\n")
  ;;

  let of_rst r =
    let r = Rst_string.to_string r
    in

    (* Split into header and determinants data *)
    let idx =
      match String_ext.substr_index r ~pos:0 ~pattern:"\nDeterminants" with
      | Some x -> x
      | None -> assert false
    in
    let (header, dets) =
       (String.sub r 0 idx, String.sub r idx (String.length r - idx) )
    in

    (* Handle header *)
    let header = r
    |> String_ext.split ~on:'\n'
    |> List.filter (fun line ->
        if (line = "") then
          false
        else
          ( (String.contains line '=') && (line.[0] = ' ') )
       )
    |> list_map (fun line ->
        "("^(
        String_ext.tr line ~target:'=' ~replacement:' '
        |> String.trim
        )^")" )
    |> String.concat ""
    in

    (* Handle determinant coefs *)
    let dets = match ( dets
      |> String_ext.split ~on:'\n'
      |> list_map String.trim
    ) with
    | _::lines -> lines
    | _ -> failwith "Error in determinants"
    in

    let psi_coef =
      let rec read_coefs accu = function
      | [] -> List.rev accu
      | ""::""::tail -> read_coefs accu tail
      | ""::c::tail ->
          let c =
            String_ext.split ~on:'\t' c
            |> list_map (fun x -> Det_coef.of_float (Float.of_string x))
            |> Array.of_list
          in
          read_coefs (c::accu) tail
      | _::tail -> read_coefs accu tail
      in
      let a =
        let buffer =
          read_coefs [] dets
        in
        let nstates =
          List.hd buffer
          |> Array.length
        in
        let extract_state i =
          let i =
            i-1
          in
          list_map (fun x -> Det_coef.to_string x.(i)) buffer
          |> String.concat " "
        in
        let rec build_result = function
        | 1 -> extract_state 1
        | i -> (build_result (i-1))^" "^(extract_state i)
        in
        build_result nstates
      in
      "(psi_coef ("^a^"))"
    in

    (* Handle determinants *)
    let psi_det =
      let n_int = N_int_number.of_int @@ (MO_number.get_max () - 1) / 64 + 1 in
      let n_alpha = Ezfio.get_electrons_elec_alpha_num ()
        |> Elec_alpha_number.of_int
      and n_beta = Ezfio.get_electrons_elec_beta_num ()
        |> Elec_beta_number.of_int
      in
      let rec read_dets accu = function
      | [] -> List.rev accu
      | ""::_::alpha::beta::tail ->
          begin
            let newdet =
               (Bitlist.of_string ~zero:'-' ~one:'+' alpha ,
                Bitlist.of_string ~zero:'-' ~one:'+' beta)
               |> Determinant.of_bitlist_couple ~n_int ~alpha:n_alpha ~beta:n_beta
               |> Determinant.sexp_of_t
               |> Sexplib.Sexp.to_string
            in
            read_dets (newdet::accu) tail
          end
      | _::tail -> read_dets accu tail
      in
      let a =
        read_dets [] dets
        |> String.concat ""
      in
      "(psi_det ("^a^"))"
    in


    let bitkind =
      Printf.sprintf "(bit_kind %d)" (Lazy.force Qpackage.bit_kind
      |> Bit_kind.to_int)
    and mo_label =
      Printf.sprintf "(mo_label %s)" (MO_label.to_string @@ read_mo_label ())
    and n_int =
      Printf.sprintf "(n_int %d)" (N_int_number.get_max ())
    and n_states =
      Printf.sprintf "(n_states %d)" (States_number.to_int @@ read_n_states ())
    and n_det_qp_edit =
      Printf.sprintf "(n_det_qp_edit %d)" (Det_number.to_int @@ read_n_det_qp_edit ())
    in
    let s =
       String.concat "" [ header ; mo_label ; bitkind ; n_int ; n_states ; psi_coef ; psi_det ; n_det_qp_edit ]
    in




    Generic_input_of_rst.evaluate_sexp t_of_sexp s
  ;;

  let update_ndet n_det_new =
    Printf.printf "Reducing n_det to %d\n" (Det_number.to_int n_det_new);
    let n_det_new =
      Det_number.to_int n_det_new
    in
    let det =
      match read () with
      | Some x -> x
      | None   -> failwith "No determinants in file"
    in
    let n_det_old, n_states =
      Det_number.to_int det.n_det,
      States_number.to_int det.n_states
    in
    if n_det_new = n_det_old then
      ()
    ;
    if n_det_new > n_det_new then
      failwith @@ Printf.sprintf "Requested n_det should be less than %d" n_det_old
    ;
    for j=0 to (n_states-1) do
      let ishift_old, ishift_new =
        j*n_det_old,
        j*n_det_new
      in
      for i=0 to (n_det_new-1) do
        det.psi_coef.(i+ishift_new) <- det.psi_coef.(i+ishift_old)
      done
    done
    ;
    let new_det =
      { det with n_det = (Det_number.of_int n_det_new) }
    in
    write ~force:true new_det
  ;;

  let extract_state istate =
    Printf.printf "Extracting state %d\n" (States_number.to_int istate);
    let det =
      match read () with
      | Some x -> x
      | None   -> failwith "No determinants in file"
    in
    let n_det, n_states =
      Det_number.to_int det.n_det,
      States_number.to_int det.n_states
    in
    if (States_number.to_int istate) > n_states then
      failwith "State to extract should not be greater than n_states"
    ;
    let j =
      (States_number.to_int istate) - 1
    in
    begin
      if (j>0) then
        let ishift =
          j*n_det
        in
        for i=0 to (n_det-1) do
          det.psi_coef.(i) <- det.psi_coef.(i+ishift)
        done
    end;
    let new_det =
      { det with n_states = (States_number.of_int 1) }
    in
    write ~force:true new_det
  ;;

  let extract_states range =
    Printf.printf "Extracting states %s\n" (Range.to_string range);
    let det =
      match read () with
      | Some x -> x
      | None   -> failwith "No determinants in file"
    in
    let n_det, n_states =
      Det_number.to_int det.n_det,
      States_number.to_int det.n_states
    in
    Range.to_int_list range
    |> List.iter (fun istate ->
      if istate > n_states then
        failwith "State to extract should not be greater than n_states")
    ;
    let sorted_list =
      Range.to_int_list range
      |> List.sort compare
    in
    let state_shift = ref 0 in
    List.iter (fun istate ->
      let j =
        istate - 1
      in
      begin
        if (j>0) then
          let ishift =
            j*n_det
          in
          for i=0 to (n_det-1) do
            det.psi_coef.(!state_shift+i) <-
            det.psi_coef.(i+ishift)
          done
          ; Printf.printf "OK\n%!" ;
      end;
      state_shift := !state_shift + n_det
    ) sorted_list
    ;
    let new_det =
      { det with n_states = (States_number.of_int @@ List.length sorted_list) }
    in
    write ~force:true new_det
  ;;

end