open Util
open Constants

exception Null_contribution

(** Array of shell pairs obtained from combining two contracted shells. The
    two integers correspond to the indices of the combined shells.
*)
type t =
{
  shell_a     : Contracted_shell.t;
  shell_b     : Contracted_shell.t;
  shell_pairs : ShellPair.t array;
  coef        : float array; 
  expo_inv    : float array; 
  center_ab   : Coordinate.t;       (* A-B *)
  norm_sq     : float;              (* |A-B|^2 *)
  norm_coef_scale : float array;    (* norm_coef.(i) / norm_coef.(0) *)
}


(** 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
  let shell_pairs = 
    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 ; 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
  in
  let coef     = Array.map (fun x -> (fun y -> y.ShellPair.coef) x) shell_pairs
  and expo_inv = Array.map (fun x -> (fun y -> y.ShellPair.expo_inv) x) shell_pairs
  in
  {
    shell_a = p_a ; shell_b = p_b ; coef ; expo_inv ;
    shell_pairs ; center_ab=shell_pairs.(0).center_ab;
    norm_coef_scale ; norm_sq=shell_pairs.(0).norm_sq
  }


(** Returns an integer characteristic of a contracted shell pair *)
let hash a = 
  Array.map Hashtbl.hash a


(** Comparison function, used for sorting *)
let cmp a b =
  if a = b then 0
  else if (Array.length a < Array.length b) then -1
  else if (Array.length a > Array.length b) then  1
  else 
    let out = ref 0 in
    begin
     try
        for k=0 to (Array.length a - 1) do
          if a.(k) < b.(k) then
            (out := (-1); raise Not_found)
          else if a.(k) > b.(k) then
            (out :=  1; raise Not_found);
        done
     with Not_found -> ();
    end;
    !out


(** The array of all shell pairs with their correspondance in the list
    of contracted shells.
*)
let shell_pairs basis =
  Array.mapi (fun i shell_a ->
    Array.mapi (fun j shell_b -> 
      create shell_a shell_b)
    (Array.sub basis 0 (i+1)) 
  ) basis


let equivalent x y =
  (Array.length x = Array.length y) &&
  (Array.init (Array.length x) (fun k -> ShellPair.equivalent x.(k) y.(k))
  |> Array.fold_left (fun accu x -> x && accu) true)


(** A list of unique shell pairs *)
let unique sp =
  let sp = 
    Array.to_list sp
    |> Array.concat
    |> Array.to_list
  in
  let rec aux accu = function
  | [] -> accu
  | x::rest -> 
      try ignore @@ List.find (fun y -> equivalent x y) accu; aux accu rest
      with Not_found -> aux (x::accu) rest
  in
  aux [] sp


(** A map from a shell pair hash to the list of indices in the array
    of shell pairs.
*)
let indices 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