open Util
open Constants

type t = ShellPair.t array

exception Null_contribution

(** Creates an contracted shell pair : an array of pairs of primitive shells.
    A contracted shell with N functions combined with a contracted
    shell with M functions generates a NxM array of shell pairs.
*)
let create ?cutoff p_a p_b =
  let cutoff, log_cutoff =
    match cutoff with
    | None -> -1., max_float
    | Some cutoff -> cutoff, -. (log cutoff)
  in

  let center_ab = Coordinate.(
    Contracted_shell.center p_a |- Contracted_shell.center p_b )
  in
  let norm_sq =
    Coordinate.dot center_ab center_ab
  in
  let norm_coef_scale_a =
    Contracted_shell.norm_coef_scale p_a 
  and norm_coef_scale_b =
    Contracted_shell.norm_coef_scale p_b 
  in
  let norm_coef_scale =
    Array.map (fun v1 ->
      Array.map (fun v2 -> v1 *. v2) norm_coef_scale_b
    ) norm_coef_scale_a
    |> Array.to_list
    |> Array.concat
  in
  Array.init (Contracted_shell.size p_a) (fun i ->
      let p_a_expo_center = Coordinate.(
        Contracted_shell.expo p_a i |. Contracted_shell.center p_a )
      in
      let norm_coef_a =
        Contracted_shell.norm_coef p_a i
      in

      Array.init (Contracted_shell.size p_b) (fun j ->
          try
            let norm_coef_b =
              Contracted_shell.norm_coef p_b j
            in
            let norm_coef = 
              norm_coef_a *. norm_coef_b
            in
            if (norm_coef < cutoff) then
              raise Null_contribution;
            let p_b_expo_center = Coordinate.(
              Contracted_shell.expo p_b j |. Contracted_shell.center p_b )
            in
            let expo = Contracted_shell.(expo p_a i +. expo p_b j) in
            let expo_inv = 1. /. expo in
            let center = Coordinate.(
              expo_inv |. (p_a_expo_center |+ p_b_expo_center ) )
            in
            let argexpo =
              Contracted_shell.(expo p_a i *. expo p_b 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 coef =
              norm_coef *. Contracted_shell.(coef p_a i *. coef p_b j) *. g
            in
            if (abs_float coef  < cutoff) then
              raise Null_contribution;
            let center_a =
              Coordinate.(center |- Contracted_shell.center p_a)
            in
            let monocentric =
              Contracted_shell.center p_a = Contracted_shell.center p_b
            in
            Some ShellPair.{ i ; j ; shell_a=p_a ; shell_b=p_b ; norm_coef ; coef ; expo ; expo_inv ; center ; center_a ; center_ab ; norm_sq ; norm_coef_scale ; monocentric }
          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


(** Returns an integer characteristic of a contracted shell pair *)
let hash a = 
  1 (*TODO*)


(** Comparison function, used for sorting *)
let cmp a b =
  hash a - hash b


(** The array of all shell pairs *)
let shell_pairs basis =
    Array.mapi (fun i shell_a -> Array.map (fun shell_b -> 
        create shell_a shell_b) (Array.sub basis 0 (i+1)) ) basis


(** A list of unique shell pairs *)
let unique_of_shell_pairs sp =
  Array.to_list sp
  |> Array.concat
  |> Array.to_list
  |> List.sort_uniq cmp


(** A map from a shell pair hash to the list of indices in the array of shell pairs. *)
let indices_of_shell_pairs sp =
  let map =
    Hashtbl.create 129
  in
  Array.iteri (fun i s ->
    Array.iteri (fun j shell_p ->
      let key =
        hash shell_p
      in
      Hashtbl.add map key (i,j); ) s ) sp;
  map