open Rresult
open Core_kernel
open Let_syntax.Result

(* FIXME: rewrite based on Dataframe? *)

type t = {
  nrows : int ;
  cols : (string * float option array) list ;
}

let of_file fn =
  let* df = Dataframe.from_file fn in
  let+ cols =
    Dataframe.columns df
    |> List.map ~f:(fun (label, col) ->
        let+ data = match (col : Dataframe.column) with
          | Floats xs -> Ok (Array.map xs ~f:Option.some)
          | Float_opts xs -> Ok xs
          | Ints xs -> Ok (Array.map xs ~f:(fun i -> Some (Float.of_int i)))
          | Int_opts xs -> Ok (Array.map xs ~f:(Option.map ~f:Float.of_int))
          | _ -> R.error_msgf "column %s is not numeric" label
        in
        label, data
      )
    |> Result.all
  in
  let cols = List.filter cols ~f:(fun (label, _) -> not (String.equal label "Sites")) in
  let nrows = match cols with
    | [] -> assert false
    | (_, h) :: _ -> Array.length h
  in
  { nrows ; cols }

let of_file_exn fn = of_file fn |> Rresult.R.failwith_error_msg

let columns x = x.cols

let make = function
  | [] -> Error (`Msg "empty column list")
  | (_, h) :: t as cols ->
    let nrows = Array.length h in
    if List.for_all t ~f:(fun (_, col) -> Array.length col = nrows) then
      Ok { nrows ; cols }
    else R.error_msg "Columns with different lengths"

let write_row oc fields =
  String.concat ~sep:"\t" fields
  |> Out_channel.fprintf oc "%s\n"

let write_columns oc res =
  Range.iter 0 res.nrows ~f:(fun i ->
      Int.to_string i
      :: List.map res.cols ~f:(fun (_, col) ->
          Option.value_map ~default:"NA" ~f:Float.to_string col.(i)
        )
      |> write_row oc
    )

let to_file table ~output =
  Out_channel.with_file output ~f:(fun oc ->
      write_row oc ("Sites" :: List.map table.cols ~f:fst) ;
      write_columns oc table
    )

let get_col cpt colname = 
  List.Assoc.find ~equal:String.equal cpt.cols colname

let get_col_exn cpt colname = 
  get_col cpt colname |> Option.value_exn