open Qcaml_linear_algebra
open Qcaml_gaussian_basis
open Qcaml_gaussian_integrals
open Qcaml_operators
  
type t =
{
  basis              :  Basis.t                     ;
  overlap            :  Overlap.t             lazy_t;
  multipole          :  Multipole.t           lazy_t;
  ortho              :  Orthonormalization.t  lazy_t;
  eN_ints            :  Electron_nucleus.t    lazy_t;
  kin_ints           :  Kinetic.t             lazy_t;
  ee_ints            :  Eri.t                 lazy_t;
  ee_lr_ints         :  Eri_long_range.t      lazy_t;
  f12_ints           :  F12.t                 lazy_t;
  f12_over_r12_ints  :  Screened_eri.t        lazy_t;
  cartesian          :  bool ;
}

let basis        t = t.basis
let overlap      t = Lazy.force t.overlap
let multipole    t = Lazy.force t.multipole
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 ee_lr_ints   t = Lazy.force t.ee_lr_ints
let f12_ints     t = Lazy.force t.f12_ints
let f12_over_r12_ints t = Lazy.force t.f12_over_r12_ints
let cartesian    t = t.cartesian

module Cs = Contracted_shell

let values t point =
  let result = Vector.create (Basis.size t.basis) in
  let resultx = Vector.to_bigarray_inplace result in
  Array.iter (fun shell ->
      Cs.values shell point
      |> Array.iteri
           (fun i_c value ->
             let i = Cs.index shell + i_c + 1 in
             resultx.{i} <- value)
    ) (Basis.contracted_shells t.basis);
  result

  
let make ~basis ?(operators=[]) ?(cartesian=false) nuclei =
  
  let overlap = lazy (
    Overlap.of_basis basis
  ) in
  
  let ortho = lazy (
    Lazy.force overlap
    |> Orthonormalization.make ~cartesian ~basis 
  ) in
  
  let eN_ints  = lazy (
    Electron_nucleus.of_basis_nuclei ~basis nuclei
  ) in
  
  let kin_ints = lazy (
    Kinetic.of_basis basis
  ) in
  
  let ee_ints  = lazy (
    Eri.of_basis basis
  ) in
  
  let rec get_f12 = function
    | (Operator.F12 _ as f) :: _ -> f
    | [] -> failwith "Missing F12 operator"
    | _ :: rest -> get_f12 rest
  in

  let rec get_rs = function
    | (Operator.Range_sep _ as r) :: _ -> r
    | [] -> failwith "Missing range-separation operator"
    | _ :: rest -> get_rs rest
  in

  let ee_lr_ints  = lazy (
    Eri_long_range.of_basis basis~operator:(get_rs operators)
  ) in
  
  let f12_ints = lazy (
    F12.of_basis basis ~operator:(get_f12 operators)
  ) in
  
  let f12_over_r12_ints = lazy (
    Screened_eri.of_basis basis ~operator:(get_f12 operators)
  ) in

  let multipole = lazy (
        Multipole.of_basis basis
      ) in

  { basis ; overlap ; multipole ; ortho ; eN_ints ; kin_ints ; ee_ints ;
    ee_lr_ints ; f12_ints ; f12_over_r12_ints ; cartesian }