open Util
open Constants

exception Null_contribution

type t =
{
  shell_a     : ContractedShell.t;
  shell_b     : ContractedShell.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) *)
  totAngMomInt : int;               (* Total angular Momentum *)
  monocentric : bool;
}


module Am = AngularMomentum
module Co = Coordinate
module Cs = ContractedShell
module Sp = ShellPair

(** 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 = Co.( p_a.Cs.center |- p_b.Cs.center )
  in
  let norm_sq =
    Co.dot center_ab center_ab
  in
  let norm_coef_scale_a =
    p_a.norm_coef_scale 
  and norm_coef_scale_b =
    p_b.norm_coef_scale 
  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 p_a.size (fun i ->
        let p_a_expo_center = Co.(p_a.Cs.expo.(i) |. p_a.Cs.center) in
        let norm_coef_a = p_a.norm_coef.(i)  in

        Array.init p_b.size (fun j ->
            try
              let norm_coef_b = p_b.norm_coef.(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 = Co.(p_b.expo.(j) |. p_b.center) in
              let expo = p_a.expo.(i) +. p_b.expo.(j) in
              let expo_inv = 1. /. expo in
              let center = Co.(expo_inv |. (p_a_expo_center |+ p_b_expo_center ) )
              in
              let argexpo =
                p_a.Cs.expo.(i) *. p_b.Cs.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 coef =
                norm_coef *. p_a.coef.(i) *. p_b.coef.(j) *. g
              in
              if abs_float coef  < cutoff then
                raise Null_contribution;
              let center_a =
                Co.(center |- p_a.center)
              in
              let monocentric =
                p_a.center  = p_b.center 
              in
              let totAngMomInt = 
                Am.to_int p_a.totAngMom +
                Am.to_int p_b.totAngMom 
              in
              Some {
                Sp.i ; j ;
                shell_a=p_a ; shell_b=p_b ;
                norm_coef ; coef ;
                expo ; expo_inv ;
                center ; center_a ; center_ab ;
                norm_sq ; monocentric ; totAngMomInt
              }
            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.Sp.coef) x) shell_pairs
  and expo_inv = Array.map (fun x -> (fun y -> y.Sp.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;
    totAngMomInt = shell_pairs.(0).Sp.totAngMomInt;
    monocentric =  shell_pairs.(0).Sp.monocentric;
  }


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