open Prelude

type tree = { car : tree String.Map.t }
type trunk = (Money.polarity * tree) String.Map.t
type path = Base of string | Sub of string * path

let rec path_to_list ?(suffix = []) p =
  match p with
  | Base x -> x :: suffix
  | Sub (x, p') -> path_to_list ~suffix:(x :: suffix) p'

let canonical : trunk =
  let make alist : tree = { car = String.Map.of_alist_exn alist } in
  String.Map.of_alist_exn
    [
      ( "Asset",
        ( Money.Increase_on_debit,
          make
            [
              ("Accounts_receivable", make []);
              ("Bank", make [ ("Savings", make []); ("Checking", make []) ]);
              ("Cash", make []);
              ("Mutual_fund", make []);
              ("Stock", make []);
            ] ) );
      ("Equity", (Money.Increase_on_credit, make []));
      ("Expense", (Money.Increase_on_debit, make []));
      ("Income", (Money.Increase_on_credit, make []));
      ( "Liability",
        ( Money.Increase_on_credit,
          make [ ("Accounts_payable", make []); ("Credit", make []) ] ) );
    ]

(* In this module, only the following function entertains the option
   that the given path may not be valid (i.e., it does not throw an
   exception for invalid paths). *)
let rec get_node : path -> tree option =
  let open Option.Let_syntax in
  function
  | Base x ->
      let%map _, t = Map.find canonical x in
      t
  | Sub (t, p) ->
      let%bind super = get_node p in
      Map.find super.car t

(** Always gives [Some] under valid paths, giving a list of valid paths *)
let children (p : path) : path list =
  let node = Option.value_exn (get_node p) in
  List.map (Map.keys node.car) ~f:(fun sub -> Sub (sub, p))

let sub (p : path) (name : string) : path option =
  let node = Option.value_exn (get_node p) in
  if Map.mem node.car name then Some (Sub (name, p)) else None

let super : path -> path option = function
  | Base _ -> None
  | Sub (_, super) -> Some super

let rec equal_path p1 p2 =
  match (p1, p2) with
  | Base x1, Base x2 -> String.(x1 = x2)
  | Sub (x1, p1'), Sub (x2, p2') -> String.(x1 = x2) && equal_path p1' p2'
  | _, _ -> false

let is_prefix (p : path) ~(prefix : path) : bool =
  List.is_prefix (path_to_list p) ~prefix:(path_to_list prefix)
    ~equal:String.equal

let rec polarity = function
  | Base x ->
      let pol, _ = Map.find_exn canonical x in
      pol
  | Sub (_, p') -> polarity p'

let assert_valid acc =
  match get_node acc with None -> failwith "invalid base account" | _ -> acc

let asset = Base "Asset" |> assert_valid
let equity = Base "Equity" |> assert_valid
let expense = Base "Expense" |> assert_valid
let income = Base "Income" |> assert_valid
let liability = Base "Liability" |> assert_valid