open Util
open Constants


type t = {
  expo            : float;              (* alpha + beta *)
  expo_inv        : float;              (* 1/(alpha + beta) *)
  center          : Coordinate.t;       (* P = (alpha * A + beta B)/(alpha+beta) *)
  center_minus_a  : Coordinate.t;       (* P - A *)
  a_minus_b       : Coordinate.t;       (* A - B *)
  a_minus_b_sq    : float;              (* |A-B|^2 *)
  norm_coef_scale : float array lazy_t;
  norm_coef       : float;              (* norm_coef_a * norm_coef_b * g, with
                     g = (pi/(alpha+beta))^(3/2) exp (-|A-B|^2 * alpha*beta/(alpha+beta)) *)
  totAngMom       : AngularMomentum.t; 
  shell_a         : PrimitiveShell.t;
  shell_b         : PrimitiveShell.t;
  (*TODO*)
  i : int; j : int;
}

exception Null_contribution

module Am = AngularMomentum
module Co = Coordinate
module Ps = PrimitiveShell

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

let equivalent a b =
  a = b
(*
  Hashtbl.hash (a.expo, a.center_a, a.center_ab, a.coef, ContractedShell.totAngMom a.shell_a, ContractedShell.totAngMom a.shell_b) 
*)


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


let create_make_of p_a p_b =

  let a_minus_b = 
    Co.( Ps.center p_a |- Ps.center p_b )
  in

  let a_minus_b_sq =
    Co.dot a_minus_b a_minus_b
  in

  let norm_coef_scale = lazy (
    Array.map (fun v1 ->
        Array.map (fun v2 -> v1 *. v2) (Ps.norm_coef_scale p_b)
      ) (Ps.norm_coef_scale p_a)
    |> Array.to_list
    |> Array.concat
  ) in

  let totAngMom =
    Am.( Ps.totAngMom p_a + Ps.totAngMom p_b )
  in

  (* TODO *)
  function i ->
  function p_a -> 

      let norm_coef_a =
        Ps.norm_coef p_a
      in

      let alpha_a = (* p_a_expo_center *)
        Co.( Ps.expo p_a |. Ps.center p_a )
      in

      (*TODO *)
      function j ->
      function p_b -> 

          let norm_coef =
            norm_coef_a *. Ps.norm_coef p_b
          in

          let expo =
            Ps.expo p_a +. Ps.expo p_b
          in

          let expo_inv = 1. /. expo in

          let norm_coef = 
            let argexpo = 
              Ps.expo p_a *. Ps.expo p_b *. a_minus_b_sq *. expo_inv
            in
            norm_coef *. (pi *. expo_inv)**1.5 *. exp (-. argexpo)
          in

          function cutoff ->

              if abs_float norm_coef > cutoff then (

                  let beta_b = 
                    Co.( Ps.expo p_b |. Ps.center p_b )
                  in

                  let center =
                    Co.(expo_inv |. (alpha_a |+ beta_b))
                  in

                  let center_minus_a = 
                    Co.(center |- Ps.center p_a)
                  in

                  Some {
                    i ; j ; totAngMom ; 
                    expo ; expo_inv ; center ; center_minus_a ; a_minus_b ; 
                    a_minus_b_sq ; norm_coef ; norm_coef_scale ; shell_a = p_a;
                    shell_b = p_b }

              )
              else None

      
let make p_a p_b =
  let f =
    create_make_of p_a p_b
  in
  match f 0 p_a 0 p_b 0. with
  | Some result -> result
  | None -> assert false


let norm_coef_scale x =
  Lazy.force x.norm_coef_scale

let expo_inv x = x.expo_inv

let monocentric x =
  Ps.center x.shell_a = Ps.center x.shell_b


let totAngMom x = x.totAngMom

let a_minus_b x = x.a_minus_b

let a_minus_b_sq x = x.a_minus_b_sq

let center_minus_a x = x.center_minus_a

let norm_coef x = x.norm_coef

let expo x = x.expo

let center x = x.center

let shell_a x = x.shell_a

let shell_b x = x.shell_b