open Util
open Lacaml.D

type t = Mat.t

module Am  = AngularMomentum
module Bs  = Basis
module Cs  = ContractedShell
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 = Mat.make ao_num ao_num 0. 
    and i = ref 0 
    and n = ref 0 in
    Array.iter (fun shell ->
      let submatrix =
          SphericalToCartesian.matrix (Cs.ang_mom shell)
        in
        ignore @@ lacpy ~b:cart_sphe ~br:(!i+1) ~bc:(!n+1) submatrix;
        i := !i + Mat.dim1 submatrix;
        n := !n + Mat.dim2 submatrix;
      ) (Bs.contracted_shells basis);
    let s = gemm ~transa:`T  ~m:!n cart_sphe overlap_matrix in
    let overlap_matrix = gemm s ~n:!n cart_sphe in
    let s = canonical_ortho ~thresh ~overlap:overlap_matrix (Mat.identity !n) in
    gemm cart_sphe ~k:!n s
  in

  if cartesian then
    canonical_ortho ~thresh ~overlap:overlap_matrix (Mat.identity @@ Mat.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 = diagonalize_symm overlap_matrix in

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

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

  let u_vec' =
    Mat.init_cols (Mat.dim1 u_vec) (Mat.dim2 u_vec) (fun i j -> u_vec.{i,j} *. u_val.{j})
  in
  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