type determinant_storage =
| Arbitrary of Determinant.t array
| 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
  | Arbitrary a -> Array.length a
  | Spin (a,b)  ->  (Array.length a) * (Array.length b)
 

let determinant_stream t =
  let imax = size t in
  match t.determinants with
  | Arbitrary a -> 
      Stream.from (fun i ->
          if i < imax then Some a.(i) else None)
  | 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 =
  match t.determinants with
  | Arbitrary a -> a
  | Spin (a,b)  -> 
      let s = determinant_stream t in
      Array.init (Array.length a * Array.length b) (fun _ ->
        Stream.next s)
      (*
    Array.to_list b
    |> List.map (fun det_b ->
        Array.map (fun det_a ->
            Determinant.of_spindeterminants det_a det_b
          ) a
        )
    |> Array.concat
    *)


let determinant t i =
  match t.determinants with
  | Arbitrary a -> a.(i)
  | Spin (a,b) ->
    let nb = Array.length b in
    let k = i / nb in
    let j = i - k * nb in
    Determinant.of_spindeterminants a.(j) b.(k)


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 a = Ss.spin_determinants det_a
    and b = Ss.spin_determinants det_b
    in Spin (a,b)
  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 "]@]"