Source file csv.ml

open! Core
open Poly

module Helper = struct

  let prepend_name acc field = (Field.name field) :: acc

  let add x acc _field = x + acc

  (*let separator = ","

  let newline = "\n"*)

  let write ~is_first:_ ~is_last:_ ~writer ~to_string _ _ field =
    writer (to_string field);

end

module type Stringable = sig
  type t
  val to_string : t -> string
  val of_string : string -> t
end

module Spec = struct
  exception Illegal_csv_file
  type t = Leaf of string | Tree of (string * t list)

  let rec depth (t : t) =
    match t with
    | Leaf _ -> 1
    | Tree (_, children) ->
        List.fold_left ~f:(fun acc child -> max acc ((depth child) + 1)) ~init:0 children

  let depth t =
    (List.fold_left ~f:(fun acc t -> max acc (depth t)) ~init:0 t)

  let rec matches' (parent_rows, current_header, children_rows) (t : t) =
    match t, current_header with
    | (Tree _ | Leaf _), [] -> raise Illegal_csv_file
    | Leaf spec_title, (real_title :: rest_of_header) ->
        if real_title <> spec_title then raise Illegal_csv_file
        else
          List.map ~f:List.tl_exn parent_rows,
          rest_of_header,
          List.map ~f:List.tl_exn children_rows
    | (Tree (spec_title, children)), (real_title :: _) ->
        if spec_title = real_title then
          match children_rows with
          | [] -> raise Illegal_csv_file
          | h :: children_rows ->
              let init = (current_header :: parent_rows), h, children_rows in
              let parent_rows, current_header, children_rows =
                List.fold_left ~f:matches' ~init children in
              match parent_rows with
              | h :: parent_rows -> parent_rows, h, (current_header :: children_rows)
              | [] -> raise Illegal_csv_file
        else raise Illegal_csv_file

  let matches header_of_csv t =
    match header_of_csv with
    | [] -> raise Illegal_csv_file
    | current_header :: children_rows ->
        let init = ([], current_header, children_rows) in
        try ([], [], []) = List.fold_left ~f:matches' ~init t with
        | _exn -> false

  let check ~csv ~header:t ~f =
    let depth = depth t in
    let rec aux depth acc csv =
      if depth = 0 then
        if matches (List.rev acc) t then f csv
        else raise Illegal_csv_file
      else
        match csv with
        | h :: csv -> aux (depth - 1) (h :: acc) csv
        | [] -> raise Illegal_csv_file
    in
    aux depth [] csv


  let prepend result depth headers =
    let len = Array.length result in
    for i = len - 1 downto 0 do
      let to_prepend =
        if i > depth || Stack.is_empty headers then
            ""
        else Stack.pop_exn headers
      in
      result.(i) <- to_prepend :: result.(i);
    done

  let str_to_human_readable str =
    let str =
      String.map str ~f:(function
        | '_' -> ' '
        | c -> c)
    in
    String.capitalize str

  let rec header_of_t ~depth ~result ~parents t =
    match t with
    | Leaf str ->
        let str = str_to_human_readable str in
        Stack.push parents str;
        prepend result depth parents
    | Tree (str, children) ->
        let str = str_to_human_readable str in
        let depth = depth + 1 in
        Stack.push parents str;
        header_of_list ~depth ~result ~parents children

  and header_of_list ~depth ~result ~parents lst =
    match lst with
    | [] -> assert false
    | h :: t ->
        List.iter
          ~f:(fun x -> header_of_t ~depth ~result ~parents:(Stack.create ()) x)
        (List.rev t);
        header_of_t ~depth ~result ~parents h

  let header lst =
    let depth = depth lst in
    let result = Array.create ~len:depth [] in
    header_of_list ~depth:0 ~result ~parents:(Stack.create ()) lst;
    Array.to_list result

end

module type Csvable_simple = sig
  type t
  val is_csv_atom : bool
  val rev_csv_header' : string list -> _ -> _ -> string list
  val rev_csv_header_spec' : Spec.t list -> _ -> _ -> Spec.t list
  val t_of_row' : _ -> string list -> (unit -> t) * (string list)
  val write_row_of_t' :
    is_first:bool -> is_last:bool -> writer:(string -> unit) -> _ -> _ -> t -> unit
end

module type Csvable = sig
  include Csvable_simple
  val csv_header : string list
  val csv_header_spec : Spec.t list
  val t_of_row : string list -> t
  val row_of_t : t -> string list
  val csv_load : ?separator:char -> string -> t list
  val csv_load_in : ?separator:char -> In_channel.t -> t list
  val csv_save_fn : ?separator:char -> (string -> unit) -> t list -> unit
  val csv_save_out : ?separator:char -> Out_channel.t -> t list -> unit
  val csv_save : ?separator:char -> string -> t list -> unit
end

exception Excess_of_elements_in_row of string list

module Record (S : Csvable_simple) : Csvable with type t = S.t = struct
  include S

  let rev_csv_header' reverse_headers _ _ =
    if S.is_csv_atom then reverse_headers
    else S.rev_csv_header' reverse_headers () ()

  let rev_csv_header_spec' specs _ _ =
    if S.is_csv_atom then specs
    else S.rev_csv_header_spec' specs () ()

  let csv_header = List.rev (rev_csv_header' [] () ())

  let rec aux_csv_header_spec (t : Spec.t) =
    match t with
    | Spec.Leaf _ -> t
    | Spec.Tree (name, children) ->
        Spec.Tree (name, aux_csv_header_spec' children)

  and aux_csv_header_spec' lst =
    List.rev_map ~f:aux_csv_header_spec lst

  let csv_header_spec = aux_csv_header_spec' (rev_csv_header_spec' [] () ())

  let of_list_without_tail aux strings =
    match aux () strings with
    | f, [] -> f ()
    | _, lst -> raise (Excess_of_elements_in_row lst)

  let t_of_row strings = of_list_without_tail S.t_of_row' strings

  (* This is not being used anymore really *)
  (*let write_row_of_t ~writer csvable =
    S.write_row_of_t' ~is_first:true ~is_last:true ~writer () () csvable*)

  let row_of_t csvable =
    let list = ref [] in
    let writer str = list := str :: !list in
    S.write_row_of_t' ~is_first:true ~is_last:true ~writer () () csvable;
    List.rev !list

  let csv_save_fn ?separator writer csvable =
    Csvlib.Csv.save_fn ?separator writer (List.map ~f:row_of_t csvable)

  let csv_save_out ?separator channel csvable =
    Csvlib.Csv.save_out ?separator channel  (List.map ~f:row_of_t csvable)

  let csv_save ?separator channel csvable =
    Csvlib.Csv.save ?separator channel (List.map ~f:row_of_t csvable)

  let csv_load_in ?separator channel =
    let list = Csvlib.Csv.load_in ?separator channel in
    List.map ~f:t_of_row list

  let csv_load ?separator file =
    let list = Csvlib.Csv.load ?separator file in
    List.map ~f:t_of_row list

end

exception Incomplete_row

module Make_csvable_simple (S : Stringable) : Csvable_simple with type t = S.t =
  struct

    type t = S.t

    let is_csv_atom = true

    let rev_csv_header' acc _ _ = acc

    let rev_csv_header_spec' acc _ _ = acc

    let t_of_row' _ strings =
      match strings with
      | [] -> raise Incomplete_row
      | value :: csvable -> (fun () -> S.of_string value), csvable

    let write_row_of_t' ~is_first ~is_last ~writer _ _ csvable =
      Helper.write ~is_first ~is_last ~writer ~to_string:S.to_string () () csvable

end

module Atom (S : Stringable) : Csvable with type t = S.t =
  Record (Make_csvable_simple (S))


let use_head f lst =
  match lst with
  | [] -> raise Incomplete_row
  | h :: t -> (fun _ -> (f h)), t

exception Illegal_atom of string

(** All the conversion functions *)
let unit_of_row  _ strings =
  use_head (fun string -> if string = "" then () else raise (Illegal_atom string)) strings

let bool_of_row _ strings =
  use_head (fun h ->
    let h = String.uppercase h in
    h = "TRUE") strings

let string_of_row _ strings =
  use_head (fun h -> h) strings

let char_of_row _ strings =
  use_head (fun string ->
    if (String.length string) = 1 then String.get string 0 else raise (Illegal_atom string))
    strings

let int_of_row _ strings =
  use_head int_of_string strings

let float_of_row _ strings =
  use_head float_of_string strings

let int32_of_row _ strings =
  use_head Int32.of_string strings

let int64_of_row _ strings =
  use_head Int64.of_string strings

let nativeint_of_row _ strings =
  use_head Nativeint.of_string strings

let big_int_of_row _ strings =
  use_head Big_int.big_int_of_string strings

let nat_of_row _ strings =
  use_head Nat.nat_of_string strings

let num_of_row _ strings =
  use_head Num.num_of_string strings

let ratio_of_row _ strings =
  use_head Ratio.ratio_of_string strings

type ('a, 'b, 'c) row_of =
  is_first:bool
  -> is_last:bool
  -> writer:(string -> unit)
  -> 'b
  -> 'c
  -> 'a
  -> unit

let write is_first is_last writer to_string t =
  Helper.write ~is_first ~is_last ~writer ~to_string () () t

let row_of_unit ~is_first ~is_last ~writer _ _ t =
  let f () = "" in
  write is_first is_last writer f t

let row_of_bool ~is_first ~is_last ~writer _ _ t =
  let f bool = if bool then "true" else "false" in
  write is_first is_last writer f t

let row_of_string ~is_first ~is_last ~writer _ _ t =
  let f string = string in
  write is_first is_last writer f t

let row_of_char ~is_first ~is_last ~writer _ _ t =
  let f char = String.make 1 char in
  write is_first is_last writer f t

let row_of_int ~is_first ~is_last ~writer _ _ t =
  write is_first is_last writer string_of_int t

let row_of_float ~is_first ~is_last ~writer _ _ t =
  write is_first is_last  writer string_of_float t

let row_of_int32 ~is_first ~is_last  ~writer _ _ t =
  write is_first is_last  writer Int32.to_string t

let row_of_int64 ~is_first ~is_last  ~writer _ _ t =
  write is_first is_last  writer Int64.to_string t

let row_of_nativeint ~is_first ~is_last  ~writer _ _ t =
  write is_first is_last  writer Nativeint.to_string t

let row_of_big_int ~is_first ~is_last  ~writer _ _ t =
  write is_first is_last  writer Big_int.string_of_big_int t

let row_of_nat ~is_first ~is_last  ~writer _ _ t =
  write is_first is_last  writer Nat.string_of_nat t

let row_of_num ~is_first ~is_last  ~writer _ _ t =
  write is_first is_last  writer Num.string_of_num t

let row_of_ratio ~is_first ~is_last  ~writer _ _ t =
  write is_first is_last  writer Ratio.string_of_ratio t