(** Atomic shell pair couple *)

open Common

(** Type *)

type t =
{
  contracted_shell_pair_couples : Contracted_shell_pair_couple.t list ;
  atomic_shell_pair_p: Atomic_shell_pair.t ;
  atomic_shell_pair_q: Atomic_shell_pair.t ;
  atomic_shell_a     : Atomic_shell.t ;
  atomic_shell_b     : Atomic_shell.t ;
  atomic_shell_c     : Atomic_shell.t ;
  atomic_shell_d     : Atomic_shell.t ;
  ang_mom            : Angular_momentum.t ;
}

module Am   = Angular_momentum
module Co   = Coordinate
module As   = Atomic_shell
module Asp  = Atomic_shell_pair
module Cspc = Contracted_shell_pair_couple


(** Access *)

let contracted_shell_pair_couples t = t.contracted_shell_pair_couples

let monocentric t =
  Asp.monocentric t.atomic_shell_pair_p && Asp.monocentric t.atomic_shell_pair_q &&
    As.center (Asp.atomic_shell_a t.atomic_shell_pair_p) = As.center (Asp.atomic_shell_a t.atomic_shell_pair_q)


let ang_mom t = t.ang_mom

let atomic_shell_pair_p t = t.atomic_shell_pair_p
let atomic_shell_pair_q t = t.atomic_shell_pair_q

let atomic_shell_a t = t.atomic_shell_a
let atomic_shell_b t = t.atomic_shell_b
let atomic_shell_c t = t.atomic_shell_c
let atomic_shell_d t = t.atomic_shell_d


let zkey_array t =
  match t.contracted_shell_pair_couples with
  | f::_ -> Cspc.zkey_array f
  |  _ -> invalid_arg "AtomicShellPairCouple.zkey_array"

let norm_scales t =
  match t.contracted_shell_pair_couples with
  | f::_ -> Cspc.norm_scales f
  |  _ -> invalid_arg "AtomicShellPairCouple.norm_scales"


(** Creation *)

let make ?(cutoff=Constants.epsilon) atomic_shell_pair_p atomic_shell_pair_q =
  let ang_mom =
    Am.(Asp.ang_mom atomic_shell_pair_p + Asp.ang_mom atomic_shell_pair_q)
  in
  let atomic_shell_a = Asp.atomic_shell_a atomic_shell_pair_p
  and atomic_shell_b = Asp.atomic_shell_b atomic_shell_pair_p
  and atomic_shell_c = Asp.atomic_shell_a atomic_shell_pair_q
  and atomic_shell_d = Asp.atomic_shell_b atomic_shell_pair_q
  in
  let contracted_shell_pair_couples =
    List.concat_map (fun ap_ab ->
      List.map (fun ap_cd ->
        Cspc.make ~cutoff ap_ab ap_cd
      ) (Asp.contracted_shell_pairs atomic_shell_pair_q)
    ) (Asp.contracted_shell_pairs atomic_shell_pair_p)
    |> Util.list_some
  in
  match contracted_shell_pair_couples with
  | [] -> None
  | _  -> Some  { atomic_shell_pair_p ; atomic_shell_pair_q ; ang_mom ;
                  atomic_shell_a ; atomic_shell_b ; atomic_shell_c ; atomic_shell_d ;
                  contracted_shell_pair_couples ;
                }

(** Printers *)
let pp ppf t =
  Format.fprintf ppf "[(%d,%d),(%d,%d)]"
    (Atomic_shell.index t.atomic_shell_a)
    (Atomic_shell.index t.atomic_shell_b)
    (Atomic_shell.index t.atomic_shell_c)
    (Atomic_shell.index t.atomic_shell_d)