let max_index = 1 lsl 14

type index_pair = { first : int ; second : int }


type storage_t = 
  | Dense  of (float, Bigarray.float32_elt, Bigarray.fortran_layout) Bigarray.Genarray.t
  | Sparse of (int, float) Hashtbl.t

type t =
  {
    size         : int ;
    four_index   : storage_t ;
  }

let key_of_indices ~r1 ~r2 =
  let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in
  let f i k = 
    let p, r = 
      if i <= k then i, k else k, i
    in p+ (r*r-r)/2
  in
  let p = f i k and q = f j l in
  f p q


let get_four_index ~r1 ~r2 t = 
  match t.four_index with
  | Dense  t -> let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in t.{i,j,k,l}
  | Sparse t -> let key = key_of_indices ~r1 ~r2 in
                try Hashtbl.find t key
                with Not_found -> 0.


let set_four_index ~r1 ~r2 ~value t = 
  match t.four_index with
  | Dense  t -> let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in
                t.{i,j,k,l} <- value;
                t.{k,j,i,l} <- value;
                t.{i,l,k,j} <- value;
                t.{k,l,i,j} <- value;
                t.{j,i,l,k} <- value;
                t.{j,k,l,i} <- value;
                t.{l,i,j,k} <- value;
                t.{l,k,j,i} <- value;

  | Sparse t -> let key = key_of_indices ~r1 ~r2 in
                Hashtbl.replace t key value
    

let increment_four_index ~r1 ~r2 ~value x = 
  match x.four_index with
  | Dense  t -> let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in
                t.{i,j,k,l} <- t.{i,j,k,l} +. value;
                t.{k,j,i,l} <- t.{k,j,i,l} +. value;
                t.{i,l,k,j} <- t.{i,l,k,j} +. value;
                t.{k,l,i,j} <- t.{k,l,i,j} +. value;
                t.{j,i,l,k} <- t.{j,i,l,k} +. value;
                t.{j,k,l,i} <- t.{j,k,l,i} +. value;
                t.{l,i,j,k} <- t.{l,i,j,k} +. value;
                t.{l,k,j,i} <- t.{l,k,j,i} +. value

  | Sparse t -> let key = key_of_indices ~r1 ~r2 in
                let old_value =
                  try Hashtbl.find t key
                  with Not_found -> 0.
                in
                Hashtbl.replace t key (old_value +. value)

let get ~r1 ~r2 = 
  get_four_index ~r1 ~r2 

let set ~r1 ~r2 = 
  set_four_index ~r1 ~r2 

let increment ~r1 ~r2 = 
  increment_four_index ~r1 ~r2 

let create ~size sparsity =
  assert (size < max_index);
  let four_index = 
    match sparsity with
    | `Dense ->
      let result = 
        Bigarray.Genarray.create Float32 Bigarray.fortran_layout [| size ; size ; size ; size |]
      in
      Bigarray.Genarray.fill result 0.;
      Dense result

    | `Sparse -> 
      let result = Hashtbl.create (size*size+13) in
      Sparse result
  in
  { size ; four_index }


let size t = t.size

(** TODO : remove epsilons *)


let get_chem t i j k l = get ~r1:{ first=i ; second=j } ~r2:{ first=k ; second=l } t
let get_phys t i j k l = get ~r1:{ first=i ; second=k } ~r2:{ first=j ; second=l } t

let set_chem t i j k l value = set ~r1:{ first=i ; second=j } ~r2:{ first=k ; second=l } ~value t
let set_phys t i j k l value = set ~r1:{ first=i ; second=k } ~r2:{ first=j ; second=l } ~value t 


(** Write all integrals to a file with the <ij|kl> convention *)
let to_file ?(cutoff=Constants.epsilon) ~filename data =
  let oc = open_out filename in
  for l_c=1 to size data do
    for k_c=1 to l_c do
      for j_c=1 to l_c do
        for i_c=1 to k_c do
          let value = get_phys data i_c j_c k_c l_c in
           if (abs_float value > cutoff) then
            Printf.fprintf oc " %5d %5d %5d %5d%20.15f\n" i_c j_c k_c l_c value;
        done;
      done;
    done;
  done;
  close_out oc