(* Algorithme de localisation des orbitales *)

(* SOMMAIRE *)

(* 0. Paramètres du calcul *)

(* I. Définition des types *)

(* II : Définition des fonctions pour la localisation *)

    (* 1. Définitions de base nécessaire pour la suite *)
    (* 2. Fonction de calcul de la matrice des angles selon la méthode de Edminston *)
    (* 3. Fonction calcul de la matrice des angles selon la méthode de Boys façon Edminston  = NON FONCTIONNELLE  *)
    (* 4. Fonction de calcul de la matrice des angles selon la méthode de Boys *)
    (* 5. Pattern matching :  calcul de alpha et du critere D *)
    (* 6. Norme d'une matrice *)
    (* 7.  Détermination de alpha_max et ses indices i et j.
           Si alpha max > pi/2 on soustrait pi/2 à la matrice des alphas de manière récursive
           Si alpha max < pi/2 on ajoute pi/2 à la matrice des alphas de manière récursive *)
    (* 8. Matrice de rotation 2 par 2 *)
    (* 9. Matrice de rotation n par n avec un angle fixe  *)
    (* 10. Fonction d'extraction des 2 vecteurs propres i et j de la matrice des OMs pour les mettres dans la matrice Ksi (n par 2)
           pour appliquer R afin d'effectuer la rotation des orbitales *) (* {1,2} -> 1ere ligne, 2e colonne *)
    (* 11. Fonction de calcul de ksi~ (matrice n par 2), nouvelle matrice par application de la matrice de rotation dans laquelle
           on obtient les deux orbitales que l'on va réinjecter dans la matrice Phi*)
    (* 12. Fonction pour la création de matrice intermédiaire pour supprimer i j et ajouter i~ et j~ *)
    (* 13. Matrice intérmédiaire où les orbitales i et j ont été supprimées et remplacées par des 0, par soustraction de la matrice Phi
           par la matrice *)

(* III. Localisation des OMs *)

    (* 1. Boucle de calcul de la nouvelle matrice des coefficient après n rotation d'orbitales  *)
    (* 2. Fonction de localisation *)

(* IV : Fonctions pour la séparation et le rassemblement des matrices *)

    (* 1. Fonction de création d'une liste d'entier à partir d'un vecteur de float*)
    (* 2. Fonction créant une liste à partir des éléments manquant d'une autre liste, dans l'intervalle [1 ; n_mo] *)
    (* 3. Fonction de séparation d'une matrice en 2 sous matrice, la première matrice correspondant aux colonnes de la matrice dont le numéro est présent
          dans la liste et la seconde à celles dont le numéro de colonne n'est pas présent dans la liste *)
    (* 4. Liste des OMs occupées *)
    (* 5. Liste des OMs de coeur *)
    (* 6. Fonction de creation d'une liste des OMs pi virtuelles *)
    (* 7. Fonction de rassemblement de 2 matrices *)

(* V. Calcul *)

    (* 1. Séparation des OMs -> m_occ : OMs occupées; m_vir : OMs vacantes *)
    (* 1.1 Liste des OMs pi virtuelles *)
    (* 2. Application d'une légère rotation d'un angle fixé aux OMs pour améliorer la convergence *)
    (* 3. Légende *)
    (* 4. Rotation des orbitales : localise localisation new_X / "methode" / epsilon = 1. si <1 diminue l'angle des rotations / nb itération / critère de convergence *)
    (* 4.1 Fonction de localisation de toutes les orbitales virtuelles sans separation *)
    (* 4.2 Fonction de localisation en separant les orbitales pi et sigma *)
    (* 4.3 Fonction de localisation en separant les orbitales pi et sigma et les 10 dernières sigma *)
    (* 4.4 Fonction séparant les orbitales pi, la premiere et la seconde moitié des orbitales sigma  *)
    (* 5. Rassemblement des occupées et des virtuelles *)
    (* 5.1 Rassemblement des orbitales occupées (coeur et non coeur) *)
    (* 5.2 Pattern matching du type de separation des orbitales virtuelles et rassemblement des orbitales occupées et virtuelles *)

(* VI. Analyse *)

    (* 1. Fonction de calcul de l'extension spatiale des orbitales *)
    (* 2. Fonction de tri par ordre croissant d'une liste *)
    (* 3. Fonction de tri par ordre croissant de l'extension spatiale des orbitales *)
    (* 4. Moyenne/variance/ecart-type/mediane de l'extension spatiale des orbitales *)
    (* 5. Fonction d'affichage *)

let methode = "Boys_er";; (* Method for th orbitals localization *)
let iteration = 10000;;   (* Maximum number of iteration for each localization *)
let cc = 10e-6;;          (* Convergence criterion *)
let pre_angle = 0.001;; (* Rotation of the matrix with a small angle to improve the convergence *)
let epsilon = 1.;;        (* Rotation angle = angle /. epsilon, default value : epsilon = 1. *)

let x_end = 1;; (* Choice of the x_end for the pi/sigma/end separation, size of the second matrix with the sigma orbitals *)
(* !!! WARNING : even if you don't use this localization method, 0 < x_end < total number of sigma orbitals.
                 If you don't use this method, use the default value, x_end = 1   WARNING !!! *)

(* Choice of the separation type of orbitals *)

(* For valence ones *)
let type_separation_occu = "pi/sigma";;      (* Separation of the pi and sigma orbitals *)
(*let type_separation_occu = "1/2";;*)           (* Separation of the first and second half of the orbitals *)
(*let type_separation_occu = "pi/sigma/end";;*)  (* Separation pi,(n-x_end) sigma and x_end sigma *)
(*let type_separation_occu = "pi/1/2";;*)        (* Separation pi, (n/2) first sigma and (n/2) last sigma *)
(*let type_separation_occu = "full";;*)          (* Without separation *)

(* For occupied ones *)
let type_separation_vir = "pi/sigma";;         (* Separation of the pi and sigma orbitals *)
(*let type_separation_vir = "1/2";;*)              (* Separation of the first and second half of the orbitals *)
(*let type_separation_vir = "pi/sigma/end";;*)     (* Separation pi,(n-x_end) sigma and x_end sigma *)
(*let type_separation_vir = "pi/1/2";;*)           (* Separation pi, (n/2) first sigma and (n/2) last sigma *)
(*let type_separation_vir = "full";;*)             (* Without separation *)

open Lacaml.D

(* I. Types definitions *)

type alphaij = {
  alpha_max : float;
  indice_ii : int;
  indice_jj : int;
}

type alphad = {
  m_alpha : Mat.t;
  d : float;
}

let localize mo_basis =

  let simulation = MOBasis.simulation mo_basis in
  let nocc =
      let elec = Simulation.electrons simulation
          in
          Electrons.n_alfa elec
  in
  let ao_basis = MOBasis.ao_basis mo_basis in
  let n_core =
    (Nuclei.small_core @@ Simulation.nuclei @@ MOBasis.simulation mo_basis) / 2
  in
  let charge = Simulation.charge simulation
               |>Charge.to_int
  in

  (* II : Definition of the fonctions for the localization *)

  (* 1. Basic definitions *)
  let m_C = MOBasis.mo_coef mo_basis in
  let n_ao = Mat.dim1 m_C in
  let pi = acos(-1.) in

  (* 2. Function for Edminston-Ruedenberg localization*)
  (* Coefficient matrix n_ao x n_mo -> n_mo x n_mo matrix with angle(rad) between orbitals, Edminston localization criterion *)
  (* matrix -> matrix, float *)
  let f_alpha_er m_C =

      let ee_ints = AOBasis.ee_ints ao_basis in
      let n_mo = Mat.dim2 m_C in

      let m_b12 = Mat.init_cols n_mo n_mo (fun i j -> 0.) in
      let m_a12 = Mat.init_cols n_mo n_mo (fun i j -> 0.) in
      let v_d = Vec.init n_mo (fun i -> 0.) in

      (* Tableaux temporaires *)
      let m_pqr =
          Bigarray.(Array3.create Float64 fortran_layout n_ao n_ao n_ao)
      in
      let m_qr_i = Mat.create (n_ao*n_ao) n_mo in
      let m_ri_j = Mat.create (n_ao*n_mo) n_mo in
      let m_ij_k = Mat.create (n_mo*n_mo) n_mo in

      Array.iter (fun s ->
         (* Grosse boucle externe sur s *)
         Array.iter (fun r ->
           Array.iter (fun q ->
             Array.iter (fun p ->
               m_pqr.{p,q,r} <- ERI.get_phys ee_ints p q r s
             ) (Util.array_range 1 n_ao)
           ) (Util.array_range 1 n_ao)
         ) (Util.array_range 1 n_ao);

       (* Conversion d'un tableau a 3 indices en une matrice nao x nao^2 *)
       let m_p_qr =
         Bigarray.reshape (Bigarray.genarray_of_array3 m_pqr) [| n_ao ; n_ao*n_ao |]
         |> Bigarray.array2_of_genarray
       in

       let m_qr_i =
         (* (qr,i) = <i r|q s> = \sum_p <p r | q s> C_{pi} *)
         gemm ~transa:`T ~c:m_qr_i m_p_qr m_C
       in

       let m_q_ri =
         (* Transformation de la matrice (qr,i) en (q,ri) *)
         Bigarray.reshape_2 (Bigarray.genarray_of_array2 m_qr_i) n_ao (n_ao*n_mo)
       in

       let m_ri_j =
         (* (ri,j) = <i r | j s> = \sum_q <i r | q s> C_{bj} *)
         gemm ~transa:`T ~c:m_ri_j m_q_ri m_C
       in

       let m_r_ij =
         (* Transformation de la matrice (ri,j) en (r,ij) *)
         Bigarray.reshape_2 (Bigarray.genarray_of_array2 m_ri_j) n_ao (n_mo*n_mo)
       in

       let m_ij_k =
         (* (ij,k) = <i k | j s> = \sum_r <i r | j s> C_{rk} *)
         gemm ~transa:`T ~c:m_ij_k m_r_ij m_C
       in

       let m_ijk  =
         (* Transformation de la matrice (ei,j) en (e,ij) *)
         Bigarray.reshape (Bigarray.genarray_of_array2 m_ij_k) [| n_mo ; n_mo ; n_mo |]
         |> Bigarray.array3_of_genarray
       in

       Array.iter (fun j ->
           Array.iter (fun i ->
             m_b12.{i,j} <- m_b12.{i,j} +. m_C.{s,j} *. (m_ijk.{i,i,i} -. m_ijk.{j,i,j});
             m_a12.{i,j} <- m_a12.{i,j} +. m_ijk.{i,i,j} *. m_C.{s,j}  -.
             0.25 *. ( (m_ijk.{i,i,i} -. m_ijk.{j,i,j}) *. m_C.{s,i} +.
                       (m_ijk.{j,j,j} -. m_ijk.{i,j,i}) *. m_C.{s,j})
                 ) (Util.array_range 1 n_mo);
                  v_d.{j} <- v_d.{j} +.  m_ijk.{j,j,j} *. m_C.{s,j}
         ) (Util.array_range 1 n_mo)
    ) (Util.array_range 1 n_ao);

    (Mat.init_cols n_mo n_mo ( fun i j ->
    if i= j
        then 0.
    else if  asin(  m_b12.{i,j} /. sqrt(m_b12.{i,j}**2. +. m_a12.{i,j}**2.)) > 0.
        then 0.25 *. acos( -. m_a12.{i,j} /. sqrt(m_b12.{i,j}**2. +. m_a12.{i,j}**2.))
        else -. 0.25 *. acos( -. m_a12.{i,j} /. sqrt(m_b12.{i,j}**2. +. m_a12.{i,j}**2.)))
    ,Vec.sum v_d)

  in

  (* 3. Function for the Boys localization like the Edminston-Ruedenberg localization *)
  (* Coefficient matrix n_ao x n_mo -> n_mo x n_mo matrix with angle between orbitals, localization Boys like ER criterion *)
  (* Matrix ->  matrix, float *)
  let f_alpha_boys_er m_C =

      let multipoles = AOBasis.multipole ao_basis in
      let n_mo = Mat.dim2 m_C in
      let phi_x_phi =
          Multipole.matrix_x multipoles
          |> MOBasis.mo_matrix_of_ao_matrix ~mo_coef:m_C
      in
      let phi_y_phi =
          Multipole.matrix_y multipoles
          |> MOBasis.mo_matrix_of_ao_matrix ~mo_coef:m_C
      in
      let phi_z_phi =
          Multipole.matrix_z multipoles
          |> MOBasis.mo_matrix_of_ao_matrix ~mo_coef:m_C
      in
      let m_b12=
          let b12 g =  Mat.init_cols n_mo n_mo ( fun i j ->
                                            (g.{i,i} -. g.{j,j}) *. g.{i,j})
          in
          Mat.add (b12 phi_x_phi) ( Mat.add (b12 phi_y_phi) (b12 phi_z_phi))
      in
      let m_a12 =
          let a12 g  = Mat.init_cols n_mo n_mo ( fun i j ->
                        g.{i,j} *. g.{i,j} -. 0.25 *. ((g.{i,i} -. g.{j,j}) *. (g.{i,i} -. g.{j,j})))
          in
          Mat.add (a12 phi_x_phi) ( Mat.add (a12 phi_y_phi) (a12 phi_z_phi))

      in
      (Mat.init_cols n_mo n_mo ( fun i j ->
          if i=j
              then 0.
          else if  +. 0.25 *. atan2 m_b12.{i,j} m_a12.{i,j} >= 0.
              then pi /. 4. -. 0.25 *. atan2 m_b12.{i,j} m_a12.{i,j}
              else -. pi /. 4. -. 0.25 *. atan2 m_b12.{i,j} m_a12.{i,j})
        ,0.5 *. Vec.sum(Vec.init n_mo ( fun i -> (phi_x_phi.{i,i})**2. +. (phi_y_phi.{i,i})**2. +. (phi_z_phi.{i,i})**2.)))
  in
(* 4. Function for the original Boys localization *)
(* Coefficient matrix n_ao x n_mo -> (n_mo x n_mo matrix with angle between orbitals, localization criterion *)
(* Matrix -> matrix, float *)
  let f_theta_boys m_C =
      let multipoles = AOBasis.multipole ao_basis in
      let n_mo = Mat.dim2 m_C
      in
      let phi_x_phi =
          Multipole.matrix_x multipoles
          |> MOBasis.mo_matrix_of_ao_matrix ~mo_coef:m_C
      in
      let phi_y_phi =
          Multipole.matrix_y multipoles
          |> MOBasis.mo_matrix_of_ao_matrix ~mo_coef:m_C
      in
      let phi_z_phi =
          Multipole.matrix_z multipoles
          |> MOBasis.mo_matrix_of_ao_matrix ~mo_coef:m_C
      in
      let phi_x2_phi =
          Multipole.matrix_x2 multipoles
          |> MOBasis.mo_matrix_of_ao_matrix ~mo_coef:m_C
      in
      let phi_y2_phi =
            Multipole.matrix_y2 multipoles
            |> MOBasis.mo_matrix_of_ao_matrix ~mo_coef:m_C
      in
      let phi_z2_phi =
            Multipole.matrix_z2 multipoles
            |> MOBasis.mo_matrix_of_ao_matrix ~mo_coef:m_C
      in
      let m_b_0 =
          let b_0 g h = Mat.init_cols n_mo n_mo (fun i j ->
              h.{i,i} +. h.{j,j} -. (g.{i,i}**2. +. g.{j,j}**2.))
          in
          Mat.add (b_0 phi_x_phi phi_x2_phi) (Mat.add (b_0 phi_y_phi phi_y2_phi) (b_0 phi_z_phi phi_z2_phi))
      in
      let m_beta =
          let beta g = Mat.init_cols n_mo n_mo (fun i j ->
              (g.{i,i} -.  g.{j,j})**2. -. 4. *. g.{i,j}**2.)
          in
          Mat.add (beta phi_x_phi) (Mat.add (beta phi_y_phi) (beta phi_z_phi))
      in
      let m_gamma =
          let gamma g = Mat.init_cols n_mo n_mo (fun i j ->
              4. *. g.{i,j} *. (g.{i,i} -. g.{j,j}) )
          in
           Mat.add (gamma phi_x_phi) (Mat.add (gamma phi_y_phi) (gamma phi_z_phi))
      in
      let m_theta = Mat.init_cols n_mo n_mo (fun i j ->
          if i = j
              then 0.
              else +. 0.25 *. atan2 m_gamma.{i,j} m_beta.{i,j})
      in
      let m_critere_B = Mat.init_cols n_mo n_mo (fun i j ->
         0.5 *. (m_b_0.{i,j} +. 0.25 *. ((1. -. cos(4. *. m_theta.{i,j})) *. m_beta.{i,j} +. sin(4. *. m_theta.{i,j}) *. m_gamma.{i,j})))

      in
      m_theta, Vec.sum(Mat.as_vec m_critere_B)
  in

  (* Function to compute the new angle after rotation between the 2 orbitals j1 and j2, by the Boys like ER localization *)
  (* n_mo x n_mo matrix with angle between orbitals, index of the orbital j1, index of the orbital j2 -> n_mo x 2 matrix with the new angles between j1 and j2, new localization criterion *)
  (* matrix, integer, integer -> matrix, float *)
  let f2_alpha_boys_er m_C j1 j2=
    let multipoles = AOBasis.multipole ao_basis in
    let n_mo = Mat.dim2 m_C
    in
    let phi_x_phi =
          Multipole.matrix_x multipoles
          |> MOBasis.mo_matrix_of_ao_matrix ~mo_coef:m_C
    in
    let phi_y_phi =
          Multipole.matrix_y multipoles
          |> MOBasis.mo_matrix_of_ao_matrix ~mo_coef:m_C
    in
    let phi_z_phi =
          Multipole.matrix_z multipoles
          |> MOBasis.mo_matrix_of_ao_matrix ~mo_coef:m_C
    in

    let m_b12=
        let b12 g =  Mat.init_cols n_mo 2 ( fun i j ->
            if j = 1
                then  (g.{i,i} -. g.{j1,j1}) *. g.{i,j1}
                else  (g.{i,i} -. g.{j2,j2}) *. g.{i,j2})

        in
        Mat.add (b12 phi_x_phi) ( Mat.add (b12 phi_y_phi) (b12 phi_z_phi))
    in
    let m_a12 =
        let a12 g  = Mat.init_cols n_mo 2 ( fun i j ->
            if j = 1
                then  g.{i,j1} *. g.{i,j1} -. 0.25 *. ((g.{i,i} -. g.{j1,j1}) *. (g.{i,i} -. g.{j1,j1}))
                else  g.{i,j2} *. g.{i,j2} -. 0.25 *. ((g.{i,i} -. g.{j2,j2}) *. (g.{i,i} -. g.{j2,j2})))
        in
        Mat.add (a12 phi_x_phi) ( Mat.add (a12 phi_y_phi) (a12 phi_z_phi))
    in
    (Mat.init_cols n_mo 2 ( fun i j ->
        if i=j1 && j=1
            then 0.
        else if i = j2 && j=2
            then 0.
        else
          let x = atan2 m_b12.{i,j} m_a12.{i,j} in
          if x >= 0.  then 0.25 *. (pi -. x)
          else -. 0.25 *. (pi +. x)
        ),
        Vec.sum @@
        Vec.init n_mo ( fun i ->
           (phi_x_phi.{i,i})**2. +. (phi_y_phi.{i,i})**2. +. (phi_z_phi.{i,i})**2.))
  in

  (* Function to compute the new angle after rotation between the 2 orbitals j1 and j2, by the Boys localization *)
  (* n_mo x n_mo matrix with angles between orbitals, index of the orbital j1, index of the orbital j2 -> n_mo x 2 matrix with the new angles between j1 and j2, new localization criterion *)
  (* matrix, integer, integer -> matrix, float *)
  let f2_theta_boys m_C j1 j2 =
        let multipoles = AOBasis.multipole ao_basis in
        let n_mo = Mat.dim2 m_C
        in
        let phi_x_phi =
              Multipole.matrix_x multipoles
              |> MOBasis.mo_matrix_of_ao_matrix ~mo_coef:m_C
        in
        let phi_y_phi =
              Multipole.matrix_y multipoles
              |> MOBasis.mo_matrix_of_ao_matrix ~mo_coef:m_C
        in
        let phi_z_phi =
              Multipole.matrix_z multipoles
              |> MOBasis.mo_matrix_of_ao_matrix ~mo_coef:m_C
        in
        let phi_r2_phi =
              Multipole.matrix_r2 multipoles
              |> MOBasis.mo_matrix_of_ao_matrix ~mo_coef:m_C
        in

        let m_beta =
            let beta g = Mat.init_cols n_mo 2 (fun i j ->
                if j = 1
                    then (g.{i,i} -. g.{j1,j1})**2. -. 4. *. g.{i,j1}**2.
                    else (g.{i,i} -. g.{j2,j2})**2. -. 4. *. g.{i,j2}**2.)
            in Mat.add (beta phi_x_phi) (Mat.add (beta phi_y_phi) (beta phi_z_phi))
        in

        let m_gamma =
            let gamma g = Mat.init_cols n_mo 2 (fun i j ->
                if j = 1
                   then 4. *. g.{i,j1} *. (g.{i,i} -. g.{j1,j1})
                   else 4. *. g.{i,j2} *. (g.{i,i} -. g.{j2,j2}))
            in Mat.add (gamma phi_x_phi) (Mat.add (gamma phi_y_phi) (gamma phi_z_phi))
        in

        let m_theta = Mat.init_cols n_mo 2 (fun i j ->
            if i=j1 && j=1
                then 0.
            else if i = j2 && j=2
                then 0.
                else +. 0.25 *. atan2 m_gamma.{i,j} m_beta.{i,j})
        in
        let m_critere_B = Vec.init n_mo (fun i ->
            phi_r2_phi.{i,i} -. phi_x_phi.{i,i}**2. -. phi_y_phi.{i,i}**2. -. phi_z_phi.{i,i}**2.)

        in m_theta, Vec.sum( m_critere_B)
   in

  (* 5. Pattern matching : choice of the localization method *)
  (* The method, coefficient matrix n_ao x n_mo -> n_mo x n_mo matrix with the angle between orbitals using previous function, localization criterion *)
  (* string, matrix -> matrix, float *)

  let m_alpha_d methode m_C =

      let alpha methode =
          match methode with
              | "Boys_er"
              | "BOYS_ER"
              | "boys_er" -> let alpha_boys , d_boys  = f_alpha_boys_er m_C in {m_alpha = alpha_boys; d = d_boys}
              | "Boys"
              | "boys"
              | "BOYS" -> let theta_boys, b_boys = f_theta_boys m_C in {m_alpha = theta_boys; d = b_boys}
              | "ER"
              | "er" -> let alpha_er , d_er = f_alpha_er m_C in {m_alpha = alpha_er; d = d_er}
              | _ -> invalid_arg "Unknown method, please enter Boys or ER"

      in
      alpha methode
  in

  (* 5. Pattern matching : choice of the localization method to compute the n_mo x 2 new angles and not the n_mo x n_mo angles, after the rotation of i and j orbitals *)
  (* The method, coefficient matrix n_ao x n_mo, index of the orbital i , index of the orbital j -> n_mo x 2 matrix with the new angle between orbitals i and j using previous function, new localization criterion *)
  (* string, matrix, integer, integer -> matrix, float *)
  let f_speed_alpha methode m_C indice_i indice_j =
    let alpha_speed methode =

        match methode with
            | "Boys_er"
            | "boys_er" -> let m_alpha, critere_D = f2_alpha_boys_er m_C indice_i indice_j in {m_alpha = m_alpha; d = critere_D}
            | "BOYS"    -> let m_alpha, critere_D = f2_theta_boys m_C indice_i indice_j in {m_alpha = m_alpha; d = critere_D}
            | _ -> invalid_arg "Unknown method, please enter Boys or ER"

      in
      alpha_speed methode

  in
  (* 6. Norm of a matrix *)
  (* Matrix -> float *)
  let norm m =

      let vec_m = Mat.as_vec m in
      let vec2 = Vec.sqr vec_m
      in
      sqrt(Vec.sum vec2)
  in

  (* 7. Calculation of the max angle[-pi/.4 , pi/.4]  and its indexes i and j in the matrix *)
  (* n_mo x n_mo matrix with angles between orbitals, coefficient matrix n_ao x n_mo, max number of iteration -> max angle, line location, column location *)
  (* matrix, matrix, integer -> float, integer, integer *)
  let rec new_m_alpha m_alpha m_C n_rec_alpha=

      let n_mo = Mat.dim2 m_C in
      let alpha_m =
          if n_rec_alpha == 0
              then m_alpha
              else
                  Mat.init_cols n_mo n_mo (fun i j ->
                  if  (m_alpha.{i,j}) >= (pi /. 2.)
                      then (m_alpha.{i,j} -. ( pi /. 2.))
                  else if m_alpha.{i,j} <= -. pi /. 2.
                      then (m_alpha.{i,j} +. ( pi /. 2.))
                      else m_alpha.{i,j} )
      in

      (* Location of the max angle in the matrix *)
      (* angle matrix -> location of the max element *)
      (* matrix -> integer *)
      let max_element alpha_m =
          Mat.as_vec alpha_m
          |> iamax
      in

      (* Indexes i and j of the max angle *)
      (* -> integer, integer *)
      let indice_ii, indice_jj =
          let max = max_element alpha_m
          in
          (max - 1) mod n_mo +1, (max - 1) / n_mo +1
      in

      (* Value of the max angle *)
      (* angle matrix -> value of the max angle*)
      (* matrix -> float *)
      let alpha alpha_m  =
          let i = indice_ii in
          let j = indice_jj
          in
          alpha_m.{i,j}
      in
      let alpha_max = alpha alpha_m

      in
      if alpha_max < pi /. 2.
          then {alpha_max; indice_ii; indice_jj}
          else new_m_alpha alpha_m m_C (n_rec_alpha-1)
  in

  (* 8. Rotation matrix 2x2 *)
  (* angle -> 2x2 rotation matrix *)
  (* float -> matrix *)
  let f_R alpha  =

      Mat.init_cols 2 2 (fun i j ->
          if i=j
              then cos alpha
          else if i>j
              then  sin alpha
              else -. sin alpha )
  in

  (* 9. Rotation matrix nxn *)
  (* Coefficient matrix n_ao x n_mo, angle -> new coefficient matrix n_ao x n_mo after a rotation with a small angle *)
  (* matrix, float -> matrix *)
  let rotate m_C angle=

      let n_mo =  Mat.dim2 m_C in

      (* Antisymmetrization of the matrix and diagonal terms = 0, X *)
      (* angle -> Antisymmetric matrix with this angle and diagonal terms = 0 *)
      (* float -> matrix *)
      let m_angle angle=

          Mat.init_cols n_mo n_mo (fun i j ->
              if i = j
                  then 0.
              else if i>j
                  then -. angle
                  else angle)
      in

      let m = m_angle angle in

      (* Square of the matrix X² *)
      (* -> X . X *)
      (* -> matrix*)
      let mm = gemm m m in

      (* Diagonalization of X² *)
      (* diag = vector that contains the eigenvalues (-tau²), mm contains the eigenvectors (W) *)
      (* -> vector  *)
      let diag = syev ~vectors:true mm in

      (* -tau² to tau² *)
      (* -> vector that contains the eigenvalues (tau²) *)
      (* -> vector *)
      let square_tau= Vec.abs diag in

      (* tau² to tau *)
      (* -> vector that contains the eigenvalues tau *)
      (* -> vector *)
      let tau = Vec.sqrt square_tau in

      (* Calculation of cos tau from the vector tau *)
      (* -> cos(tau) *)
      (* -> integer *)
      let cos_tau = Vec.cos tau in

      (* Matrix cos tau *)
      (* -> matrix with diagonal terms cos(tau) and 0 also *)
      (* -> matrix *)
      let m_cos_tau =

          Mat.init_cols n_mo n_mo (fun i j ->
              if i=j
                  then cos_tau.{i}
                  else 0.)
      in

      (* Calcul of sin(tau) from the vector tau *)
      let sin_tau = Vec.sin tau in

      (* Matrix sin tau *)
      (* -> matrix with diagonal terms sin(tau) and 0 also *)
      (* -> matrix *)
      let m_sin_tau =

          Mat.init_cols n_mo n_mo (fun i j ->
              if i=j
                  then sin_tau.{i}
                  else 0.)
      in

      (* Calculation of the transposed matrix of W (X²) *)
      (* -> transposed matrix of W *)
      (* -> matrix *)
      let transpose_mm = Mat.transpose_copy mm in

      (* Calculation of the vector tau^-1 *)
      (* -> vector tau^-1 *)
      (* -> vector *)
      let tau_1 =

          Vec.init n_mo (fun i ->
              if tau.{i}<= 0.001
                  then 1.
                  else 1. /. tau.{i})
      in
      (* Calculation of the matrix tau^⁻1 *)
      (* -> matrix tau^-1 *)
      (* -> matrix *)
      let m_tau_1 =

          Mat.init_cols n_mo n_mo (fun i j ->
              if i=j
                  then tau_1.{i}
                  else 0.)
      in

      (* gemm mm (gemm m_cos_tau transpose_mm) -> W cos(tau) Wt *)
      (* -> matrix *)
      let a = gemm mm (gemm m_cos_tau transpose_mm) in

      (* gemm mm (gemm m_tau_1 ( gemm m_sin_tau (gemm transpose_mm m))) -> W tau^-1 sin(tau) Wt X *)
      (* -> matrix *)
      let b = gemm mm (gemm m_tau_1 ( gemm m_sin_tau (gemm transpose_mm m))) in

      (* Sum of a + b -> Rotation matrix *)
      (* -> Rotation matrix *)
      (* -> matrix *)
      let m_r = Mat.add a b
      in
      gemm m_C m_r
  in

  (* Function to extract 2 vectors i and j in a matrix *)
  (* Coefficient matrix n_ao x n_mo, index of the orbital i,  index of the orbital j -> n_mo x 2 matrix with the eigenvetors i and j*)
  (* matrix, integer, integer -> matrix *)
  let f_Ksi indice_i indice_j m_C  =
      let n_ao = Mat.dim1 m_C

      in
      Mat.init_cols n_ao 2 (fun i j ->
          if j=1
              then m_C.{i,indice_i}
              else m_C.{i,indice_j} )
  in

  (* Function to apply a rotation with a 2x2 rotation matrix on a n_mo x 2 matrix that contains i and j *)
  (* 2x2 rotation matrix, n_mo x 2 eigenvetors(i and j) matrix -> n_mo x 2 new eigenvetors(i~ and j~) after rotation *)
  (* matrix, matrix -> matrix *)
  let f_Ksi_tilde m_R m_Ksi = gemm m_Ksi m_R
  in

  (* Function to create intermediate coefficient matrix in order to delete the i j orbitals and add i~ j~ in the coefficient matrix *)
  (* n_mo x 2 matrix with k and l orbitals, index of k, index of l, n_ao x n_mo matrix coefficient -> matrix where all terms are equal to 0 except the k and l columns that contain the eigenvetors k and l *)
  (* matrix, integer, integer, matrix -> matrix *)
  let f_k mat indice_i indice_j m_C =
      let n_mo = Mat.dim2 m_C in
      let n_ao = Mat.dim1 m_C

      in
      Mat.init_cols n_ao n_mo (fun i j ->
          if j=indice_i
              then mat.{i,1}
          else if j=indice_j
              then mat.{i,2}
          else 0.)
  in

  (* Function to create a intermediate angle matrix in order to delete all the previous angle between i and all the other orbitals, j and all the other orbitals.
  And add the new angles between these orbitals *)
  (* n_mo x 2 matrix with angles between k and all the other orbitals l and all the other orbitals, index of k, index of l, coefficent matrix n_ao x n_mo -> matrix where all terms are equal to 0 except the k and l lines and columns that contains the angles between the orbitals *)
  (* matrix, integer, integer, matrix -> matrix  *)
  let f_k_angle mat indice_i indice_j m_C =
      let n_mo = Mat.dim2 m_C

      in
      Mat.init_cols n_mo n_mo (fun i j ->
          if j=indice_i
              then mat.{i,1}
          else if j=indice_j
              then mat.{i,2}
          else if i = indice_i
              then -. mat.{j,1}
          else if i = indice_j
              then -. mat.{j,2}
          else 0.)
  in

  (* Intermediate matrix where the i and j orbitals are equal to 0 *)
  (* Coefficient matrix n_ao x n_mo, matrix where terms are equal to 0 except the k and l columns that contain the eigenvetors k and l *)
  (* matrix, matrix -> matrix *)
  let f_interm m_C mat =

      Mat.sub m_C mat
  in

  (* Function to compute the new coefficient matrix after a rotation between 2 orbitals *)
  (* coefficient matrix n_ao x n_mo, angle matrix, parameter -> new coefficient matrix n_ao x n_mo *)
  let new_m_C m_C m_alpha epsilon =

      (* alphaij contains the max angle of the angle matrix and the indexes *)
      let n_rec_alpha = 10 in
      let alphaij = new_m_alpha m_alpha m_C n_rec_alpha in

      (* Value of the angle *)
      let alpha = (alphaij.alpha_max) *. epsilon in

      (* Localtion of the max angle *)
      let indice_i = alphaij.indice_ii in
      let indice_j = alphaij.indice_jj in

      (* Rotation matrix *)
      let m_R = f_R alpha in

      (* Matrix with i and j orbitals *)
      let m_Ksi = f_Ksi indice_i indice_j m_C in

      (* Matrix with the new i~ and j~ orbitals *)
      let m_Ksi_tilde = f_Ksi_tilde m_R m_Ksi in

      (* Matrix to delete the i and j orbitals in the coefficient matrix *)
      let m_Psi = f_k m_Ksi indice_i indice_j m_C in

      (* Matrix to add the i~ and j~ orbitals in the coefficient matrix *)
      let m_Psi_tilde = f_k m_Ksi_tilde indice_i indice_j m_C in

      (* Coefficient matrix without the i and j orbitals *)
      let m_interm = f_interm m_C m_Psi

      in
      (* Matrice after rotation, max angle, index i, index j *)
      (Mat.add m_Psi_tilde m_interm, alpha, indice_i, indice_j)
  in

  (* Function to compute the new angle matrix after one rotation between 2 orbitals *)
  (* Previous angle matrix n_mo x n_mo, new coefficient matrix n_ao x n_mo, previous angle matrix, index of the orbital i, index of the orbital j -> new angle matrix, new localization criterion *)
  (* matrix, matrix, matrix, integer, integer -> matrix, float *)
  let m_alpha m_new_m_C prev_m_alpha indice_i indice_j =

      (* n_mo x 2 matrix that contains the new angles between i,j with the other orbitals *)
      let speed_alphad = f_speed_alpha methode m_new_m_C indice_i indice_j in

      (* New localization criterion *)
      let critere_D = speed_alphad.d in

      (* Previous angles *)
      let ksi_angle = f_Ksi indice_i indice_j prev_m_alpha in

      (* New angles *)
      let ksi_angle_tilde = speed_alphad.m_alpha in

      (* Matrix to delete the previous angles *)
      let psi_angle = f_k_angle ksi_angle indice_i indice_j m_new_m_C in

      (* Matrix to add the new angles *)
      let psi_tilde_angle = f_k_angle ksi_angle_tilde indice_i indice_j m_new_m_C in

      (* Matrix without the angles between i, j and the orbitals *)
      let m_interm = f_interm prev_m_alpha psi_angle

      in
      (* New angle matrix, new localization criterion *)
      (Mat.add m_interm psi_tilde_angle, critere_D)
  in



(* III. Localization *)

  (* Loop to compute the new coefficient matrix avec n orbital rotations *)
  (* Coefficeient matrix n_ao x n_mo, localization method, parameter, max number of iteration, previous localization criterion, previous angle matrix, convergence criterion -> new coefficient matrix after n orbital rotations *)
  (* matrix, string, float, integer, float, matrix, float -> matrix *)
  let rec localisation m_C methode epsilon n prev_critere_D prev_m_alpha cc =

      if n == 0
        then m_C
        else

           (* New coefficient matrix after one rotation *)
           let m_new_m_C, alpha_max, indice_i, indice_j = new_m_C m_C prev_m_alpha epsilon in

           (* New angle matrix after one rotation *)
           let m_alpha, critere_D = m_alpha m_new_m_C prev_m_alpha indice_i indice_j  in

           (* Matrix norm, angle average  *)

           let norm_alpha = norm m_alpha in
           let moyenne_alpha = norm_alpha /. (float_of_int((Mat.dim1 m_C)* (Mat.dim2 m_C))) in
           let alpha = alpha_max /. epsilon in

           Printf.printf "%i %f %f %f %f\n%!" (iteration-n) critere_D alpha norm_alpha moyenne_alpha;

           (* Convergence criterion *)
           if alpha_max**2.  < cc**2.
               then m_new_m_C
               else localisation m_new_m_C methode epsilon (n-1) critere_D m_alpha cc
  in

  (* 2. Localization function *)
  let localise localisation m_C methode epsilon n cc =

      let alphad = m_alpha_d methode m_C in
      let prev_m_alpha = alphad.m_alpha in
      let prev_critere_D = 0.
      in
      localisation m_C methode epsilon n prev_critere_D prev_m_alpha cc
  in

(* IV : Functions to split and assemble matrix *)

  (* Function to create a integer list from a vector of float *)
  (* float vector -> integer list *)
  let int_list vec =

      let float_list = Vec.to_list vec in
      let g a = int_of_float(a)
      in
      List.map g float_list
  in

  (* Function to create a list from the missing elements of an other list*)
  (* Coefficient matrix n_ao x n_mo, list -> list of missing element of the previous list *)
  (* matrix, integer list -> integer list *)
  let miss_elem mat list =

      let n_mo = Mat.dim2 mat in
      let vec = Vec.init (n_mo) (fun i ->
          if List.mem i list
              then 0.
              else float_of_int(i))
      in
      let list_int = int_list vec
      in
      List.filter (fun x -> x > 0) list_int
  in

  (* Function to split a matrix in 2 matrix, the first matrix corresponds to the column number whose number is in the list, the second matrix corresponds to the column which are not in the list *)
  (* Coefficient matrix n_ao x n_mo, list -> matrix, matrix *)
  (* matrix, integer list -> matrix, matrix*)
  let split_mat mat list =

      let vec_of_mat = Mat.to_col_vecs mat in
      let f a = vec_of_mat.(a-1) in
      let vec_list_1 = List.map f list in
      let list_2 = miss_elem mat list in
      let vec_list_2 = List.map f list_2
      in
      (Mat.of_col_vecs_list vec_list_1,Mat.of_col_vecs_list vec_list_2)
  in


  (* Function to create a matrix from a list*)
  (* Coefficient matrix n_ao x n_mo, list of orbitals -> matrix with these orbitals *)
  (* matrix, integer list -> matrix*)
  let create_mat mat list =

      let vec_of_mat = Mat.to_col_vecs mat in
      let f a = vec_of_mat.(a-1) in
      let vec_list = List.map f list
      in
      Mat.of_col_vecs_list vec_list
  in

  (* List of the occupied MOs *)
  let list_occ =

      let vec_occ = Vec.init (nocc) (fun i -> float_of_int(i))
      in
      int_list vec_occ
  in

  (* List of the core MOs *)
  let list_core =
      let vec_core = Vec.init (n_core) (fun i -> float_of_int(i))
      in int_list vec_core
  in

  (* Function to create a list of pi MOs *)
  (* Coefficient matrix n_ao x n_mo -> list of pi orbitals *)
  (* matrix -> integer list *)
  let list_pi mat =
      let n_mo = Mat.dim2 mat in
      let m_pi = Mat.init_cols n_ao (n_mo+1) ( fun i j ->
          if j=1 && m_C.{i,j}**2. < 10e-8**2.
              then 0.
          else if j=1
              then m_C.{i,j}
          else if mat.{i,j-1}**2. < 10e-8**2.
              then 0.
              else mat.{i,j-1})
      in
      let vec = Mat.to_col_vecs m_pi in
      let vec_dot = Vec.init ((Mat.dim2 mat)) (fun i -> dot vec.(0) vec.(i)) in
      let vec_pi = Vec.init ((Mat.dim2 mat)) (fun i ->
          if vec_dot.{i} = 0.
              then float_of_int(i)
              else 0.)
      in
      let list_pi = int_list vec_pi in
      let rec remove x list =
          match list with
            | [] -> []
            | var :: tail -> if var = x
                       then remove x tail
                       else var :: (remove x tail)
      in
      remove 0 list_pi
  in

  (* Function to create a list of the (n-x_end) first MOs of a matrix *)
  (* Coefficient matrix n_ao x n_mo -> list of the (n-x_end) first MOs of the matrix *)
  (* matrix -> integer list *)
  let list_x_end mat =
      let n = Mat.dim2 mat in
      let vec_x_end = Vec.init (n-x_end)  (fun i -> float_of_int(i))
      in
      int_list vec_x_end
  in

  (* Function to create a list of the (n/2) first MOs of a matrix *)
  (* Coefficient matrix n_ao x n_mo -> list of the (n/2) first MOs of the matrix *)
  (* matrix -> integer list *)
  let list_s_12 mat =
      let n = Mat.dim2 mat in
      let vec_12 = Vec.init (n/2)  (fun i -> float_of_int(i))
      in
      int_list vec_12
  in

  (* Function to assemble to matrix with the same number of line *)
  (* Coefficient matrix n_ao x n, coefficient matrix n_ao x m -> Coefficient matrix n_ao x (n+m) *)
  (* matrix, matrix -> matrix*)
  let assemble mat_1 mat_2 =

      let v_mat_1 = Mat.to_col_vecs mat_1 in
      let v_mat_2 = Mat.to_col_vecs mat_2 in
      let m = Array.append v_mat_1 v_mat_2
      in
      Mat.of_col_vecs m
  in

(* V. Calculation  *)

   (* 1. Caption *)
  let text = "Method : " ^ methode ^ "; Separation type of the valence occupied MOs : " ^ type_separation_occu  ^ "; Separation type of the virtual MOs : " ^ type_separation_vir ^ "; Max number of iteration : " ^ string_of_int(iteration) ^ "; cc : " ^ string_of_float(cc)^"; epsilon : "^string_of_float(epsilon)^"; charge : "^string_of_int(charge)^"; Rotation pre angle : "^string_of_float(pre_angle)
  in
  let caption = "n Criterion max_angle norm_angle average_angle"
  in
  Printf.printf "%s\n" "";
  Printf.printf "%s\n" "Orbitals localization ";
  Printf.printf "%s\n" "";
  Printf.printf "%s\n" text;
  Printf.printf "%s\n" "";
  Printf.printf "%s\n" caption;
  Printf.printf "%s\n" "";

  (* Separation of the occupied and virtual MOs, valence and core MOs *)
  let m_occ , m_vir = split_mat m_C list_occ in

  let m_core, m_occu = split_mat m_occ list_core in

  (* Localization function *)
  (* Coefficient matrix n_ao x n_mo -> localized coefficient matrix n_ao x n_mo *)
  (* matrix -> matrix *)
  let f_localise mat =

      if Mat.dim2 mat = 0
          then Printf.printf "%s\n" "0 orbital, no localization";

      let new_mat =
          if Mat.dim2 mat = 0
              then mat
              else rotate mat pre_angle
      in
      if Mat.dim2 mat = 0
          then new_mat
          else localise localisation new_mat methode epsilon iteration cc
  in

  (* Core MOs localization *)
  let loc_m_core =
          Printf.printf "%s\n" "I/III";
          Printf.printf "%s\n" "Localization of core orbitals";
          Printf.printf "%s %i\n" "Number of orbitals" (Mat.dim2 m_core);
          Printf.printf "%s\n" "";

          f_localise m_core in

          Printf.printf "%s\n" "";
          Printf.printf "%s\n" "End I/III";
          Printf.printf "%s\n" "";

  (* Localization function without separation *)
  let f_loc_assemble_1 mat =

      Printf.printf "%s\n" "";
      Printf.printf "%s\n" "Localization 1/1";
      Printf.printf "%s %i\n" "Number of orbitals" (Mat.dim2 mat);

      f_localise mat in

      Printf.printf "%s\n" "";

  (* Localization function with 1 MOs separation *)
  let f_loc_assemble_12 mat list_12  =

      if List.length(list_12 mat) + List.length(miss_elem mat (list_12 mat)) <> (Mat.dim2 mat)
          then Printf.printf "%s\n" "Bug : List length problem (list_12 or miss_elem list_12)";

      let mat_1, mat_2 = split_mat mat (list_12 mat) in

      Printf.printf "%s\n" "";
      Printf.printf "%s\n" "Localization 1/2";
      Printf.printf "%s %i\n" "Number of orbitals" (Mat.dim2 mat_1);
      let m_loc_1 = f_localise mat_1 in
      Printf.printf "%s\n" "";

      Printf.printf "%s\n" "Localization 2/2";
      Printf.printf "%s %i\n" "Number of orbitals" (Mat.dim2 mat_2);
      let m_loc_2 = f_localise mat_2 in
      Printf.printf "%s\n" "";

      assemble m_loc_1 m_loc_2
  in

  (* Localization function with 2 MOs separations *)
  let f_loc_assemble_123 mat list_12 list_2ab =

      if List.length(list_12 mat) + List.length(miss_elem mat (list_12 mat)) <> (Mat.dim2 mat)
          then Printf.printf "%s\n" "Bug : List length problem (list_12 or miss_elem list_12)";

      let mat_1, mat_2 = split_mat mat (list_12 mat) in

      if List.length(list_2ab mat_2) + List.length(miss_elem mat_2 (list_2ab mat_2)) <> (Mat.dim2 mat_2)
          then Printf.printf "%s\n" "Bug : List length problem (list_2ab or miss_elem list_2ab)";

      let mat_2a, mat_2b = split_mat mat_2 (list_2ab mat_2) in

      Printf.printf "%s\n" "";
      Printf.printf "%s\n" "Localization 1/3";
      Printf.printf "%s %i\n" "Number of orbitals" (Mat.dim2 mat_1);
      let m_loc_1 = f_localise mat_1 in
      Printf.printf "%s\n" "";

      Printf.printf "%s\n" "Localization 2/3";
      Printf.printf "%s %i\n" "Number of orbitals" (Mat.dim2 mat_2a);
      let m_loc_2a = f_localise mat_2a in
      Printf.printf "%s\n" "";

      Printf.printf "%s\n" "Localization 3/3";
      Printf.printf "%s %i\n" "Number of orbitals" (Mat.dim2 mat_2b);
      let m_loc_2b = f_localise mat_2b in
      Printf.printf "%s\n" "";

      let m_loc_2ab = assemble m_loc_2a m_loc_2b
      in

      assemble m_loc_1 m_loc_2ab
  in

  (* Localization function with 3 MOs separations *)
  let f_loc_assemble_1234 mat list_12 list_1ab list_2ab  =

      if List.length(list_12 mat) + List.length(miss_elem mat (list_12 mat)) <> (Mat.dim2 mat)
          then Printf.printf "%s\n" "Bug : List length problem (list_12 or miss_elem list_12)";

      let mat_1, mat_2 = split_mat mat (list_12 mat) in

      if List.length(list_1ab mat_1) + List.length(miss_elem mat (list_1ab mat_1)) <> (Mat.dim2 mat_1)
          then Printf.printf "%s\n" "Bug : List length problem (list_1ab or miss_elem list_1ab)";

      let mat_1a, mat_1b = split_mat mat_1 (list_1ab mat_1) in

      if List.length(list_2ab mat_2) + List.length(miss_elem mat (list_2ab mat_1)) <> (Mat.dim2 mat_2)
          then Printf.printf "%s\n" "Bug : List length problem (list_2ab or miss_elem list_2ab)";

      let mat_2a, mat_2b = split_mat mat_2 (list_2ab mat_2) in

      Printf.printf "%s\n" "";
      Printf.printf "%s\n" "Localization 1/4";
      Printf.printf "%s %i\n" "Number of orbitals" (Mat.dim2 mat_1a);
      let m_loc_1a = f_localise mat_1a in
      Printf.printf "%s\n" "";

      Printf.printf "%s\n" "Localization 2/4";
      Printf.printf "%s %i\n" "Number of orbitals" (Mat.dim2 mat_1b);
      let m_loc_1b = f_localise mat_1b in
      Printf.printf "%s\n" "";

      Printf.printf "%s\n" "Localization 3/4";
      Printf.printf "%s %i\n" "Number of orbitals" (Mat.dim2 mat_2a);
      let m_loc_2a = f_localise mat_2a in
      Printf.printf "%s\n" "";

      Printf.printf "%s\n" "Localization 4/4";
      Printf.printf "%s %i\n" "Number of orbitals" (Mat.dim2 mat_2b);
      let m_loc_2b = f_localise mat_2b in
      Printf.printf "%s\n" "";

      let m_loc_1ab = assemble m_loc_1a m_loc_1b in
      let m_loc_2ab = assemble m_loc_2a m_loc_2b
      in
      assemble m_loc_1ab m_loc_2ab
  in

  (* Localization function with 4 separations "by hands" *)
  let f_loc_assemble_main mat list_1 list_2 list_3 list_4 =

      if (List.length(list_1) + List.length(list_2) + List.length(list_3) + List.length(list_4)) <> (Mat.dim2 mat)
          then Printf.printf "%s\n" "Bug : List length problem";

      let mat_1 = create_mat mat list_1 in
      let mat_2 = create_mat mat list_2 in
      let mat_3 = create_mat mat list_3 in
      let mat_4 = create_mat mat list_4 in

      Printf.printf "%s\n" "";
      Printf.printf "%s\n" "Localization 1/4";
      Printf.printf "%s %i\n" "Number of orbitals" (Mat.dim2 mat_1);
      let loc_mat_1 = f_localise mat_1 in
      Printf.printf "%s\n" "";

      Printf.printf "%s\n" "Localization 2/4";
      Printf.printf "%s %i\n" "Number of orbitals" (Mat.dim2 mat_2);
      let loc_mat_2 = f_localise mat_2 in
      Printf.printf "%s\n" "";

      Printf.printf "%s\n" "Localization 3/4";
      Printf.printf "%s %i\n" "Number of orbitals" (Mat.dim2 mat_3);
      let loc_mat_3 = f_localise mat_3 in
      Printf.printf "%s\n" "";

      Printf.printf "%s\n" "Localization 4/4";
      Printf.printf "%s %i\n" "Number of orbitals" (Mat.dim2 mat_4);
      let loc_mat_4 = f_localise mat_4

      in
      Printf.printf "%s\n" "";
      assemble (assemble loc_mat_1 loc_mat_2) (assemble loc_mat_3 loc_mat_4)
  in

  (* Pattern matching for the separation type of the occupied MOs *)
  let m_assemble_occ, m_assemble_occu =

      let m_assemble_occu =
          Printf.printf "%s\n" "II/III";
          Printf.printf "%s\n" "Localization of valence orbitals";
          Printf.printf "%s %i\n" "Number of orbitals" (Mat.dim2 m_occu);
          match type_separation_occu with
              | "full"         -> let m_assemble_occu = f_loc_assemble_1 m_occu in m_assemble_occu
              | "pi/sigma"     -> let m_assemble_occu = f_loc_assemble_12 m_occu list_pi in m_assemble_occu
              | "pi/sigma/end" -> let m_assemble_occu = f_loc_assemble_123 m_occu list_pi list_x_end in m_assemble_occu
              | "pi/1/2"       -> let m_assemble_occu = f_loc_assemble_123 m_occu list_pi list_s_12 in m_assemble_occu
              | "1/2"          -> let m_assemble_occu = f_loc_assemble_12 m_occu list_s_12 in m_assemble_occu
              | "pi1/pi2/1/2"  -> let m_assemble_occu = f_loc_assemble_1234 m_occu list_pi list_s_12 list_s_12 in m_assemble_occu
              | "By_hands"    -> let m_assemble_occu = f_loc_assemble_main m_occu (miss_elem m_occ []) [] [] [] in m_assemble_occu
              | _              -> invalid_arg "Unknown separation type of valence orbitals"
      in
      Printf.printf "%s\n" "End II/III";
      Printf.printf "%s\n" "";
      (assemble loc_m_core m_assemble_occu, m_assemble_occu)
  in

  (* Pattern matching for the separation type of the virtual MOs *)
  let m_assemble_vir =

      let m_assemble_vir =
          Printf.printf "%s\n" "III/III";
          Printf.printf "%s\n" "Localization of virtual orbitals";
          Printf.printf "%s %i\n" "Number of orbitals" (Mat.dim2 m_vir);
          match type_separation_vir with
              | "full"         -> let m_assemble_vir = f_loc_assemble_1 m_vir in m_assemble_vir
              | "pi/sigma"     -> let m_assemble_vir = f_loc_assemble_12 m_vir list_pi in m_assemble_vir
              | "pi/sigma/end" -> let m_assemble_vir = f_loc_assemble_123 m_vir list_pi list_x_end in m_assemble_vir
              | "pi/1/2"       -> let m_assemble_vir = f_loc_assemble_123 m_vir list_pi list_s_12 in m_assemble_vir
              | "1/2"          -> let m_assemble_vir = f_loc_assemble_12 m_vir list_s_12 in m_assemble_vir
              | "pi1/pi2/1/2"  -> let m_assemble_vir = f_loc_assemble_1234 m_vir list_pi list_s_12 list_s_12 in m_assemble_vir
              | "By_hands"    -> let m_assemble_vir = f_loc_assemble_main m_vir (miss_elem m_vir []) [] [] [] in m_assemble_vir
              | _              -> invalid_arg "Unknown separation type of virtual orbitals"
      in
      Printf.printf "%s\n" "End III/III";
      Printf.printf "%s\n" "";
      m_assemble_vir
  in

  (* Tack occupied and virtual MOs together*)
  let m_assemble_loc = assemble m_assemble_occ m_assemble_vir
  in

  (* VI. Analysis *)

  (* Function to compute the spatial extent of MOs *)
  (* Coefficient matrix n_ao x n_mo -> vector where the i th element is equal to the spatial extent of the i th orbital *)
  (* Matrix -> vector *)
  let distrib mat=
      if Mat.dim1 mat = 0
      then Mat.as_vec mat
      else
      (let multipoles = AOBasis.multipole ao_basis in
      let n_mo = Mat.dim2 mat in
      let phi_x_phi =
          Multipole.matrix_x multipoles
          |> MOBasis.mo_matrix_of_ao_matrix ~mo_coef:mat
      in
      let phi_y_phi =
          Multipole.matrix_y multipoles
          |> MOBasis.mo_matrix_of_ao_matrix ~mo_coef:mat
      in
      let phi_z_phi =
          Multipole.matrix_z multipoles
          |> MOBasis.mo_matrix_of_ao_matrix ~mo_coef:mat
      in
      let phi_x2_phi =
          Multipole.matrix_x2 multipoles
          |> MOBasis.mo_matrix_of_ao_matrix ~mo_coef:mat
      in
      let phi_y2_phi =
          Multipole.matrix_y2 multipoles
          |> MOBasis.mo_matrix_of_ao_matrix ~mo_coef:mat
      in
      let phi_z2_phi =
          Multipole.matrix_z2 multipoles
          |> MOBasis.mo_matrix_of_ao_matrix ~mo_coef:mat
      in
      Vec.init n_mo (fun i ->
      phi_x2_phi.{i,i} -. phi_x_phi.{i,i}**2. +. phi_y2_phi.{i,i} -. phi_y_phi.{i,i}**2. +. phi_z2_phi.{i,i} -. phi_z_phi.{i,i}**2.))
  in

  (* Sorting function *)
  (* tri insertion (<) list -> list in ascending order *)
  (* parameter, list -> list *)
  let rec insere comparaison elem liste = match liste with
          |  [] -> elem::[]
          |  tete::queue -> if comparaison elem tete
                                then elem :: liste
                                else tete :: insere comparaison elem queue
  in
  let rec tri_insertion comp = function
          |  [] -> []
          |  tete::queue -> insere comp tete (tri_insertion comp queue)
  in

  (* Sorting function to sort orbital spatial extent by ascending order *)
  (* coefficient matrix n_ao x n_mo -> matrix n_ao x 2, first column corresponds to the spatial extent in function of the MOs number, the second column corresponds to the spatial extent in ascending order *)
  (* matrix -> matrix *)
  let m_distribution mat =
      let vec_distrib= distrib mat in
      let list_tri_distrib = tri_insertion (<) (Vec.to_list vec_distrib) in
      let vec_tri_distrib = Vec.of_list list_tri_distrib
      in
      Mat.init_cols (Vec.dim vec_distrib) 2 (fun i j ->
              if j=1
                  then vec_distrib.{i}
                  else vec_tri_distrib.{i})

  in

  (* Average/variance/standart deviation/median of spatial extent *)
  (* Function to compute the average *)
  let f_average mat =
          if Mat.dim1 mat = 0
          then 0.
          else  Vec.sum (distrib mat) /. float_of_int(Vec.dim (distrib mat))
  in

  (* Function to compute the variance *)
  let f_variance mat =
          if (Vec.dim (distrib mat)) = 0
          then 0.
          else  Vec.sum(Vec.init (Vec.dim (distrib mat)) (fun i ->
     ((distrib mat).{i}-. (f_average mat))**2.)) /. float_of_int(Vec.dim (distrib mat))
  in

  (* Function to compute the standard deviation *)
  let f_stand_dev mat =
          if (Vec.dim (distrib mat)) = 0
          then 0.
          else  sqrt(abs_float(f_variance mat))
  in

  (* Fonction de calcul de la mediane *)
  let f_median mat =
          if (Vec.dim (distrib mat)) = 0
          then 0.
          else
      let vec_distrib= distrib mat in
      let list_tri_distrib = tri_insertion (<) (Vec.to_list vec_distrib) in
      let vec_tri_distrib = Vec.of_list list_tri_distrib
      in
      if (Vec.dim vec_distrib) mod 2 = 0
          then (vec_tri_distrib.{(Vec.dim vec_distrib)/2} +. vec_tri_distrib.{((Vec.dim vec_distrib)/2)+1}) /. 2.
          else vec_tri_distrib.{int_of_float(float_of_int(Vec.dim vec_distrib)/. 2.) + 1}
  in
  (* Display function *)
  (* Coefficient matrix n_ao x n_mo -> string *)
  (* matrix -> string *)
  let analyse mat =
      let average = f_average mat in
      let variance = f_variance mat in
      let stand_dev = f_stand_dev mat in
      let median = f_median mat
      in
      "Average : "^string_of_float(average)^"; Median : "^string_of_float(median)^"; Variance : "^string_of_float(variance)^"; Standard deviation : "^string_of_float(stand_dev)
  in
  (* Display the numer of MOs *)
  (* Coefficient matrix n_ao x n_mo -> string *)
  (* matrix -> string *)
  let n_orb mat = "Number of molecular orbitals : "^string_of_int(Mat.dim2 mat) in

  (* Display *)
  Printf.printf "%s\n" "";
  Printf.printf "%s\n" (n_orb m_C);
  Printf.printf "%s %i\n" "Number of pi orbitals : " (List.length (list_pi m_C));
  Printf.printf "%s %i\n" "Number of occupied pi orbitals : " (List.length (list_pi m_occ));
  Printf.printf "%s %i\n" "Number of virtual pi orbitals : " (List.length (list_pi m_vir));
  Printf.printf "%s %i\n" "Number of sigma orbitals : " (List.length (miss_elem m_C (list_pi m_C)));
  Printf.printf "%s %i\n" "Number of occupied sigma orbitals : " (List.length (miss_elem m_occ (list_pi m_occ)));
  Printf.printf "%s %i\n" "Number of virtual sigma orbitals : " (List.length (miss_elem m_vir (list_pi m_vir)));

  Printf.printf "%s\n" "";

  Printf.printf "%s %s\n" "All MOs before the localization" (analyse m_C);
  Printf.printf "%s %s\n" "All MOs after the localization" (analyse m_assemble_loc);
  Printf.printf "%s\n" "";
  Util.debug_matrix "Distribution of the spatial extent, before localization / sorting before localization / after localization / sorting after localization" (assemble(m_distribution m_C) (m_distribution m_assemble_loc));

  Printf.printf "%s\n" "";
  Printf.printf "%s %s\n" "Occupied orbitals : " (n_orb m_occ);
  Printf.printf "%s %s\n" "Occupied orbitals before localization"  (analyse m_occ);
  Printf.printf "%s %s\n" "Occupied orbitals after localization" (analyse m_assemble_occ);
  Printf.printf "%s\n" "";

  Printf.printf "%s\n" "";
  Printf.printf "%s %s\n" "Core orbitals : " (n_orb m_core);
  Printf.printf "%s %s\n" "Core orbitalses before localization : " (analyse m_core);
  Printf.printf "%s %s\n" "Core orbitals after localization : " (analyse loc_m_core);
  Printf.printf "%s\n" "";

  Printf.printf "%s\n" "";
  Printf.printf "%s %s\n" "Valence orbitals : " (n_orb m_occu);
  Printf.printf "%s %s\n" "Valence orbitals before localization : " (analyse m_occu);
  Printf.printf "%s %s\n" "Valence rbitals after localization : " (analyse m_assemble_occu);
  Printf.printf "%s\n" "";

  Printf.printf "%s\n" "";
  Printf.printf "%s %s\n" "Virtual orbitals : " (n_orb m_vir);
  Printf.printf "%s %s\n" "Virtual orbitals before localization : " (analyse m_vir);
  Printf.printf "%s %s\n" "Virtual orbitals after localization : " (analyse m_assemble_vir);
  Printf.printf "%s\n" "";

m_assemble_loc