open Util
open Lacaml.D
open Constants

let max_index = 1 lsl 14

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

type array2 = (float, Bigarray.float64_elt, Bigarray.fortran_layout) Bigarray.Array2.t

type storage_t = 
  | Dense  of array2
  | Svd of Vec.t * array2 * array2
  | Sparse of (int, float) Hashtbl.t

type t =
  {
    size             : int ;
    two_index        : array2;
    two_index_anti   : array2;
    three_index      : array2;
    three_index_anti : array2;
    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 check_bounds r1 r2 t =
  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 dense_index i j size =
  (j-1)*size + i


let sym_index i j =
  if i < j then
    (j*(j-1))/2 + i
  else
    (i*(i-1))/2 + j


let unsafe_get_four_index ~r1 ~r2 t = 

  let open Bigarray.Array2 in
  let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in 

  if i=k then
    if j=l then
      unsafe_get t.two_index i j
    else
      unsafe_get t.three_index (dense_index j l t.size) i
  else if j=l then
      unsafe_get t.three_index (dense_index i k t.size) j

  else if i=l then
    if k=j then 
      unsafe_get t.two_index_anti i j
    else
      unsafe_get t.three_index_anti (dense_index j k t.size) i 

  else if j=k then
      unsafe_get t.three_index_anti (dense_index i l t.size) j

  else if i=j then
    if k=l then 
      unsafe_get t.two_index_anti i k
    else
      unsafe_get t.three_index_anti (dense_index k l t.size) i

  else if k=l then
      (* <ij|kk> *)
      unsafe_get t.three_index_anti (dense_index i j t.size) k

  else
    match t.four_index with
    | Dense  a -> unsafe_get a (dense_index i k t.size) (sym_index j l)
    | Svd _ -> assert false
    | Sparse a -> let key = key_of_indices ~r1 ~r2 in
                  try Hashtbl.find a key
                  with Not_found -> 0.


let get_four_index ~r1 ~r2 t = 
  check_bounds r1 r2 t;
  unsafe_get_four_index ~r1 ~r2 t  


let unsafe_set_four_index ~r1 ~r2 ~value t = 


  let open Bigarray.Array2 in
  let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in 

  let () = 

    if i=k then
      begin
        if j=l then
          begin
            unsafe_set t.two_index i j value;
            unsafe_set t.two_index j i value;
            unsafe_set t.three_index (dense_index i i t.size) j value;
          end;
        unsafe_set t.three_index (dense_index j l t.size) i value;
        unsafe_set t.three_index (dense_index l j t.size) i value;
      end
    else if j=l then
        begin
          unsafe_set t.three_index (dense_index i k t.size) j value;
          unsafe_set t.three_index (dense_index k i t.size) j value;
        end

    else if i=l then
      begin
        if j=k then
          begin
            unsafe_set t.two_index_anti i j value;
            unsafe_set t.two_index_anti j i value;
            unsafe_set t.three_index_anti (dense_index i i t.size) j value;
          end;
        unsafe_set t.three_index_anti (dense_index j k t.size) i value;
        unsafe_set t.three_index_anti (dense_index k j t.size) i value;
      end
    else if j=k then
        begin
          unsafe_set t.three_index_anti (dense_index i l t.size) j value;
          unsafe_set t.three_index_anti (dense_index l i t.size) j value;
        end

    else if i=j then 
      begin
        if k=l then
          begin
            unsafe_set t.two_index_anti i k value;
            unsafe_set t.two_index_anti k i value;
            unsafe_set t.three_index_anti (dense_index i i t.size) k value;
          end;
        unsafe_set t.three_index_anti (dense_index k l t.size) i value;
        unsafe_set t.three_index_anti (dense_index l k t.size) i value;
      end

    else if k=l then
        (* <ij|kk> *)
        begin
          unsafe_set t.three_index_anti (dense_index i j t.size) k value;
          unsafe_set t.three_index_anti (dense_index j i t.size) k value;
        end
  in

  match t.four_index with
  | Dense  a -> let ik = (dense_index i k t.size)
                and jl = (dense_index j l t.size)
                and ki = (dense_index k i t.size)
                and lj = (dense_index l j t.size)
                and ik_s = (sym_index i k)
                and jl_s = (sym_index j l)
                in
                begin
                  unsafe_set a ik jl_s value;
                  unsafe_set a ki jl_s value;
                  unsafe_set a jl ik_s value;
                  unsafe_set a lj ik_s value;
                end
  | Sparse a -> let key = key_of_indices ~r1 ~r2 in
                Hashtbl.replace a key value
  | Svd _ -> assert false
    

let set_four_index ~r1 ~r2 ~value t = 
  check_bounds r1 r2 t;
  unsafe_set_four_index ~r1 ~r2 ~value t


let unsafe_increment_four_index ~r1 ~r2 ~value t = 
  let updated_value =
    value +. unsafe_get_four_index ~r1 ~r2 t
  in
  unsafe_set_four_index ~r1 ~r2 ~value:updated_value t


let increment_four_index ~r1 ~r2 ~value t = 
  check_bounds r1 r2 t;
  unsafe_increment_four_index ~r1 ~r2 ~value t


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 ?(temp_dir="/dev/shm") sparsity =
  assert (size < max_index);
  let two_index = 
      SharedMemory.create ~temp_dir Float64 [| size ; size |]
      |> Bigarray.array2_of_genarray
  in
  let two_index_anti = 
      SharedMemory.create ~temp_dir Float64 [| size ; size |]
      |> Bigarray.array2_of_genarray
  in
  let three_index = 
      SharedMemory.create ~temp_dir Float64 [| size * size ; size |]
      |> Bigarray.array2_of_genarray
  in
  let three_index_anti = 
      SharedMemory.create ~temp_dir Float64 [| size * size ; size |]
      |> Bigarray.array2_of_genarray
  in
  let four_index = 
    match sparsity with
    | `Dense  -> let result =
                     SharedMemory.create ~temp_dir Float64 [| size*size ; (size*(size+1))/2 |]
                     |> Bigarray.array2_of_genarray
                 in
                 Dense result
    | `Sparse -> let result = Hashtbl.create (size*size+13) in
      Sparse result
  in
  { size ; two_index ; two_index_anti ; three_index ; three_index_anti ; 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 

let increment_chem t i j k l value = increment ~r1:{ first=i ; second=j } ~r2:{ first=k ; second=l } ~value t
let increment_phys t i j k l value = increment ~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  =
  Vec.init d.size (fun i -> get_phys d i j k l)
    

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


let get_phys_all_ij d ~k ~l  =
  Mat.init_cols d.size d.size (fun i j -> get_phys d i j k l)
    

let get_chem_all_ij d ~k ~l  =

  if k = l then

      let result = 
        Mat.col d.three_index k
        |> Bigarray.genarray_of_array1 
      in
      Bigarray.reshape_2 result d.size d.size

  else

      match d.four_index with
      | Dense  a -> 
          let kl = sym_index k l in
          let result = 
            Mat.col a kl
            |> Bigarray.genarray_of_array1
          in
          Bigarray.reshape_2 result d.size d.size
      | Sparse a ->
          Mat.init_cols d.size d.size (fun i j -> get_chem d i j k l)
      | Svd _ -> assert false
    

    
let to_stream d =

  let i = ref 0
  and j = ref 1
  and k = ref 1
  and l = ref 1
  in
  let f_dense _ =
    incr i;
    if !i > !k then begin
      i := 1;
      incr j;
      if !j > !l then begin
        j := 1;
        incr k;
        if !k > !l then begin
          k := 1;
          incr l;
        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.integrals_cutoff) ~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 [@tailcall]) ()
    | 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 [@tailcall]) (d :: accu) 
  in
  append [] 



let broadcast t = 
  let size =
    Parallel.InterNode.broadcast (lazy t.size)
  in
  let bufsize = size * size * size in
  let stream = to_stream t in
  let rec iterate () =
    let buffer =
      Parallel.InterNode.broadcast (lazy (
        if Stream.peek stream = None then None else
        Some (Array.init bufsize (fun _ -> 
           try Some (Stream.next stream)
           with _ -> None
         ))
      ) )
    in
    match buffer with
    | None -> ()
    | Some buffer ->
      begin
        if not Parallel.InterNode.master then
          Array.iter (fun x ->
            match x with
            | Some {i_r1 ; j_r2 ; k_r1 ; l_r2 ; value} ->
                  set_phys t i_r1 j_r2 k_r1 l_r2 value
            | None -> () ) buffer;
        (iterate [@tailcall]) ()
      end
  in iterate ();
  t

     

let four_index_transform_dense_sparse ds coef source =

  let mo_num = Mat.dim2 coef in
  let destination = create ~size:mo_num ds in

  let ao_num = Mat.dim1 coef in
  let mo_num_2  = mo_num * mo_num in
  let ao_num_2  = ao_num * ao_num in
  let ao_mo_num = ao_num * mo_num in

  let range_mo = list_range 1 mo_num in
  let range_ao = list_range 1 ao_num in

  let u = Mat.create mo_num_2 mo_num
  and o = Mat.create ao_num ao_num_2 
  and p = Mat.create ao_num_2 mo_num
  and q = Mat.create ao_mo_num mo_num
  in

  if Parallel.master then Printf.eprintf "4-idx transformation \n%!";

  let task delta =
      Mat.fill u 0.;

      List.iter (fun l ->
          if abs_float coef.{l,delta} > epsilon then
            begin
              let jk = ref 1 in
              List.iter (fun k ->
                  get_chem_all_ij source ~k ~l 
                  |> lacpy ~b:o ~bc:!jk 
                  |> ignore;
                  jk := !jk + ao_num;
                ) range_ao;
              (* o_i_jk *)

              let p =
                gemm ~transa:`T ~c:p o coef
                (* p_jk_alpha  = \sum_i o_i_jk c_i_alpha *)
              in
              let p' = 
                Bigarray.reshape_2 (Bigarray.genarray_of_array2 p) ao_num ao_mo_num
                (* p_j_kalpha *)
              in

              let q =
                gemm ~transa:`T ~c:q p' coef
                (* q_kalpha_beta  = \sum_j p_j_kalpha c_j_beta *)
              in
              let q' =
                Bigarray.reshape_2 (Bigarray.genarray_of_array2 q) ao_num mo_num_2
                (* q_k_alphabeta  = \sum_j p_j_kalpha c_j_beta *)
              in

              ignore @@
              gemm ~transa:`T ~beta:1. ~alpha:coef.{l,delta} ~c:u q' coef ;
              (* u_alphabeta_gamma = \sum_k q_k_alphabeta c_k_gamma *)
            end
        ) range_ao;
      let u =
        Bigarray.reshape 
          (Bigarray.genarray_of_array2 u)
          [| mo_num ; mo_num ; mo_num |]
        |> Bigarray.array3_of_genarray 
      in
      let result = ref [] in
      List.iter (fun gamma ->
          List.iter (fun beta ->
              List.iter (fun alpha ->
                  let x = u.{alpha,beta,gamma} in
                  if abs_float x > Constants.integrals_cutoff then
                    result := (alpha, beta, gamma, delta, x) :: !result;
                ) (list_range 1 beta) 
            ) range_mo
        ) (list_range 1 delta);
      Array.of_list !result
  in


  let n = ref 0 in
  Stream.of_list range_mo
  |> Farm.run ~f:task ~ordered:true ~comm:Parallel.Node.comm
  |> Stream.iter (fun l ->
     if Parallel.master then (incr n ; Printf.eprintf "\r%d / %d%!" !n mo_num);
     Array.iter (fun (alpha, beta, gamma, delta, x) ->
      set_chem destination alpha beta gamma delta x) l);

  if Parallel.master then Printf.eprintf "\n%!";
  broadcast destination


let svd_of_dense t =
  let four_index =
    let f =
      match t.four_index with
      | Dense a -> a
      | _ -> invalid_arg "svd_of_dense expects a Dense storage"
    in
    let size = t.size in
    let b = Mat.create (size*size) (size*size) in
    for k = 1 to size do
      for l = 1 to size do
        let ac = sym_index k l in
        lacpy f ~n:1 ~ac ~b ~bc:(dense_index k l size)
        |> ignore;
        lacpy f ~n:1 ~ac ~b ~bc:(dense_index l k size)
        |> ignore
      done
    done;
    let d, u, v = gesdd b in
    let rank =
      let w = copy d in
      scal (1. /. d.{1}) w;
      let rec aux i = 
        if i = Vec.dim w then i else
        if w.{i} < epsilon_float then i else
          aux (i+1)
      in aux 1
    in
    let d = copy d ~n:rank in
    let u = lacpy ~n:rank u in
    let v = lacpy ~m:rank v in
    Svd (d,u,v)
  in
  { t with four_index }


let dense_of_svd t = 
  let four_index =
    let d, u, v =
      match t.four_index with
      | Svd (d, u, v) -> d, u, v
      | _ -> invalid_arg "dense_of_svd expects a Svd storage"
    in
    let f = 
      gemm u @@ gemm (Mat.of_diag d) v 
    in
    let size = t.size in
    let b = Mat.create (size*size) ((size*(size+1))/2) in
    for k = 1 to size do
      for l = 1 to k do
        let bc = sym_index k l in
        lacpy f ~ac:(dense_index l k size) ~n:1 ~bc ~b
        |> ignore
      done
    done;
    Dense b
  in
  { t with four_index }


let four_index_transform_svd coef source =
  let t0 = Unix.gettimeofday () in
  let svd = svd_of_dense source in
  Printf.printf "Computed SVD in %f seconds\n%!" (Unix.gettimeofday () -. t0);

  let u0, d0, v0 = 
    match svd.four_index with
    | Svd (d, u, v) -> u, d, (Mat.transpose_copy v)
    | _ -> assert false
  in

  let t0 = Unix.gettimeofday () in

  let ao_num = Mat.dim1 coef in
  let mo_num = Mat.dim2 coef in
  let svd_num = Vec.dim d0 in

  let destination = create ~size:mo_num `Dense in

  let mo_num_2  = mo_num * mo_num in

  if Parallel.master then Printf.eprintf "4-idx transformation \n%!";

  let u_vecs = Mat.to_col_vecs u0
  and v_vecs = Mat.to_col_vecs v0
  in


  Printf.printf "%d %d %d %d\n" (Mat.dim1 u0) (Mat.dim2 u0)  (Mat.dim1 v0) (Mat.dim2 v0);
  let task a =
    let o =
      Bigarray.reshape_2 (Bigarray.genarray_of_array1 u_vecs.(a-1)) ao_num ao_num
    in

    let u =
      let p = gemm ~transa:`T coef @@ gemm o coef in
      Bigarray.reshape_1 (Bigarray.genarray_of_array2 p) mo_num_2
    in

    let o =
      Bigarray.reshape_2 (Bigarray.genarray_of_array1 v_vecs.(a-1)) ao_num ao_num
    in

    let v =
      let p = gemm ~transa:`T coef @@ gemm o coef in
      Bigarray.reshape_1 (Bigarray.genarray_of_array2 p) mo_num_2
    in
    (u, v)
  in


  (*
  let four_index = Svd (d,u'',v'') in
  let t = make_from_four_index

  { source with
    size = mo_num;
    two_index;
    two_index_anti;
    three_index;
    three_index_anti;
    four_index
  }
  *)

  let n = ref 0 in
  let u_list = ref []
  and v_list = ref []
  in
  Util.list_range 1 svd_num
  |> Stream.of_list 
  |> Farm.run ~f:task ~ordered:true ~comm:Parallel.Node.comm
  |> Stream.iter (fun (u,v) ->
     if Parallel.master then (incr n ; Printf.eprintf "\r%d / %d%!" !n mo_num);
     u_list := u :: !u_list;
     v_list := v :: !v_list);

  let u = Mat.of_col_vecs_list (List.rev !u_list)
  and v = Mat.of_col_vecs_list (List.rev !v_list)
  in
  Printf.printf "Computed 4-idx in %f seconds\n%!" (Unix.gettimeofday () -. t0);

  let result = 
    let p = gemm u @@ gemm ~transb:`T (Mat.of_diag d0) v in
    Bigarray.reshape (Bigarray.genarray_of_array2 p) [| mo_num ; mo_num ; mo_num ; mo_num |]
  in
  for a=1 to mo_num do
    for b=a to mo_num do
      for c=1 to mo_num do
        for d=c to mo_num do
          let x = result.{a,b,c,d} in
          if abs_float x > Constants.integrals_cutoff then
            set_chem destination a b c d x;
        done
      done
    done
  done;

  Printf.printf "Computed 4-idx in %f seconds\n%!" (Unix.gettimeofday () -. t0);

  if Parallel.master then Printf.eprintf "\n%!";
  broadcast destination









let four_index_transform coef source =
  match source.four_index with
  | Sparse   _ -> four_index_transform_dense_sparse `Sparse coef source
  | Svd      _ -> assert false (* four_index_transform_svd coef source *)
  | Dense    _ -> four_index_transform_dense_sparse `Dense  coef source