open Util
open Constants
open Powers
open Coordinate

exception NullPair

(** 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 (angMom_a, angMom_b) 
    zero_m_array (expo_b) (expo_inv_p) (center_ab, center_pa, center_pc)
    map
  =

  let maxm = angMom_a.tot + angMom_b.tot  in
  let maxsze = maxm+1 in
  let empty = Array.make maxsze 0. in

  let get_xyz angMom =
    match angMom with
    | { y=0 ; z=0 ; _ } -> X
    | {       z=0 ; _ } -> Y
    | _                 -> Z
  in


  (** Vertical recurrence relations *)
  let rec vrr angMom_a =
    let { x=ax ; y=ay ; z=az } = angMom_a in
    if ax < 0 || ay < 0 || az < 0 then raise Exit
    else
      match angMom_a.tot  with
      | 0 -> zero_m_array
      | _ ->
        let key = Zkey.of_powers (Zkey.Three angMom_a) in

        try Zmap.find map key with
        | Not_found -> 
          let result = 
            let xyz = get_xyz angMom_a in
            let am = Powers.decr xyz angMom_a in
            let amxyz = Powers.get xyz am in
            if amxyz < 0 then empty else
              let f1 = Coordinate.get xyz center_pa 
              and f2 = expo_inv_p *.  (Coordinate.get xyz center_pc)
              in
              if amxyz < 1 then
                let v1 =
                  vrr am
                in
                Array.init maxsze (fun m ->
                    if m = maxm then (f1 *. v1.(m) ) else
                    (f1 *. v1.(m) ) -. f2 *. v1.(m+1) )
              else
                let v3 =
                  let amm = Powers.decr xyz am in
                  vrr amm
                in
                let v1 =
                  vrr am
                in
                let f3 = (float_of_int amxyz) *. expo_inv_p *. 0.5 in
                Array.init maxsze (fun m -> f1 *. v1.(m) -.
                    (if m = maxm then 0. else
                        f2 *. v1.(m+1) )
                    +. f3 *. (v3.(m) -.  if m = maxm then 0. else
                                expo_inv_p *. v3.(m+1))
                  )
          in Zmap.add map key result;
          result


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

    let { x=bx ; y=by ; z=bz } = angMom_b in
    if bx < 0 || by < 0 || bz < 0 then raise Exit
    else
      match angMom_b.tot  with
      | 0 -> (vrr angMom_a).(0)
      | _ ->
        let xyz = get_xyz angMom_b in
        let bxyz = Powers.get xyz angMom_b in
        if (bxyz < 1) then 0. else
          let ap = Powers.incr xyz angMom_a in
          let bm = Powers.decr xyz angMom_b in
          let h1 =
            hrr ap bm
          in
          let f2 =
            Coordinate.get xyz center_ab
          in
          if abs_float f2 < cutoff then h1 else
            let h2 =
              hrr angMom_a bm
            in
            h1 +. f2 *. h2

  in
  hrr angMom_a angMom_b









(** Computes all the one-electron integrals of the contracted shell pair *)
let contracted_class_shell_pair ~zero_m shell_p geometry : float Zmap.t =

  let shell_a = shell_p.ContractedShellPair.shell_a
  and shell_b = shell_p.ContractedShellPair.shell_b
  in
  let 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 *)

  let norm_coef_scale_p = shell_p.ContractedShellPair.norm_coef_scale 
  in
  for ab=0 to (Array.length shell_p.ContractedShellPair.shell_pairs - 1)
  do
    let b = shell_p.ContractedShellPair.shell_pairs.(ab).ShellPair.j in
    try
      begin
        let coef_prod = shell_p.ContractedShellPair.coef.(ab) in

        (** Screening on the product of coefficients *)
        if (abs_float coef_prod) < 1.e-4*.cutoff then
          raise NullPair;


        let expo_pq_inv =
          shell_p.ContractedShellPair.expo_inv.(ab) 
        in

        let center_ab = 
          shell_p.ContractedShellPair.center_ab
        in
        let center_p = 
          shell_p.ContractedShellPair.shell_pairs.(ab).ShellPair.center
        in
        let center_pa = 
          Coordinate.(center_p |- Contracted_shell.center shell_a) 
        in

        for c=0 to Array.length geometry - 1 do
          let element, nucl_coord = geometry.(c) in
          let charge = Element.to_charge element |> Charge.to_float in
          let center_pc =
            Coordinate.(center_p |- nucl_coord ) 
          in
          let norm_pq_sq = 
            Coordinate.dot center_pc center_pc
          in

          let zero_m_array = 
            zero_m ~maxm ~expo_pq_inv ~norm_pq_sq
          in
          match Contracted_shell.(totAngMom shell_a, totAngMom shell_b) with
          | Angular_momentum.(S,S) -> 
            let integral =
              zero_m_array.(0)
            in
            contracted_class.(0) <- contracted_class.(0) -. coef_prod *. integral *. charge
          | _ ->
            let map = Zmap.create (2*maxm) in
            let norm_coef_scale = norm_coef_scale_p in
            (* Compute the integral class from the primitive shell quartet *)
            class_indices
            |> Array.iteri (fun i key ->
                let (angMomA,angMomB) =
                  match Zkey.to_powers ~kind:Zkey.Kind_6 key with
                  | Zkey.Six x -> x
                  | _          -> assert false
                in
                let norm = norm_coef_scale.(i) in
                let coef_prod = coef_prod *. norm in
                let integral = 
                  hvrr_one_e (angMomA, angMomB)
                     zero_m_array
                    (Contracted_shell.expo shell_b b)
                    (shell_p.ContractedShellPair.expo_inv.(ab))
                    (center_ab, center_pa, center_pc)
                    map
                in
                contracted_class.(i) <- contracted_class.(i) -. coef_prod *. integral *. charge
              )
      done
  end
  with NullPair -> ()
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