open Lacaml.D
open Util

type t =
{
  basis         : Basis.t ;
  overlap       : Overlap.t lazy_t;
  ortho         : Orthonormalization.t lazy_t;
  eN_ints       : NucInt.t lazy_t;
  kin_ints      : KinInt.t lazy_t;
  ee_ints       : ERI.t lazy_t;
  f12_ints      : F12.t lazy_t;
  cartesian     : bool;
}

let basis     t = t.basis
let overlap   t = Lazy.force t.overlap
let ortho     t = Lazy.force t.ortho
let eN_ints   t = Lazy.force t.eN_ints
let kin_ints  t = Lazy.force t.kin_ints
let ee_ints   t = Lazy.force t.ee_ints
let f12_ints  t = Lazy.force t.f12_ints
let cartesian t = t.cartesian



let make ~cartesian ~basis nuclei =

  let overlap = lazy (
    Overlap.of_basis basis
  ) in

  let ortho = lazy (
    Orthonormalization.make ~cartesian ~basis (Lazy.force overlap)
  ) in

  let eN_ints  = lazy (
    NucInt.of_basis_nuclei ~basis nuclei
  ) in

  let kin_ints = lazy (
    KinInt.of_basis basis
  ) in

  let ee_ints  = lazy (
    ERI.of_basis basis
  ) in

  let f12_ints  = lazy (
    F12.of_basis basis
  ) in

  { basis ; overlap ; ortho ; eN_ints ; kin_ints ; ee_ints ; f12_ints ;
    cartesian ;
  }



let test_case t =

  let check_matrix title a r =
    let a = Mat.to_array a in
    Array.iteri (fun i x ->
      let message =
        Printf.sprintf "%s line %d" title i
      in
      Alcotest.(check (array (float 1.e-10))) message a.(i) x
    ) (Mat.to_array r)
  in

  let check_eri title a r =
    let f { ERI.i_r1 ; j_r2 ; k_r1 ; l_r2 ; value } =
      (i_r1, (j_r2, (k_r1, (l_r2, value)))) 
    in
    let a = ERI.to_list a |> List.map f
    and r = ERI.to_list r |> List.map f
    in
    Alcotest.(check (list (pair int (pair int (pair int (pair int (float 1.e-12))))))) "ERI" a r
  in

  let test_overlap () =
    let reference =
      sym_matrix_of_file "test_files/ao_overlap.ref"
    in
    let overlap =
      Lazy.force t.overlap |> Overlap.matrix
    in
    check_matrix "Overlap" overlap reference
  in

  let test_eN_ints () =
    let reference =
      sym_matrix_of_file "test_files/ao_nuc.ref"
    in
    let eN_ints =
      Lazy.force t.eN_ints |> NucInt.matrix
    in
    check_matrix "eN_ints" eN_ints reference
  in

  let test_kin_ints () =
    let reference =
      sym_matrix_of_file "test_files/ao_kin.ref"
    in
    let kin_ints =
      Lazy.force t.kin_ints |> KinInt.matrix
    in
    check_matrix "kin_ints" kin_ints reference
  in

  let test_ee_ints () =
    let reference =
      ERI.of_file "test_files/ao_eri.ref" ~sparsity:`Dense ~size:(Basis.size t.basis)
    in
    let ee_ints =
      Lazy.force t.ee_ints 
    in
    check_eri "ee_ints" ee_ints reference
  in

  [
    "Overlap", `Quick, test_overlap;
    "eN_ints", `Quick, test_eN_ints;
    "kin_ints", `Quick, test_kin_ints;
    "ee_ints", `Quick, test_ee_ints;
  ]