(** Data structures for storing the determinant space.

If the space is built as the outer product of all {% $\alpha$ %} and {%
$\beta$ %} determinants, the storage is of type [Spin]. It is sufficient
to have the arrays of {% $\alpha$ %} and {% $\beta$ %} spindeterminants.

Otherwise, the space is of type [Arbitrary]. 

*)

type arbitrary_space =
  { 
    det         : int array array ;
    det_alfa    : Spindeterminant.t array ;
    det_beta    : Spindeterminant.t array ;
    index_start : int array;
  }


type determinant_storage =
  | Arbitrary of arbitrary_space
  | Spin      of (Spindeterminant.t array * Spindeterminant.t array)


type t = 
  {
    n_alfa       : int ;
    n_beta       : int ;
    mo_class     : MOClass.t ;
    mo_basis     : MOBasis.t ;
    determinants : determinant_storage;
  }


module Ss = SpindeterminantSpace

let n_alfa   t = t.n_alfa
let n_beta   t = t.n_beta
let mo_class t = t.mo_class
let mo_basis t = t.mo_basis

let active_mos t =
  mo_class t 
  |> MOClass.active_mos


let inactive_mos t =
  mo_class t 
  |> MOClass.inactive_mos


let virtual_mos t =
  mo_class t 
  |> MOClass.inactive_mos


let mo_class_array t =
  mo_class t
  |> MOClass.mo_class_array 


let size t = 
  match t.determinants with
  | Spin (a,b)  ->  (Array.length a) * (Array.length b)
  | Arbitrary a ->  
      let ndet_a = Array.length a.det_alfa in
      a.index_start.(ndet_a)
 

let determinant_stream t =
  match t.determinants with
  | Arbitrary a -> 
    let det_beta = a.det_beta
    and det_alfa = a.det_alfa
    and det = a.det in
    let n_alfa = Array.length det_alfa in
    let alfa = ref det_alfa.(0)
    and det_i_alfa = ref det.(0) in
    let i_alfa = ref 0
    and k_beta = ref 0
    in
    Stream.from (fun _ ->
        if !i_alfa = n_alfa then None else
          begin
            let i_beta = (!det_i_alfa).(!k_beta) in
            let beta = det_beta.(i_beta) in
            let result = 
              Some (Determinant.of_spindeterminants (!alfa) beta)
            in
            incr k_beta;
            if !k_beta = Array.length !det_i_alfa then
              begin
                k_beta := 0;
                incr i_alfa;
                if !i_alfa < n_alfa then 
                  begin
                    alfa := det_alfa.(!i_alfa);
                    det_i_alfa := det.(!i_alfa)
                  end
              end;
            result
          end
        )

  | Spin (a,b) -> 
    let na = Array.length a
    and nb = Array.length b in
    let i = ref 0
    and j = ref 0 in
    Stream.from (fun k ->
        if !j < nb then
          let result = 
            Determinant.of_spindeterminants a.(!i) b.(!j)
          in
          incr i;
          if !i = na then (i := 0 ; incr j);
          Some result
        else
          None)


let determinants t = t.determinants


let determinants_array t =
  let s = determinant_stream t in
  Array.init (size t) (fun _ -> Stream.next s)


let determinant t i =
  let alfa, beta =
    match t.determinants with
    | Arbitrary a ->
      let i_alfa =
        let index_start = a.index_start in
        let rec loop i_alfa =
          if index_start.(i_alfa) <= i then
            (loop [@tailcall]) (i_alfa+1) 
          else i_alfa
        in loop 0
      in
      let i_beta = i - a.index_start.(i_alfa) in
      let alfa = a.det_alfa.(i_alfa) in
      let beta = a.det_beta.(i_beta) in
      alfa, beta

    | Spin (a,b) ->
      let nb = Array.length b in
      let k = i / nb in
      let j = i - k * nb in
      a.(j), b.(k)

    in
    Determinant.of_spindeterminants alfa beta



let fock_diag det_space det =

  let alfa_list = 
    Determinant.alfa det
    |> Spindeterminant.to_list
  in
  let beta_list = 
    Determinant.beta det
    |> Spindeterminant.to_list
  in
  let mo_basis = mo_basis det_space in
  let mo_num = MOBasis.size mo_basis in
  let one_e_ints  = MOBasis.one_e_ints mo_basis
  and two_e_ints  = MOBasis.two_e_ints mo_basis
  in
  let two_e i j k l = ERI.get_phys two_e_ints i j k l in
  let build_array list1 list2 =      
    let result = Array.make (mo_num+1) 0. in

    (* Occupied *)
    List.iter (fun i ->
        let x = one_e_ints.{i,i} in
        result.(i) <- result.(i) +. x;
        result.(0) <- result.(0) +. x;
        List.iter (fun j ->
            if j <> i then
              begin
                let x = two_e i j i j -. two_e i j j i in
                result.(i) <- result.(i) +. x;
                result.(0) <- result.(0) +. 0.5 *. x;
              end
          ) list1;
        List.iter (fun j ->
              begin
                let x = two_e i j i j in
                result.(i) <- result.(i) +. x;
                result.(0) <- result.(0) +. 0.5 *. x;
              end
          ) list2;
      ) list1;

    (* Virtuals*)
    List.iter (fun i ->
        if result.(i) = 0. then
          begin
            let x = one_e_ints.{i,i} in
            result.(i) <- result.(i) +. x;
            List.iter (fun j ->
                let x = two_e i j i j -. two_e i j j i in
                result.(i) <- result.(i) +. x;
              ) list1;
            List.iter (fun j ->
                  begin
                    let x = two_e i j i j in
                    result.(i) <- result.(i) +. x;
                  end
              ) list2;
          end
     ) (Util.list_range 1 mo_num);
     result
  in
  let alfa, beta = 
    build_array alfa_list beta_list, 
    build_array beta_list alfa_list 
  in
  let e = alfa.(0) +. beta.(0) in
  alfa.(0) <- e;
  beta.(0) <- e;
  alfa, beta



let spin_of_mo_basis mo_basis f =

  let s = MOBasis.simulation mo_basis in
  let e = Simulation.electrons s in
  let n_alfa = Electrons.n_alfa e 
  and n_beta = Electrons.n_beta e in
  let det_a = f n_alfa
  and det_b = f n_beta
  in
  let mo_class = Ss.mo_class det_a in
  let determinants =
    let det_alfa = Ss.spin_determinants det_a 
    and det_beta = Ss.spin_determinants det_b
    in
    let n_det_beta = Array.length det_beta in
    let n_det_alfa = Array.length det_alfa in

    let ndet = n_det_alfa * n_det_beta in
    if Parallel.master then
      Format.printf "Number of determinants : %d %d %d\n%!" 
        n_det_alfa n_det_beta ndet;
    Spin (det_alfa, det_beta)
  in
  { n_alfa ; n_beta ; mo_class ; mo_basis ; determinants }


let arbitrary_of_mo_basis mo_basis f =

  let s = MOBasis.simulation mo_basis in
  let e = Simulation.electrons s in
  let n_alfa = Electrons.n_alfa e 
  and n_beta = Electrons.n_beta e in
  let det_a = f n_alfa
  and det_b = f n_beta
  in
  let mo_class = Ss.mo_class det_a in
  let determinants =
    let det_alfa = Ss.spin_determinants det_a 
    and det_beta = Ss.spin_determinants det_b
    in
    let n_det_beta = Array.length det_beta in
    let n_det_alfa = Array.length det_alfa in

    let det = Array.make n_det_alfa
        (Array.init n_det_beta (fun i -> i))
    in
    let index_start = Array.init (n_det_alfa+1) (fun i -> i*n_det_beta) in
    let ndet = (index_start.(n_det_alfa)) in

    if Parallel.master then
      Format.printf "Number of determinants : %d %d %d\n%!" 
        n_det_alfa n_det_beta ndet;
    Arbitrary {
      det_alfa ; det_beta ; det ; index_start
    }
  in
  { n_alfa ; n_beta ; mo_class ; mo_basis ; determinants }



let cas_of_mo_basis mo_basis ~frozen_core n m =
  let f n_alfa =
    Ss.cas_of_mo_basis mo_basis ~frozen_core n_alfa n m
  in
  spin_of_mo_basis mo_basis f


let fci_of_mo_basis mo_basis ~frozen_core =
  let f n_alfa =
    Ss.fci_of_mo_basis mo_basis ~frozen_core n_alfa
  in
  spin_of_mo_basis mo_basis f


let fci_f12_of_mo_basis mo_basis ~frozen_core mo_num =
  let s = MOBasis.simulation mo_basis in
  let e = Simulation.electrons s in
  let n_alfa = Electrons.n_alfa e 
  and n_beta = Electrons.n_beta e in
  let n_core = 
    if frozen_core then
      (Nuclei.small_core @@ Simulation.nuclei @@ MOBasis.simulation mo_basis) / 2
    else 0
  in
  let n, m = 
      (n_alfa + n_beta - n_core),
      (mo_num - n_core)
  in
  let f n_alfa =
    Ss.cas_of_mo_basis mo_basis ~frozen_core n_alfa n m
  in
  let r = 
    spin_of_mo_basis mo_basis f
  in
  { r with mo_class =
    MOClass.to_list r.mo_class
    |> List.rev_map (fun i ->
      match i with
      | MOClass.Virtual i when i > mo_num -> MOClass.Auxiliary i 
      | i -> i)
    |> List.rev
    |> MOClass.of_list }
    

let cas_f12_of_mo_basis mo_basis ~frozen_core n m mo_num =
  let f n_alfa =
    Ss.cas_of_mo_basis mo_basis ~frozen_core n_alfa n m
  in
  let r = 
    spin_of_mo_basis mo_basis f
  in
  { r with mo_class =
    MOClass.to_list r.mo_class
    |> List.rev_map (fun i ->
      match i with
      | MOClass.Virtual i when i > mo_num -> MOClass.Auxiliary i 
      | i -> i)
    |> List.rev
    |> MOClass.of_list
    }



let pp ppf t =
  Format.fprintf ppf "@[<v 2>[ ";
  let i = ref 0 in
  determinant_stream t
  |> Stream.iter (fun d -> Format.fprintf ppf "@[<v>@[%8d@]@;@[%a@]@]@;" !i
                    (Determinant.pp (MOBasis.size (mo_basis t))) d; incr i) ;
  Format.fprintf ppf "]@]"