open Core_kernel

type t = {
  oracle : bool array option;
  scores_per_meth : (string * float option array) list;
}

(* invariant: all columns have the same length *)

let of_file fn =
  match
    In_channel.read_lines fn
    |> List.map ~f:(String.split ~on:'\t')
    |> List.transpose
  with
  | None -> failwith "Variable number of fields"
  | Some cols ->
      let cols =
        List.map cols ~f:(function [] -> assert false | h :: t -> (h, t))
      in
      let oracle =
        List.Assoc.find cols "Oracle" ~equal:String.equal
        |> Option.map
             ~f:
               (List.map ~f:(function
                 | "0" -> false
                 | "1" -> true
                 | _ -> failwith "Unexpected element in oracle column"))
        |> Option.map ~f:Array.of_list
      in
      let scores_per_meth =
        List.filter_map cols ~f:(fun (label, values) ->
            if String.(label <> "Oracle" && label <> "Sites") then
              let values =
                List.map values ~f:(function
                  | "NA" -> None
                  | x -> Some (Float.of_string x))
                |> Array.of_list
              in
              Some (label, values)
            else None)
      in
      { oracle; scores_per_meth }

let nrows results =
  match results.oracle with
  | Some oracle -> Array.length oracle
  | None -> (
      match results.scores_per_meth with
      | (_, col) :: _ -> Array.length col
      | [] -> 0 )

type cell = Nullable_float of float option | Bool of bool | Int of int

let string_of_cell = function
  | Bool b -> if b then "1" else "0"
  | Nullable_float None -> "NA"
  | Nullable_float (Some f) -> Float.to_string f
  | Int i -> Int.to_string i

let line_of_cells xs = List.map xs ~f:string_of_cell |> String.concat ~sep:"\t"

let maybe_add o xs = match o with None -> xs | Some h -> h :: xs

let headers results =
  "Site"
  :: maybe_add
       (Option.map results.oracle ~f:(Fn.const "Oracle"))
       (List.map results.scores_per_meth ~f:fst)

let get_oracle res i = Option.map res.oracle ~f:(fun o -> Bool o.(i))

let row res i =
  Int (i + 1)
  :: maybe_add (get_oracle res i)
       (List.map res.scores_per_meth ~f:(fun (_, s) -> Nullable_float s.(i)))

let write_columns oc res =
  let nrows = nrows res in
  Utils.range_iter 0 nrows ~f:(fun i ->
      row res i |> line_of_cells |> Out_channel.fprintf oc "%s\n")

let to_file results ~output =
  let header_str = headers results |> String.concat ~sep:"\t" in
  Out_channel.with_file output ~f:(fun oc ->
      Out_channel.fprintf oc "%s\n" header_str;
      write_columns oc results)