open Lacaml.D

(** All utilities which should be included in all source files are defined here *)

(** {2 Functions from libm} *)

external erf_float : float -> float = "erf_float_bytecode" "erf_float"
  [@@unboxed] [@@noalloc]
  (** Error function [erf] from [libm] *)

external erfc_float : float -> float = "erfc_float_bytecode" "erfc_float"
  [@@unboxed] [@@noalloc]
  (** Complementary error function [erfc] from [libm] *)

external gamma_float : float -> float = "gamma_float_bytecode" "gamma_float"
  [@@unboxed] [@@noalloc]
  (** Gamma function [gamma] from [libm] *)

external vfork : unit -> int = "unix_vfork" "unix_vfork"

(*
external popcnt : int64 -> int32 = "popcnt_bytecode" "popcnt"
  [@@unboxed] [@@noalloc]
  (** popcnt instruction *)

external trailz : int64 -> int32 = "trailz_bytecode" "trailz"
  [@@unboxed] [@@noalloc]
  (** ctz instruction *)
*)

val popcnt : int64 -> int
(** popcnt instruction *)

val trailz : int64 -> int
(** ctz instruction *)

val leadz : int64 -> int
(** ctz instruction *)



(** {2 General functions} *)

val fact : int -> float
(** Factorial function.
  @raise Invalid_argument for negative arguments or arguments >100.
  *)

val binom : int -> int -> int
(** Binomial coefficient. [binom n k] {% $= C_n^k$ %}. *)

val binom_float : int -> int -> float
(** Binomial coefficient. [binom n k] {% $= C_n^k$ %}. *)

val pow : float -> int -> float
(** Fast implementation of the power function for small integer powers *)

val chop : float -> (unit -> float) -> float
(** In [chop a f], evaluate [f] only if the absolute value of [a] is larger
  than {!Constants.epsilon}, and return [a *. f ()].
  *)

val float_of_int_fast : int -> float
(* Faster implementation of float_of_int for small positive ints *)

val of_some : 'a option -> 'a

(** {2 Functions related to the Boys function} *)

val incomplete_gamma : alpha:float -> float -> float
(** {{:https://en.wikipedia.org/wiki/Incomplete_gamma_function}
    Lower incomplete gamma function}
   @raise Failure when the calculation doesn't converge.
  *)

val boys_function : maxm:int -> float -> float array
(** {{:https://link.springer.com/article/10.1007/s10910-005-9023-3}
    Generalized Boys function}.
  @param maxm Maximum total angular momentum.
  *)


(** {2 Extension of the Array module} *)

val array_range : int -> int -> int array
(** [array_range first last] returns an array [| first; first+1 ; ... ; last-1 ; last |]. *)

val array_sum : float array -> float
(** Returns the sum of all the elements of the array *)

val array_product : float array -> float
(** Returns the product of all the elements of the array *)


(** {2 Extension of the List module} *)

val list_some : 'a option list -> 'a list
(** Filters out all [None] elements of the list, and returns the elements without
    the [Some]. *)

val list_range : int -> int -> int list
(** [list_range first last] returns a list [first; first+1 ; ... ; last-1 ; last ]. *)

val list_pack : int -> 'a list -> 'a list list
(** Example:
{[
    list_pack 3 [ 1; 2; 3; ...; 18; 19; 20 ] =
      [[1; 2; 3]; [4; 5; 6]; [7; 8; 9]; [10; 11; 12]; [13; 14; 15];
      [16; 17; 18]; [19; 20]]
]}
*)

(** {2 Useful streams} *)
val stream_range : int -> int -> int Stream.t
(** [stream_range first last] returns a stream <first ; first+1 ; ... ; last-1 ; last>. *)

val stream_to_list : 'a Stream.t -> 'a list
(** Read a stream and put items in a list. *)

val stream_fold : ('a -> 'b -> 'a) -> 'a -> 'b Stream.t -> 'a
(** Apply a fold to the elements of the stream. *)


(** {2 Linear algebra } *)

val normalize : Vec.t -> Vec.t
(** Returns a the vector normalized *)

val normalize_mat : Mat.t -> Mat.t
(** Returns the matrix with all the column vectors normalized *)

val diagonalize_symm : Mat.t -> Mat.t * Vec.t
(** Diagonalize a symmetric matrix. Returns the eigenvectors and the eigenvalues. *)

val xt_o_x : o:Mat.t -> x:Mat.t -> Mat.t
(** Computes {% $\mathbf{X^\dag\, O\, X}$ %} *)

val x_o_xt : o:Mat.t -> x:Mat.t -> Mat.t
(** Computes {% $\mathbf{X\, O\, X^\dag}$ %} *)

val remove_epsilons : Mat.t -> Mat.t
(** Remove values below machine precision *)

val qr_ortho : Mat.t -> Mat.t
(** QR orthogonalization of the input matrix *)

val canonical_ortho: ?thresh:float -> overlap:Mat.t -> Mat.t -> Mat.t
(** Canonical orthogonalization. [overlap] is the overlap matrix {% $\mathbf{S}$ %},
    and the last argument contains the vectors {% $\mathbf{C}$ %} to orthogonalize.

{%
$$
\mathbf{C_\bot} = \mathbf{C\, U\, D^{-1/2}}, \;
\mathbf{U\, D\, V^\dag} = \mathbf{S}
$$
%}
    
    *)

val debug_matrix: string -> Mat.t -> unit
(** Prints a matrix in stdout for debug *)

val matrix_of_file : string -> Mat.t
(** Reads a matrix from a file with format "%d %d %f" corresponding to
    [i, j, A.{i,j}]. *)

val sym_matrix_of_file : string -> Mat.t
(** Reads a symmetric matrix from a file with format "%d %d %f" corresponding to
    [i, j, A.{i,j}]. *)


(** {2 Printers} *)
val pp_float_array_size : Format.formatter -> float array -> unit
(** Example:
{[
[ 6:   1.000000   1.732051   1.732051   1.000000   1.732051   1.000000
  ]
]}
*)

val pp_float_array : Format.formatter -> float array -> unit
(** Example:
{[
[    1.000000   1.732051   1.732051   1.000000   1.732051   1.000000
  ]
]}
*)

val pp_float_2darray_size : Format.formatter -> float array array -> unit
(** Example:
{[
[
  2:[ 6:  1.000000   1.732051   1.732051   1.000000   1.732051   1.000000 ]
    [ 4:  1.000000   2.000000   3.000000   4.000000 ] ]
]}
*)

val pp_float_2darray : Format.formatter -> float array array -> unit
(** Example:
{[
[ [   1.000000   1.732051   1.732051   1.000000   1.732051   1.000000 ]
  [   1.000000   2.000000   3.000000   4.000000 ] ]
]}
*)

val pp_bitstring : int -> Format.formatter -> Z.t -> unit
(** Example: [ pp_bitstring 14 ppf z -> +++++------+-- ] *)

val pp_matrix : Format.formatter -> Mat.t -> unit


(** {1 Unit tests} *)

val test_case : unit -> (string * [> `Quick ] * (unit -> unit)) list