open Util

type t = {
  expo     : float;
  expo_inv : float;
  center_ab: float array;
  center   : float array;
  norm_sq  : float;
  norm     : float;
  coef     : float;
  norm_fun : int array -> int array -> float;
  i        : int;
  j        : int;
}

exception Null_contribution

let create_array ?(cutoff=0.) p_a p_b =
  let log_cutoff =
    if (cutoff = 0.) then infinity
    else -. (log cutoff)
  in

  let x_a = Coordinate.x p_a.Contracted_shell.center
  and y_a = Coordinate.y p_a.Contracted_shell.center
  and z_a = Coordinate.z p_a.Contracted_shell.center
  and x_b = Coordinate.x p_b.Contracted_shell.center
  and y_b = Coordinate.y p_b.Contracted_shell.center
  and z_b = Coordinate.z p_b.Contracted_shell.center
  in
  (*
  match p_a.Contracted_shell.center, p_b.Contracted_shell.center with
  | [|x_a; y_a; z_a|], [|x_b; y_b; z_b|] ->
  *)
    let center_ab =
      Coordinate.(p_a.Contracted_shell.center |- p_b.Contracted_shell.center)
    in
    let norm_sq =
      Coordinate.dot center_ab center_ab
    in
    Array.init p_a.Contracted_shell.size (fun i ->
        let p_a_expo_center =
          [| p_a.Contracted_shell.expo.(i) *. x_a  ; p_a.Contracted_shell.expo.(i) *. y_a  ; p_a.Contracted_shell.expo.(i) *. z_a  |]
        in

        Array.init p_b.Contracted_shell.size (fun j ->
            try
              let f1 = 
                p_a.Contracted_shell.norm_coef.(i)
              in
              let f2 =
                p_b.Contracted_shell.norm_coef.(j)
              in
              let norm_fun a b =
                f1 a *. f2 b
              in
              let norm =
                norm_fun
                  [| Angular_momentum.to_int p_a.Contracted_shell.totAngMom ; 0 ; 0 |]
                  [| Angular_momentum.to_int p_b.Contracted_shell.totAngMom ; 0 ; 0 |]
              in
              if (norm < cutoff) then
                raise Null_contribution;
              let p_b_expo_center =
                [| p_b.Contracted_shell.expo.(j) *. x_b  ; p_b.Contracted_shell.expo.(j) *. y_b  ; p_b.Contracted_shell.expo.(j) *. z_b  |]
              in
              let expo = p_a.Contracted_shell.expo.(i) +. p_b.Contracted_shell.expo.(j) in
              let expo_inv = 1. /. expo in
              let center =
                [| (p_a_expo_center.(0) +. p_b_expo_center.(0)) *. expo_inv;
                   (p_a_expo_center.(1) +. p_b_expo_center.(1)) *. expo_inv;
                   (p_a_expo_center.(2) +. p_b_expo_center.(2)) *. expo_inv |]
              in
              let argexpo =
                p_a.Contracted_shell.expo.(i) *. p_b.Contracted_shell.expo.(j)
                *. norm_sq *. expo_inv
              in
              if (argexpo > log_cutoff) then
                raise Null_contribution;
              let g =
                (pi *. expo_inv)**(1.5) *. exp(-. argexpo)
              in
              let norm_inv = 1./.norm in
              let norm_fun a b =
                norm_inv *. norm_fun a b
              in
              let coef =
                norm *. p_a.Contracted_shell.coef.(i) *. p_b.Contracted_shell.coef.(j) *. g
              in
              if (abs_float coef  < cutoff) then
                raise Null_contribution;
              Some { i ; j ; norm_fun ; norm ; coef ; expo ; expo_inv ; center ; center_ab=(Coordinate.to_float_array center_ab) ; norm_sq }
            with
            | Null_contribution -> None
          )
      )
    |> Array.to_list
    |> Array.concat
    |> Array.to_list
    |> List.filter (function Some _ -> true | None -> false)
    |> List.map (function Some x -> x | None -> assert false)
    |> Array.of_list