open Util
open Constants

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) (totAngMom_a, totAngMom_b)
    (maxm, zero_m_array) (expo_b) (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
  let maxm = totAngMom_a+totAngMom_b in
  let maxsze = maxm+1 in
  let empty = Array.make maxsze 0. in

  (** Vertical recurrence relations *)
  let rec vrr angMom_a totAngMom_a =
    let ax,ay,az = angMom_a in
    if (ax < 0) || (ay < 0) || (az < 0) then
      empty
    else
      match totAngMom_a with
      | 0 -> zero_m_array
      | _ ->
        let key = Zkey.of_int_tuple (Zkey.Three angMom_a) in

        try Zmap.find map key with
        | Not_found -> 
          let result = 
            let am, amm, amxyz, xyz  =
              match angMom_a with
              | (x,0,0) -> (x-1,0,0),(x-2,0,0), x-1, 0
              | (x,y,0) -> (x,y-1,0),(x,y-2,0), y-1, 1
              | (x,y,z) -> (x,y,z-1),(x,y,z-2), z-1, 2
            in
            if amxyz < 0 then empty else
              let f1 = Coordinate.coord center_pa xyz 
              and f2 = expo_inv_p *.  (Coordinate.coord center_pc xyz)
              in
              if amxyz < 1 then
                let v1 =
                  vrr am (totAngMom_a-1)
                in
                Array.init maxsze (fun m ->
                    if m = maxm then (f1 *. v1.(m) ) else
                    (f1 *. v1.(m) ) -. f2 *. v1.(m+1) )
              else
                let v3 =
                  vrr amm (totAngMom_a-2)
                in
                let v1 =
                  vrr am (totAngMom_a-1)
                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 totAngMom_a totAngMom_b =

    let bx,by,bz = angMom_b in
    if (bx < 0) || (by < 0) || (bz < 0) then 0.
    else
      match totAngMom_b with
      | 0 -> (vrr angMom_a totAngMom_a).(0)
      | _ ->
        let angMom_ax, angMom_ay, angMom_az = angMom_a
        and angMom_bx, angMom_by, angMom_bz = angMom_b in
        let bxyz, xyz =
          match angMom_b with
          | (_,0,0) -> angMom_bx, 0
          | (_,_,0) -> angMom_by, 1
          | (_,_,_) -> angMom_bz, 2
        in
        if (bxyz < 1) then 0. else
          let ap, bm =
            match xyz with
            | 0 -> (angMom_ax+1,angMom_ay,angMom_az),(angMom_bx-1,angMom_by,angMom_bz)
            | 1 -> (angMom_ax,angMom_ay+1,angMom_az),(angMom_bx,angMom_by-1,angMom_bz)
            | _ -> (angMom_ax,angMom_ay,angMom_az+1),(angMom_bx,angMom_by,angMom_bz-1)
          in
          let h1 =
            hrr ap bm (totAngMom_a+1) (totAngMom_b-1) 
          in
          let f2 =
            (Coordinate.coord center_ab xyz)
          in
          if (abs_float f2 < cutoff) then h1 else
            let h2 =
              hrr angMom_a bm totAngMom_a (totAngMom_b-1)
            in
            h1 +. f2 *. h2

  in
  hrr angMom_a angMom_b totAngMom_a totAngMom_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_int_tuple ~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)
                    (Contracted_shell.totAngMom shell_a, Contracted_shell.totAngMom shell_b)
                    (maxm, 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