open Qcaml_common
open Qcaml_linear_algebra
open Qcaml_gaussian_basis

type t = Matrix.t

module Am  = Angular_momentum
module Bs  = Basis
module Cs  = Contracted_shell
module Ov  = Overlap


  
let make_canonical ~thresh ~basis ~cartesian ~overlap =

  let overlap_matrix =  Ov.matrix overlap in

  let make_canonical_spherical basis =
    let ao_num = Bs.size basis in
    let cart_sphe = Matrix.make ao_num ao_num 0. 
    and i = ref 0 
    and n = ref 0 in
    Array.iter (fun shell ->
      let submatrix =
          Spherical_to_cartesian.matrix (Cs.ang_mom shell)
        in
        Matrix.copy_inplace ~b:cart_sphe ~br:(!i+1) ~bc:(!n+1) submatrix;
        i := !i + Matrix.dim1 submatrix;
        n := !n + Matrix.dim2 submatrix;
      ) (Bs.contracted_shells basis);
    let s = Matrix.gemm ~transa:`T  ~m:!n cart_sphe overlap_matrix in
    let overlap_matrix = Matrix.gemm s ~n:!n cart_sphe in
    let s = Orthonormalization.canonical_ortho ~thresh ~overlap:overlap_matrix (Matrix.identity !n) in
    Matrix.gemm cart_sphe ~k:!n s
  in

  if cartesian then
    Orthonormalization.canonical_ortho ~thresh ~overlap:overlap_matrix (Matrix.identity @@ Matrix.dim1 overlap_matrix)
  else
    match basis with
    | None -> invalid_arg
              "Basis.t is required when cartesian=false in make_canonical"
    | Some basis -> make_canonical_spherical basis



let make_lowdin ~thresh ~overlap =
  
  let overlap_matrix = Ov.matrix overlap in
  let u_vec, u_val =
    Matrix.diagonalize_symm overlap_matrix
  in
  let u_vec_x = Matrix.to_bigarray_inplace u_vec in

  Vector.iter (fun x -> if x < thresh then
    invalid_arg (__FILE__^": make_lowdin") ) u_val;

  let u_val = Vector.reci (Vector.sqrt u_val) in

  let u_vec' =
    Matrix.init_cols (Matrix.dim1 u_vec) (Matrix.dim2 u_vec) (fun i j -> u_vec_x.{i,j} *. u_val.{j})
  in
  Matrix.gemm u_vec' ~transb:`T u_vec



let make ?(thresh=1.e-12) ?basis ~cartesian overlap =
  (*
  make_lowdin ~thresh ~overlap
  *)
  make_canonical ~thresh ~basis ~cartesian ~overlap