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 *)
}


module Am  = AngularMomentum
module Co  = Coordinate
module Cs  = ContractedShell
module Ps  = PrimitiveShell
module Psp = PrimitiveShellPair
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=1.e-32) s_a s_b =

(*
Format.printf "@[<2>shell_a :@ %a@]@;" Cs.pp s_a;
Format.printf "@[<2>shell_b :@ %a@]@;" Cs.pp s_b;
*)

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

let center_ab = Co.( Cs.center s_a |- Cs.center s_b ) in

(*
Format.printf "@[center_ab :@ %a@]@;" Coordinate.pp_angstrom center_ab;
Format.printf "@[a_minus_b :@ %a@]@." Coordinate.pp_angstrom (Psp.a_minus_b (
    match make 0 (Cs.prim s_a).(0) 0 (Cs.prim s_b).(0) 0.
  with Some x -> x | _ -> assert false));
  *)

let norm_sq = Co.dot center_ab center_ab in

let norm_coef_scale_a =
  Cs.norm_coef_scale s_a
and norm_coef_scale_b =
  Cs.norm_coef_scale s_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

assert (norm_coef_scale =  Psp.norm_coef_scale (
  match make 0 (Cs.prim s_a).(0) 0 (Cs.prim s_b).(0) 0.
  with Some x -> x | _ -> assert false));

  let shell_pairs = 
    Array.init (Cs.size s_a) (fun i ->
let p_a = (Cs.prim s_a).(i) in
let p_a_expo_center = Co.( (Cs.expo s_a).(i) |. Cs.center s_a ) in
let norm_coef_a = (Cs.norm_coef s_a).(i)  in
assert (norm_coef_a = Ps.norm_coef p_a);

        let make = make 0 p_a in

        Array.init (Cs.size s_b) (fun j ->
let p_b = (Cs.prim s_b).(j) in

            try
              let sp = make 0 p_b cutoff in

let sp_ = match sp with Some x -> x | None -> raise Null_contribution in

              let norm_coef_b = (Cs.norm_coef s_b).(j) in
assert (norm_coef_b = Ps.norm_coef p_b);

              let norm_coef   = norm_coef_a *. norm_coef_b in

              if norm_coef < cutoff  then
                raise Null_contribution;
              let p_b_expo_center = Co.( (Cs.expo s_b).(j) |. Cs.center s_b ) in
              let expo = (Cs.expo s_a).(i) +. (Cs.expo s_b).(j) in
              let expo_inv = 1. /. expo in
              let center = Co.(expo_inv |. (p_a_expo_center |+ p_b_expo_center ) )
              in
              let argexpo =
                (Cs.expo s_a).(i) *. (Cs.expo s_b).(j) *. norm_sq *. expo_inv
              in
              let g =
                (pi *. expo_inv)**(1.5) *. exp (-. argexpo)
              in
              let coef =
                norm_coef *.  g
              in
              if abs_float coef  < cutoff then
                raise Null_contribution;
              let center_a =
                Co.(center |- Cs.center s_a)
              in
              let monocentric =
                Cs.(center s_a = center s_b)
              in
              let totAngMomInt = 
                Am.(Cs.totAngMom s_a + Cs.totAngMom s_b)
                |> Am.to_int
              in
assert (expo= Psp.expo sp_ );
assert (expo_inv= Psp.expo_inv sp_ );
assert (center= Psp.center sp_ );
              Some ( (Cs.coef s_a).(i) *. (Cs.coef s_b).(j), {
                Sp.i ; j ;
                shell_a=s_a ; shell_b=s_b ;
                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 (c,y) -> c *. y.Sp.coef) shell_pairs
  and expo_inv = Array.map (fun (_,y) -> y.Sp.expo_inv) shell_pairs
  in
  let shell_pairs = Array.map snd shell_pairs in
  {
    shell_a = s_a ; shell_b = s_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;
  }


let shell_a         x = x.shell_a
let shell_b         x = x.shell_b
let shell_pairs     x = x.shell_pairs
let coef            x = x.coef
let expo_inv        x = x.expo_inv
let center_ab       x = x.center_ab
let norm_sq         x = x.norm_sq
let totAngMomInt    x = x.totAngMomInt
let norm_coef_scale x = x.norm_coef_scale

let monocentric x = x.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 of_basis 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