open Util

(** In chop f g, evaluate g only if f is non zero, and return f *. (g ()) *)
let chop f g =
  if (abs_float f) < cutoff then 0.
  else f *. (g ())



(** Horizontal and Vertical Recurrence Relations (HVRR) *)
let hvrr_one_e
    m (angMom_a, angMom_b) (totAngMom_a, totAngMom_b)
    (maxm, zero_m_array) (expo_inv_p) (center_ab, center_pa, center_pc)
    map
  =

  let totAngMom_a = Angular_momentum.to_int totAngMom_a
  and totAngMom_b = Angular_momentum.to_int totAngMom_b
  in

  (** Vertical recurrence relations *)
  let rec vrr m angMom_a totAngMom_a =

    if    angMom_a.(0) < 0 || angMom_a.(1) < 0 || angMom_a.(2) < 0 then 0.
    else
      match totAngMom_a with
      | 0 -> zero_m_array.(m)
      | _ ->

        let key = [| angMom_a.(0)+1; angMom_a.(1)+1; angMom_a.(2)+1; |]
                  |> Zkey.(of_int_array ~kind:Kind_3)
        in

        let (found, result) =
          try (true, Zmap.find map.(m) key) with
          | Not_found -> (false,
                          let am  = [| angMom_a.(0) ; angMom_a.(1) ; angMom_a.(2) |]
                          and amm = [| angMom_a.(0) ; angMom_a.(1) ; angMom_a.(2) |]
                          and xyz =
                            match angMom_a with
                            | [|0;0;_|] -> 2
                            | [|0;_;_|] -> 1
                            | _         -> 0
                          in
                          am.(xyz)  <- am.(xyz) - 1;
                          amm.(xyz) <- amm.(xyz) - 2;
                          chop (Coordinate.coord center_pa xyz)
                            (fun () -> vrr m am (totAngMom_a-1))
                          +. chop (0.5 *. (float_of_int am.(xyz)) *. expo_inv_p) 
                            (fun () -> vrr m amm (totAngMom_a-2))
                          -. chop ((Coordinate.coord center_pc xyz) *. expo_inv_p)
                            (fun () -> vrr (m+1) am (totAngMom_a-1))
                          -. chop (0.5 *. (float_of_int am.(xyz)) *. expo_inv_p *. expo_inv_p)
                            (fun () -> vrr (m+1) amm (totAngMom_a-2))
        in
        if not found then
          Zmap.add map.(m) key result;
        result


  (** Horizontal recurrence relations *)
  and hrr angMom_a angMom_b totAngMom_a totAngMom_b =

    if    angMom_b.(0) < 0 || angMom_b.(1) < 0 || angMom_b.(2) < 0 then 0.
    else
      match totAngMom_b with
      | 0 -> vrr 0 angMom_a 
      | _ ->

        let key = [| angMom_a.(0)+1; angMom_a.(1)+1; angMom_a.(2)+1;
                     angMom_b.(0)+1; angMom_b.(1)+1; angMom_b.(2)+1 |]
                  |> Zkey.(of_int_array ~kind:Kind_6)
        in

        let (found, result) =
          try (true, Zmap.find map.(m) key) with
          | Not_found -> (false,
                          begin
                              let ap  = [| angMom_a.(0) ; angMom_a.(1) ; angMom_a.(2) |]
                              and bm  = [| angMom_b.(0) ; angMom_b.(1) ; angMom_b.(2) |]
                              and xyz =
                                match angMom_b with
                                | [|0;0;_|] -> 2
                                | [|0;_;_|] -> 1
                                | _        -> 0
                              in
                              ap.(xyz) <- ap.(xyz) + 1;
                              bm.(xyz) <- bm.(xyz) - 1;
                              hrr ap bm (totAngMom_a+1) (totAngMom_b-1)
                              +. chop (Coordinate.coord center_ab xyz) (fun () ->
                                  hrr angMom_a bm totAngMom_a (totAngMom_b-1) )
                          end)
        in
        if not found then
          Zmap.add map.(m) key result;
        result

  in
  hrr m angMom_a angMom_b angMom_c angMom_d totAngMom_a totAngMom_b
    totAngMom_c totAngMom_d










(** Computes all the one-electron integrals of the contracted shell pair *)
let contracted_class_nuc ~zero_m shell_a shell_b shell_c shell_d : float Zmap.t =

  let shell_p = Shell_pair.create_array shell_a shell_b
  and maxm = 
    let open Angular_momentum in
    (to_int @@ Contracted_shell.totAngMom shell_a) + (to_int @@ Contracted_shell.totAngMom shell_b) 
  in

  (* Pre-computation of integral class indices *)
  let class_indices = 
    Angular_momentum.zkey_array
      (Angular_momentum.Doublet
         Contracted_shell.(totAngMom shell_a, totAngMom shell_b))
  in

  let contracted_class = 
    Array.make (Array.length class_indices) 0.;
  in
  ()

  (* Compute all integrals in the shell for each pair of significant shell pairs *)

  for ab=0 to (Array.length shell_p - 1)
  do
      let coef_prod =
         shell_p.(ab).Shell_pair.coef 
      in
      (** Screening on thr product of coefficients *)
      if (abs_float coef_prod) > 1.e-4*.cutoff then
        begin

            let expo_pq_inv =
              shell_p.(ab).Shell_pair.expo_inv 
            in
            let center_ab = 
              Coordinate.shell_p.(ab).Shell_pair.center_ab 
            in
            let norm_pq_sq =
              Coordinate.dot center_pq center_pq
            in

            let zero_m_array = 
              zero_m ~maxm ~expo_pq_inv ~norm_pq_sq
            in

              match Contracted_shell.(totAngMom shell_a, totAngMom shell_b,
                                      totAngMom shell_c, totAngMom shell_d) with
              | Angular_momentum.(S,S,S,S) -> Array.iteri (fun i key ->
                  let coef_prod =
                    shell_p.(ab).Shell_pair.coef *. shell_q.(cd).Shell_pair.coef 
                  in
                  let integral =
                    zero_m_array.(0)
                  in
                  contracted_class.(i) <- contracted_class.(i) +. coef_prod *. integral
                ) class_indices
              | _ ->
                  let d = shell_q.(cd).Shell_pair.j in
                  let map = Array.init maxm (fun _ -> Zmap.create (Array.length class_indices)) in
                  (* Compute the integral class from the primitive shell quartet *)
                  Array.iteri (fun i key ->
                      let (angMomA,angMomB,angMomC,angMomD) =
                        let a = Zkey.to_int_array Zkey.Kind_12 key in
                        ( [| a.(0) ;  a.(1) ;  a.(2) |], 
                          [| a.(3) ;  a.(4) ;  a.(5) |], 
                          [| a.(6) ;  a.(7) ;  a.(8) |], 
                          [| a.(9) ; a.(10) ; a.(11) |] )
                      in
                      let norm = 
                        shell_p.(ab).Shell_pair.norm_fun angMomA angMomB *. shell_q.(cd).Shell_pair.norm_fun angMomC angMomD 
                      in
                      let integral = chop norm (fun () ->
                        ghvrr 0 (angMomA, angMomB, angMomC, angMomD)
                          (Contracted_shell.totAngMom shell_a, Contracted_shell.totAngMom shell_b,
                          Contracted_shell.totAngMom shell_c, Contracted_shell.totAngMom shell_d)
                          (maxm, zero_m_array)
                          (Contracted_shell.expo shell_b b, Contracted_shell.expo shell_d d)
                          (shell_p.(ab).Shell_pair.expo_inv, shell_q.(cd).Shell_pair.expo_inv)
                          (shell_p.(ab).Shell_pair.center_ab, shell_q.(cd).Shell_pair.center_ab, center_pq)
                          map )
                      in
                      contracted_class.(i) <- contracted_class.(i) +. coef_prod *. integral
                ) class_indices
        end
    done
  done;
  let result = 
    Zmap.create (Array.length contracted_class)
  in
  Array.iteri (fun i key -> Zmap.add result key contracted_class.(i)) class_indices;
  result

*)