open Util

let max_index = 1 lsl 14

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


type storage_t = 
  | Dense  of (float, Bigarray.float64_elt, Bigarray.fortran_layout) Bigarray.Array2.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-1))/2
  in
  let p = f i k and q = f j l in
  f p q


let dense_index i j size =
  (j-1)*size + i


let get_four_index ~r1 ~r2 t = 
  match t.four_index with
  | Dense  a -> (let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in 
                 let size = t.size in
                 assert ( (i lor j lor k lor l) > 0 );
                 assert ( i <= size && j <= size && k <= size && l <= size );
                 Bigarray.Array2.unsafe_get a (dense_index i j size) (dense_index k l size)
                )
  | Sparse a -> let key = key_of_indices ~r1 ~r2 in
                try Hashtbl.find a key
                with Not_found -> 0.


let set_four_index ~r1 ~r2 ~value t = 
  match t.four_index with
  | Dense  a -> (let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in
                 let size = t.size in
                 assert ( (i lor j lor k lor l) > 0 );
                 assert ( i <= size && j <= size && k <= size && l <= size);
                 let ij = (dense_index i j size)
                 and kl = (dense_index k l size)
                 and il = (dense_index i l size)
                 and kj = (dense_index k j size)
                 and ji = (dense_index j i size)
                 and lk = (dense_index l k size)
                 and li = (dense_index l i size)
                 and jk = (dense_index j k size)
                 in
                 let open Bigarray.Array2 in
                   unsafe_set a ij kl value;
                   unsafe_set a kj il value;
                   unsafe_set a il kj value;
                   unsafe_set a kl ij value;
                   unsafe_set a ji lk value;
                   unsafe_set a li jk value;
                   unsafe_set a jk li value;
                   unsafe_set a lk ji value
                )

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

let increment_four_index ~r1 ~r2 ~value t = 
  match t.four_index with
  | Dense  a -> (let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in
                 let size = t.size in
                 assert ( (i lor j lor k lor l) > 0 );
                 assert ( i <= size && j <= size && k <= size && l <= size);
                 let ij = (dense_index i j size)
                 and kl = (dense_index k l size)
                 and il = (dense_index i l size)
                 and kj = (dense_index k j size)
                 and ji = (dense_index j i size)
                 and lk = (dense_index l k size)
                 and li = (dense_index l i size)
                 and jk = (dense_index j k size)
                 in
                 let open Bigarray.Array2 in
                   unsafe_set a ij kl (value +. unsafe_get a ij kl) ;
                   unsafe_set a kj il (value +. unsafe_get a kj il) ;
                   unsafe_set a il kj (value +. unsafe_get a il kj) ;
                   unsafe_set a kl ij (value +. unsafe_get a kl ij) ;
                   unsafe_set a ji lk (value +. unsafe_get a ji lk) ;
                   unsafe_set a li jk (value +. unsafe_get a li jk) ;
                   unsafe_set a jk li (value +. unsafe_get a jk li) ;
                   unsafe_set a lk ji (value +. unsafe_get a lk ji) 
                )

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

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

let set ~r1 ~r2 ~value = 
  match classify_float value with
  | FP_normal -> set_four_index ~r1 ~r2 ~value
  | FP_zero    
  | FP_subnormal -> fun _ -> ()
  | FP_infinite
  | FP_nan ->
    let msg =
      Printf.sprintf "FourIdxStorage.ml : set : r1 = (%d,%d) ; r2 = (%d,%d)"
        r1.first r1.second r2.first r2.second
    in
    raise (Invalid_argument msg)

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.Array2.create Float64 Bigarray.fortran_layout (size*size) (size*size)
      in
      Bigarray.Array2.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 



type element = (** Element for the stream *)
{
  i_r1: int ;
  j_r2: int ;
  k_r1: int ;
  l_r2: int ;
  value: float
}


let get_phys_all_i d ~j ~k ~l  =
  Array.init d.size (fun i -> get_phys d (i+1) j k l)
    

let get_chem_all_i d ~j ~k ~l  =
  Array.init d.size (fun i -> get_chem d (i+1) j k l)


let get_phys_all_ji d ~k ~l  =
  Array.init d.size (fun j -> get_phys_all_i d ~j:(j+1) ~k ~l)
    

let get_chem_all_ji d ~k ~l  =
  Array.init d.size (fun j -> get_chem_all_i d ~j:(j+1) ~k ~l)
    

    
let to_stream d =

  let i = ref 0
  and j = ref 1
  and k = ref 1
  and l = ref 1
  in
  let rec f_dense _ =
    i := !i+1;
    if !i > !k then begin
      i := 1;
      j := !j + 1;
      if !j > !l then begin
        j := 1;
        k := !k + 1;
        if !k > !l then begin
          k := 1;
          l := !l + 1;
        end;
      end;
     end;
     if !l <= d.size then 
       Some { i_r1 = !i ; j_r2 = !j ;
             k_r1 = !k ; l_r2 = !l ;
             value = get_phys d !i !j !k !l
           }
     else
       None
  in
  Stream.from f_dense
    

(** 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
  to_stream data 
  |> Stream.iter (fun {i_r1 ; j_r2 ; k_r1 ; l_r2 ; value} ->
           if (abs_float value > cutoff) then
            Printf.fprintf oc " %5d %5d %5d %5d%20.15f\n" i_r1 j_r2 k_r1 l_r2 value);
  close_out oc



let of_file ~size ~sparsity filename =
  let result = create ~size sparsity in
  let ic = Scanf.Scanning.open_in filename in
  let rec read_line () =
    let result =
      try
        Some (Scanf.bscanf ic " %d %d %d %d %f" (fun i j k l v ->
          set_phys result i j k l v))
      with End_of_file -> None
    in
    match result with
    | Some () -> read_line ()
    | None    -> ()
  in
  read_line ();
  Scanf.Scanning.close_in ic;
  result


let to_list data =
  let s =
    to_stream data
  in
  let rec append accu = 
    let d = 
      try Some (Stream.next s) with
      | Stream.Failure -> None
    in
    match d with
    | None -> List.rev accu
    | Some d -> append (d :: accu) 
  in
  append []