(** Electron-electron repulsion integrals *)

open Util

(** (00|00)^m : Fundamental electron repulsion integral
    $ \int \int \phi_p(r1) 1/r_{12} \phi_q(r2) dr_1 dr_2 $

    maxm        : Maximum total angular momentum
    expo_pq_inv : $1./p + 1./q$ where $p$ and $q$ are the exponents of
                $\phi_p$ and $\phi_q$
    norm_pq_sq  : square of the distance between the centers of $\phi_p$
                and $\phi_q$
*)
let zero_m ~maxm ~expo_pq_inv ~norm_pq_sq =
  let exp_pq = 1. /. expo_pq_inv in
  let t = norm_pq_sq *. exp_pq in
  boys_function ~maxm t
  |> Array.mapi (fun m fm ->
      two_over_sq_pi *. (if m mod 2 = 0 then fm else -.fm) *.
      (pow exp_pq m) *. (sqrt exp_pq)
    )


(** Compute all the integrals of a contracted class *)
(*
let contracted_class  shell_a shell_b shell_c shell_d : float Zmap.t =
  TwoElectronRRVectorized.contracted_class ~zero_m shell_a shell_b shell_c shell_d 
  *)
let contracted_class  shell_a shell_b shell_c shell_d : float Zmap.t =
  TwoElectronRR.contracted_class ~zero_m shell_a shell_b shell_c shell_d 

(** Compute all the integrals of a contracted class *)
let contracted_class_shell_pairs_vec ?schwartz_p ?schwartz_q shell_p shell_q : float Zmap.t =
  TwoElectronRRVectorized.contracted_class_shell_pairs ~zero_m ?schwartz_p ?schwartz_q shell_p shell_q

let contracted_class_shell_pairs ?schwartz_p ?schwartz_q shell_p shell_q : float Zmap.t =
  TwoElectronRR.contracted_class_shell_pairs ~zero_m ?schwartz_p ?schwartz_q shell_p shell_q


let cutoff2 = cutoff *. cutoff
(*
type n_cls = { n : int ; cls : Z.t array }
*)
exception NullIntegral

(** Write all integrals to a file with the <ij|kl> convention *)
let to_file ~filename basis =
  let to_int_tuple x = 
    let open Zkey in
    match to_int_tuple Kind_3 x with
    | Three x -> x
    | _ -> assert false
  in

  let oc = open_out filename in

  (* Pre-compute all shell pairs *)
  let shell_pairs =
    Array.mapi (fun i shell_a -> Array.map (fun shell_b -> 
        Shell_pair.create_array shell_a shell_b) (Array.sub basis 0 (i+1)) ) basis
  in
  Printf.printf "%d shells\n" (Array.length basis);

  (* Pre-compute diagonal integrals for Schwartz *)
  let schwartz = 
    Array.map (fun pair_array -> Array.map (fun pair ->
        let cls = 
          contracted_class_shell_pairs pair pair
        in
        (cls, Zmap.fold (fun key value accu -> max (abs_float value) accu) cls 0. )
      ) pair_array ) shell_pairs
  in

  let icount = ref 0 in
  for i=0 to (Array.length basis) - 1 do
    print_int basis.(i).Contracted_shell.indice ; print_newline ();
    for j=0 to i do
      let schwartz_p, schwartz_p_max = schwartz.(i).(j) in
        if (schwartz_p_max >= cutoff) then
          icount := !icount + 1;
    done;
  done;
  Printf.printf "%d shell pairs\n" !icount;

  let inn = ref 0 and out = ref 0 in

  for i=0 to (Array.length basis) - 1 do
    print_int basis.(i).Contracted_shell.indice ; print_newline ();
    for j=0 to i do
      let schwartz_p, schwartz_p_max = schwartz.(i).(j) in
      try
        if (schwartz_p_max < cutoff) then raise NullIntegral;
        let
          shell_p = shell_pairs.(i).(j)
        in

        for k=0 to i do
          for l=0 to k do
            let schwartz_q, schwartz_q_max = schwartz.(k).(l) in
            try
            if schwartz_p_max *. schwartz_q_max < cutoff2 then
              raise NullIntegral;
            let
              shell_q = shell_pairs.(k).(l)
            in

            let swap =
              Array.length shell_q < Array.length shell_p 
            in

            (* Compute all the integrals of the class *)
            let cls =
              if swap then
                if Array.length shell_p < 2 then
                  contracted_class_shell_pairs ~schwartz_p:schwartz_q ~schwartz_q:schwartz_p shell_q shell_p 
                else
                  contracted_class_shell_pairs_vec ~schwartz_p:schwartz_q ~schwartz_q:schwartz_p shell_q shell_p 
              else
                if Array.length shell_q < 2 then
                  contracted_class_shell_pairs ~schwartz_p ~schwartz_q shell_p shell_q 
                else
                  contracted_class_shell_pairs_vec ~schwartz_p ~schwartz_q shell_p shell_q 
            in

            (* Write the data in the output file *)
            Array.iteri (fun i_c powers_i ->
                let i_c = basis.(i).Contracted_shell.indice + i_c + 1 in
                let xi = to_int_tuple powers_i in
                Array.iteri (fun j_c powers_j ->
                    let j_c = basis.(j).Contracted_shell.indice + j_c + 1 in
                    let xj = to_int_tuple powers_j in
                    Array.iteri (fun k_c powers_k ->
                        let k_c = basis.(k).Contracted_shell.indice + k_c + 1 in 
                        let xk = to_int_tuple powers_k in
                        Array.iteri (fun l_c powers_l ->
                            let l_c = basis.(l).Contracted_shell.indice + l_c + 1 in
                            let xl = to_int_tuple powers_l in
                            let key = 
                              if swap then
                                Zkey.of_int_tuple (Zkey.Twelve (xk,xl,xi,xj))
                              else
                                Zkey.of_int_tuple (Zkey.Twelve (xi,xj,xk,xl))
                            in
                            let value = 
                              Zmap.find cls key 
                            in
                            if (abs_float value > cutoff) then
                              (inn := !inn + 1;
                              (*
                              assert ((i_c, k_c, j_c, l_c) <> (-1,0,0,0));
                                *)
                              Printf.fprintf oc "%4d %4d %4d %4d %20.12e\n"
                                i_c k_c j_c l_c value
                                )
                            else
                              out := !out + 1;
                          ) basis.(l).Contracted_shell.powers 
                      ) basis.(k).Contracted_shell.powers
                  ) basis.(j).Contracted_shell.powers
              ) basis.(i).Contracted_shell.powers;
              with NullIntegral ->  ()
          done;
        done;
      with NullIntegral -> ()
    done;
  done;
  Printf.printf "In: %d  Out:%d\n" !inn !out ;
  close_out oc

(*

module IntegralMap = Zmap
let index i j k l =
  Zkey.of_int_array Zkey.Kind_4 [| i;j;k;l |]


module Key = struct
 type t=int
 let equal (x:int) (y:int) = x = y 
 let hash (x:int) = x
end
module IntegralMap = Hashtbl.Make(Key)
let index i j k l =
  let f i k = 
    let (p,r) = 
      if i <= k then (i,k) else (k,i)
    in p+ (r*r-r)/2
  in
  let p = f i k and q = f j l in
  f p q


let to_file ~filename basis =
  let oc = open_out filename in
  let zkey = Array.map (fun b -> 
      let result =
        Angular_momentum.(zkey_array (Kind_1 b.Contracted_shell.totAngMom))
      in
      { n=Array.length result ;  cls=result }
    ) basis
  in 

  let key_array =  Array.make 12 0 in
  let result = ref [] in

  let i_shift = ref 1 in
  for i=0 to (Array.length basis) - 1 do
    print_int !i_shift ; print_newline ();
    let j_shift = ref 1 in
    for j=0 to i do
      let k_shift = ref 1 in
      for k=0 to i do
        let l_shift = ref 1 in
        for l=0 to k do
          let cls =
            contracted_class basis.(i) basis.(j) basis.(k) basis.(l)
          in

          for i_c = 0 to zkey.(i).n - 1 do
            let x = Zkey.(to_int_array Kind_3 zkey.(i).cls.(i_c)) in
            key_array.(0) <- x.(0);
            key_array.(1) <- x.(1);
            key_array.(2) <- x.(2);
            for j_c = 0 to zkey.(j).n - 1 do
              let x = Zkey.(to_int_array Kind_3 zkey.(j).cls.(j_c)) in
              key_array.(3) <- x.(0);
              key_array.(4) <- x.(1);
              key_array.(5) <- x.(2);
              for k_c = 0 to zkey.(k).n - 1 do
                let x = Zkey.(to_int_array Kind_3 zkey.(k).cls.(k_c)) in
                key_array.(6) <- x.(0);
                key_array.(7) <- x.(1);
                key_array.(8) <- x.(2);
                for l_c = 0 to zkey.(l).n - 1 do
                  let x = Zkey.(to_int_array Kind_3 zkey.(l).cls.(l_c)) in
                  key_array.( 9) <- x.(0);
                  key_array.(10) <- x.(1);
                  key_array.(11) <- x.(2);
                  let key = 
                    Zkey.(of_int_array Kind_12  key_array)
                  in
                  let value = 
                    Zmap.find cls key 
                  in
                  if (abs_float value > cutoff) then
                    let key =
                      index (!i_shift+i_c) (!j_shift+j_c) (!k_shift+k_c) (!l_shift+l_c) 
                    in
                    result := (key, value) :: !result
                done
              done
            done
          done;
          l_shift := !l_shift + zkey.(l).n
        done;
        k_shift := !k_shift + zkey.(k).n
      done;
      j_shift := !j_shift + zkey.(j).n
    done;
    i_shift := !i_shift + zkey.(i).n
  done
  ;
  print_endline "Computation Done";

  let result = Array.of_list !result in
  let result =
    let a = IntegralMap.create (Array.length result) in
    Array.iter (fun (k,v) -> IntegralMap.add a k v) result;
    a
  in
  print_endline "Map formed";

  for i=1 to !i_shift - 1 do
    for k=1 to !i_shift - 1 do
      for j=1 to !i_shift - 1 do
        for l=1 to !i_shift - 1 do
          let key =
            index i j k l
          in
          try
            let value = 
              IntegralMap.find result key 
            in
            Printf.fprintf oc " %5d %5d %5d %5d%20.15f\n" i j k l value
          with Not_found -> ()
        done
      done
    done
  done;
  close_out oc

let xto_file ~filename basis =
  let zkey = Array.map (fun b -> 
      let result =
        Angular_momentum.(zkey_array (Kind_1 b.Contracted_shell.totAngMom))
      in
      { n=Array.length result ;  cls=result }
    ) basis
  in 

  let key_array =  Array.make 12 0 in
  let result = ref [] in

  let (i,j,k,l) = (1,1,1,18) in
  let (i,j,k,l) = (i-1,j-1,k-1,l-1) in
  basis.(i) |> Contracted_shell.to_string |> print_endline;
  basis.(j) |> Contracted_shell.to_string |> print_endline;
  basis.(k) |> Contracted_shell.to_string |> print_endline;
  basis.(l) |> Contracted_shell.to_string |> print_endline;
  let bi, bj, bk, bl =
    basis.(i), basis.(j), basis.(k), basis.(l)
  in
  let cls =
    (*contracted_class basis.(i) basis.(j) basis.(k) basis.(l) *)
    contracted_class bi bj bk bl
  in
  Zmap.iter (fun  k v  -> Printf.printf "%50s %20e\n" Zkey.(to_string Kind_12 k) v) cls;

          for i_c = 0 to zkey.(i).n - 1 do
            let x = Zkey.(to_int_array Kind_3 zkey.(i).cls.(i_c)) in
            key_array.(0) <- x.(0);
            key_array.(1) <- x.(1);
            key_array.(2) <- x.(2);
            for j_c = 0 to zkey.(j).n - 1 do
              let x = Zkey.(to_int_array Kind_3 zkey.(j).cls.(j_c)) in
              key_array.(3) <- x.(0);
              key_array.(4) <- x.(1);
              key_array.(5) <- x.(2);
              for k_c = 0 to zkey.(k).n - 1 do
                let x = Zkey.(to_int_array Kind_3 zkey.(k).cls.(k_c)) in
                key_array.(6) <- x.(0);
                key_array.(7) <- x.(1);
                key_array.(8) <- x.(2);
                for l_c = 0 to zkey.(l).n - 1 do
                  let x = Zkey.(to_int_array Kind_3 zkey.(l).cls.(l_c)) in
                  key_array.( 9) <- x.(0);
                  key_array.(10) <- x.(1);
                  key_array.(11) <- x.(2);
                  let key = 
                    Zkey.(of_int_array Kind_12  key_array)
                  in
                  let value = 
                    Zmap.find cls key 
                  in
                  if (abs_float value > cutoff) then
                    result := (key, value) :: !result
                done
              done
            done
          done;

  List.iter (fun (k,v) -> Printf.printf "%60s %e\n" Zkey.(to_string Kind_12 k) v) !result
  ;
  *)