open Prelude

(** The 'kernel' of account types: a hierarchy of valid types. A valid type is a
    path that leads to a node in the hierarchy. *)
module Type_hierarchy : sig
  type path
  type polarity = Increase_on_debit | Increase_on_credit

  val children : path -> path list
  val sub : path -> string -> path option
  val super : path -> path option
  val equal_path : path -> path -> bool
  val is_prefix : path -> prefix:path -> bool
  val polarity : path -> polarity
  val asset : path
  val equity : path
  val expense : path
  val income : path
  val liability : path
end = struct
  type polarity = Increase_on_debit | Increase_on_credit
  type tree = { car : tree String.Map.t }
  type trunk = { car : (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 mk alist : tree = { car = String.Map.of_alist_exn alist } in
    {
      car =
        String.Map.of_alist_exn
          [
            ( "Asset",
              ( Increase_on_debit,
                mk
                  [
                    ("Accounts_receivable", mk []);
                    ("Bank", mk [ ("Savings", mk []); ("Checking", mk []) ]);
                    ("Cash", mk []);
                    ("Mutual_fund", mk []);
                    ("Stock", mk []);
                  ] ) );
            ("Equity", (Increase_on_credit, mk []));
            ("Expense", (Increase_on_debit, mk []));
            ("Income", (Increase_on_credit, mk []));
            ( "Liability",
              ( Increase_on_credit,
                mk [ ("Accounts_payable", mk []); ("Credit", mk []) ] ) );
          ];
    }

  (* In this module, only the following two 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.car 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.car 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
end

module Type = struct
  type t = Type_hierarchy.path [@@deriving equal]

  let rec base (t : t) : t option =
    match Type_hierarchy.super t with
    | None -> (* [t] is the root type *) None
    | Some t' ->
        (* [t] is a base type if its supertype is the root type *)
        Some (Option.value (base t') ~default:t)

  (** [a] is a strict supertype of [b] *)
  let is_strict_super a b =
    Type_hierarchy.is_prefix b ~prefix:a && not ([%equal: t] a b)
end

module Debit_credit = struct
  type t = Debit | Credit [@@deriving string, sexp_of]

  (*  let opposite = function Debit -> Credit | Credit -> Debit *)
end

module Money = struct
  module Amount : sig
    type t

    val equal : t -> t -> bool
    val compare : t -> t -> int
    val of_bigint : Bigint.t -> t option
    val to_bigint : t -> Bigint.t
    val ( + ) : t -> t -> t
    val ( = ) : t -> t -> bool
    val sexp_of_t : t -> Sexp.t
    val zero : t
  end = struct
    type t = Bigint.t [@@deriving sexp_of]

    let equal = Bigint.equal
    let compare = Bigint.compare
    let of_bigint x = if Bigint.(zero <= x) then Some x else None
    let to_bigint x = x
    let ( + ) x y = Bigint.(x + y)
    let ( = ) = equal
    let zero = Bigint.zero
  end

  module Diff : sig
    type t

    val equal : t -> t -> bool
    val compare : t -> t -> int
    val of_bigint : Bigint.t -> t
    val to_bigint : t -> Bigint.t
    val ( + ) : t -> t -> t
    val ( +% ) : t -> Amount.t -> t
    val ( - ) : t -> t -> t
    val ( -% ) : t -> Amount.t -> t
    val ( = ) : t -> t -> bool
    val neg : t -> t
    val ( ~$ ) : int -> t
    val sexp_of_t : t -> Sexp.t

    val of_amount :
      Amount.t -> Debit_credit.t -> on_debit:[ `Incr | `Decr ] -> t
  end = struct
    type t = Bigint.t [@@deriving sexp_of]

    let equal = Bigint.equal
    let compare = Bigint.compare
    let of_bigint x = x
    let to_bigint x = x
    let ( + ) x y = Bigint.(x + y)
    let ( +% ) x y = x + of_bigint (Amount.to_bigint y)
    let ( - ) x y = Bigint.(x - y)
    let ( -% ) x y = x - of_bigint (Amount.to_bigint y)
    let ( = ) = equal
    let neg = Bigint.neg
    let ( ~$ ) = Fn.compose of_bigint Bigint.of_int

    let of_amount x (dc : Debit_credit.t) ~on_debit =
      match (dc, on_debit) with
      | Debit, `Incr -> of_bigint (Amount.to_bigint x)
      | Credit, `Incr -> neg (of_bigint (Amount.to_bigint x))
      | Debit, `Decr -> neg (of_bigint (Amount.to_bigint x))
      | Credit, `Decr -> of_bigint (Amount.to_bigint x)
  end
end

module Commodity_id = struct
  type t = string [@@deriving equal, compare, sexp]

  module Map = Map.Make (struct
    type nonrec t = t [@@deriving equal, compare, sexp]
  end)
end

module Account = struct
  type t = Type.t * node

  and node =
    (* Balance in some commodity *)
    | Leaf of Commodity_id.t * Money.Diff.t
    | Subtree of node String.Map.t
end