open Util
open Constants

exception Null_contribution

type t =
{
  shell_a         : ContractedShell.t;
  shell_b         : ContractedShell.t;
  shell_pairs     : PrimitiveShellPair.t array;
  coefficients    : float array; 
  exponents_inv   : float array; 
  center_ab       : Coordinate.t;       (* A-B *)
  norm_sq         : float;              (* |A-B|^2 *)
  norm_scales     : float array;        (* norm_coef.(i) / norm_coef.(0) *)
  totAngMom       : AngularMomentum.t;  (* Total angular Momentum *)
}


module Am  = AngularMomentum
module Co  = Coordinate
module Cs  = ContractedShell
module Ps  = PrimitiveShell
module Psp = PrimitiveShellPair

(** 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 make ?(cutoff=1.e-32) s_a s_b =

  let make = Psp.create_make_of (Cs.primitives s_a).(0) (Cs.primitives s_b).(0) in

  let shell_pairs = 
    Array.mapi (fun i p_a ->
        let c_a = (Cs.coefficients s_a).(i) in
        let make = make p_a in
        Array.mapi (fun j p_b ->
            let c_b = (Cs.coefficients s_b).(j) in
            let coef = c_a *. c_b in
            assert (coef <> 0.);
            let cutoff = cutoff /. abs_float coef in
            coef, make p_b cutoff) (Cs.primitives s_b)) (Cs.primitives s_a)
    |> Array.to_list
    |> Array.concat
    |> Array.to_list
    |> List.filter (function (_, Some _) -> true | _ -> false)
    |> List.map (function (c, Some x) -> (c,x) | _ -> assert false)
    |> Array.of_list
  in


  let coefficients = Array.map (fun (c,y) -> c *. Psp.normalization y) shell_pairs
  and exponents_inv = Array.map (fun (_,y) -> Psp.exponent_inv y) shell_pairs
  in
  let shell_pairs = Array.map snd shell_pairs in
  if Array.length shell_pairs = 0 then
    None
  else
    let root = shell_pairs.(0) in
    Some {
      shell_a = s_a ; shell_b = s_b ; coefficients ; exponents_inv ; shell_pairs ;
      center_ab = Psp.a_minus_b root;
      norm_scales = Psp.norm_scales root;
      norm_sq=Psp.a_minus_b_sq root;
      totAngMom = Psp.totAngMom root;
    }


let shell_a         x = x.shell_a
let shell_b         x = x.shell_b
let shell_pairs     x = x.shell_pairs
let coefficients    x = x.coefficients
let exponents_inv   x = x.exponents_inv
let center_ab       x = x.center_ab
let norm_sq         x = x.norm_sq
let totAngMom       x = x.totAngMom
let norm_scales     x = x.norm_scales

let monocentric x = Psp.monocentric x.shell_pairs.(0)

(** 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 of_contracted_shell_array basis =
  Array.mapi (fun i shell_a ->
    Array.mapi (fun j shell_b -> 
      make shell_a shell_b)
    (Array.sub basis 0 (i+1)) 
  ) basis


let equivalent x y =
  (Array.length x = Array.length y) &&
  let rec eqv = function
  | 0 -> true
  | k -> if Psp.equivalent x.(k) y.(k) then
            eqv (k-1)
         else false
  in eqv (Array.length x - 1)


(** 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