(** 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 = Spindeterminant_space

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 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 - 1) + Array.length a.det.(ndet_a - 1)
 

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 (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 fci_of_mo_basis ?(frozen_core=true) mo_basis =
  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 =
    Ss.fci_of_mo_basis ~frozen_core mo_basis n_alfa
  and det_b =
    Ss.fci_of_mo_basis ~frozen_core mo_basis 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
    Spin (det_alfa, det_beta)
    (*
    let n_det_beta = Array.length det_beta in
    Arbitrary {
      det_alfa ; det_beta ;
      det = Array.make (Array.length det_alfa) (Array.init (Array.length det_beta) (fun i -> i));
      index_start = Array.mapi (fun i _ -> i*n_det_beta) det_alfa;
    }
    *)
  in
  { n_alfa ; n_beta ; mo_class ; mo_basis ; determinants }


let pp_det_space 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_det (MOBasis.size (mo_basis t))) d; incr i) ;
  Format.fprintf ppf "]@]"