path: root/lib/ledger.ml

open Prelude

type tx_type =
  | Interest_tx
  | Online_banking_tx
  | Recurrent_direct_tx
  | Payment_terminal_tx
  | Cash_payment_tx
  | Atm_tx
  | Auto_save_rounding_tx
  | Batch_tx
  | Direct_debit_tx
  | Periodic_tx

type iban_tag = Account_tag | Counterparty_iban_tag [@@deriving compare, sexp]

type unit_tag = Filed_tag | Google_pay_tag | Auto_round_savings_tag
[@@deriving compare, sexp]

type string_tag =
  | Desc_tag
  | User_tag
  | Counterparty_name_tag
  | Reference_tag
  | Mandate_id_tag
  | Creditor_id_tag
  | Other_party_tag
  | Transaction_tag
  | Terminal_tag
  | Card_seq_no_tag
  | Savings_account_tag
[@@deriving compare, sexp]

module Label = struct
  type 'a t =
    | Iban_label : iban_tag -> Iban.t t
    | String_label : string_tag -> string t
    | Timestamp_label : Time_ns.t t
    | Unit_label : unit_tag -> unit t

  let int_to_cmp x : ('a, 'a) Dmap.cmp =
    if x < 0 then Lt else if x > 0 then Gt else Eq

  let compare (type a1 a2) (v1 : a1 t) (v2 : a2 t) : (a1, a2) Dmap.cmp =
    match (v1, v2) with
    | Iban_label t1, Iban_label t2 -> int_to_cmp @@ [%compare: iban_tag] t1 t2
    | String_label t1, String_label t2 ->
        int_to_cmp @@ [%compare: string_tag] t1 t2
    | Timestamp_label, Timestamp_label -> Eq
    | Unit_label t1, Unit_label t2 -> int_to_cmp @@ [%compare: unit_tag] t1 t2
    | Iban_label _, _ -> Lt
    | String_label _, Iban_label _ -> Gt
    | String_label _, _ -> Lt
    | Timestamp_label, Unit_label _ -> Lt
    | Timestamp_label, _ -> Gt
    | Unit_label _, _ -> Gt
end

module Labels = struct
  include Dmap.Make (Label)

  let sexp_of_binding = function
    | Binding (Iban_label tag, iban) ->
        Sexp.List
          [
            Sexp.Atom "iban"; [%sexp_of: iban_tag] tag; [%sexp_of: Iban.t] iban;
          ]
    | Binding (String_label tag, s) ->
        Sexp.List
          [ Sexp.Atom "string"; [%sexp_of: string_tag] tag; Sexp.Atom s ]
    | Binding (Timestamp_label, ts) ->
        Sexp.List [ Sexp.Atom "timestamp"; [%sexp_of: Time_ns_unix.t] ts ]
    | Binding (Unit_label tag, ()) ->
        Sexp.List [ Sexp.Atom "unit"; [%sexp_of: unit_tag] tag ]

  let binding_of_sexp sexp =
    match sexp with
    | Sexp.List [ Sexp.Atom "iban"; tag_sexp; iban_sexp ] ->
        Binding
          ( Iban_label ([%of_sexp: iban_tag] tag_sexp),
            [%of_sexp: Iban.t] iban_sexp )
    | Sexp.List [ Sexp.Atom "string"; tag_sexp; Sexp.Atom s ] ->
        Binding (String_label ([%of_sexp: string_tag] tag_sexp), s)
    | Sexp.List [ Sexp.Atom "timestamp"; ts_sexp ] ->
        Binding (Timestamp_label, [%of_sexp: Time_ns_unix.t] ts_sexp)
    | Sexp.List [ Sexp.Atom "unit"; tag_sexp ] ->
        Binding (Unit_label ([%of_sexp: unit_tag] tag_sexp), ())
    | _ -> of_sexp_error "Labels.binding_of_sexp: invalid binding" sexp

  let sexp_of_t m = Sexp.List (bindings m |> List.map ~f:sexp_of_binding)

  let t_of_sexp sexp =
    match sexp with
    | Sexp.List labels ->
        Sequence.(of_list labels >>| binding_of_sexp |> to_seq) |> of_seq
    | Sexp.Atom _ -> of_sexp_error "Labels.t_of_sexp: list needed" sexp
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

(*
type scalar =
  | Amount of Money.Amount.t
  | Rate of { in_primary_commodity : Money.Amount.t; rate : Bigdecimal.t }
[@@deriving equal, compare, sexp_of] *)

module Gh = struct
  (* The category of the five top-level categories *)
  type global

  (* The five top-level categories *)
  type asset [@@deriving sexp_of]
  type equity [@@deriving sexp_of]
  type expense [@@deriving sexp_of]
  type income [@@deriving sexp_of]
  type liability [@@deriving sexp_of]

  (* Subcategories of assets *)
  type bank [@@deriving sexp_of]

  (* No subcategories *)
  type final [@@deriving sexp_of]
end

module Account_structure0 (F : sig
  type 'a t [@@deriving sexp_of]
end) =
struct
  type 'a f =
    | Accounts_payable : Gh.final f F.t -> Gh.liability f
    | Accounts_receivable : Gh.final f F.t -> Gh.asset f
    | Bank : Gh.bank f F.t -> Gh.asset f
    | Cash : Gh.final f F.t -> Gh.asset f
    | Credit : Gh.final f F.t -> Gh.liability f
    | Mutual_fund : Gh.final f F.t -> Gh.asset f
    | Stock : Gh.final f F.t -> Gh.asset f
    | Savings : Gh.final f F.t -> Gh.bank f
    | Checking : Gh.final f F.t -> Gh.bank f
  [@@deriving sexp_of]

  type t0 =
    | Asset of Gh.asset f F.t
    | Equity of Gh.equity f F.t
    | Expense of Gh.expense f F.t
    | Income of Gh.income f F.t
    | Liability of Gh.liability f F.t
  [@@deriving sexp_of]
end

module Account_type = struct
  type 'a elem = Leaf | Node of 'a [@@deriving sexp_of]

  include Account_structure0 (struct
    type 'a t = 'a elem [@@deriving sexp_of]
  end)
end

module Account_structure (F : sig
  type 'a t [@@deriving sexp_of]
end) =
struct
  include Account_structure0 (F)

  module Mapper = struct
    type nonrec 'b t = {
      car :
        'a.
        'a f F.t ->
        ('a Account_type.f Account_type.elem -> Account_type.t0) ->
        ('b * 'a f F.t) option;
    }

    let map (type b c) (f : c t) (mkt : b Account_type.f -> Account_type.t0) :
        b f -> (c * b f) option =
      let open Option.Let_syntax in
      function
      | Accounts_payable v ->
          let%map c, v' =
            f.car v (fun el -> mkt (Account_type.Accounts_payable el))
          in
          (c, Accounts_payable v')
      | Accounts_receivable v ->
          let%map c, v' =
            f.car v (fun el -> mkt (Account_type.Accounts_receivable el))
          in
          (c, Accounts_receivable v')
      | Bank v ->
          let%map c, v' = f.car v (fun el -> mkt (Account_type.Bank el)) in
          (c, Bank v')
      | Cash v ->
          let%map c, v' = f.car v (fun el -> mkt (Account_type.Cash el)) in
          (c, Cash v')
      | Credit v ->
          let%map c, v' = f.car v (fun el -> mkt (Account_type.Credit el)) in
          (c, Credit v')
      | Mutual_fund v ->
          let%map c, v' =
            f.car v (fun el -> mkt (Account_type.Mutual_fund el))
          in
          (c, Mutual_fund v')
      | Stock v ->
          let%map c, v' = f.car v (fun el -> mkt (Account_type.Stock el)) in
          (c, Stock v')
      | Savings v ->
          let%map c, v' = f.car v (fun el -> mkt (Account_type.Savings el)) in
          (c, Savings v')
      | Checking v ->
          let%map c, v' = f.car v (fun el -> mkt (Account_type.Checking el)) in
          (c, Checking v')
  end

  module Folder = struct
    type nonrec 'b t = { car : 'a. 'a f F.t -> 'b }

    let fold (type b c) (f : c t) : b f -> c = function
      | Accounts_payable v -> f.car v
      | Accounts_receivable v -> f.car v
      | Bank v -> f.car v
      | Cash v -> f.car v
      | Credit v -> f.car v
      | Mutual_fund v -> f.car v
      | Stock v -> f.car v
      | Savings v -> f.car v
      | Checking v -> f.car v
  end
end

module Typed_account_path = struct
  type 'a elem = Leaf | Node of string * 'a | Ind of string * 'a elem
  [@@deriving sexp_of]

  include Account_structure (struct
    type 'a t = 'a elem [@@deriving sexp_of]
  end)
end

module Account_path = struct
  type t = string list [@@deriving compare, sexp]

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

module Account_hierarchy = struct
  (* The contents of an account of category 'a *)
  type 'a core =
    (* Comprises of subaccounts of its subcategories *)
    | Node of 'a String.Map.t
    (* Comprises of subaccounts of its own category *)
    | Ind of 'a account String.Map.t
    (* Has no subaccounts, has a balance in a certain commodity *)
    | Leaf of Commodity_id.t * Money.Diff.t
  [@@deriving sexp_of]

  and extra = { description : String.t } [@@deriving sexp_of]
  and 'a account = extra * 'a core

  include Account_structure (struct
    type 'a t = 'a account [@@deriving sexp_of]
  end)

  (* All accounts *)
  type world = t0 String.Map.t

  let rec unsafe_alter_aux (subaid : Account_path.t)
      (f :
        Account_type.t0 ->
        extra ->
        Commodity_id.t ->
        Money.Diff.t ->
        'a * extra * Money.Diff.t) : 'a Mapper.t =
    {
      car =
        (fun in_acc mkt ->
          let open Option.Let_syntax in
          match (subaid, in_acc) with
          | [], (extra, Leaf (acc_comm, acc_bal)) ->
              let x, extra', acc_bal' =
                f (mkt Account_type.Leaf) extra acc_comm acc_bal
              in
              Some (x, (extra', Leaf (acc_comm, acc_bal')))
          | [], _ -> None
          | subaid0 :: subaid, (extra, Node subaccs) ->
              let%bind subacc = Map.find subaccs subaid0 in
              let%map x, subacc' =
                Mapper.map
                  (unsafe_alter_aux subaid f)
                  (fun k -> mkt (Node k))
                  subacc
              in
              (x, (extra, Node (Map.set subaccs ~key:subaid0 ~data:subacc')))
          | subaid0 :: subaid, (extra, Ind subaccs) ->
              let%bind subacc = Map.find subaccs subaid0 in
              let%map x, subacc' = (unsafe_alter_aux subaid f).car subacc mkt in
              (x, (extra, Ind (Map.set subaccs ~key:subaid0 ~data:subacc')))
          | _ :: _, (_, Leaf _) -> None);
    }

  let unsafe_alter (aid : Account_path.t)
      (f :
        Account_type.t0 ->
        extra ->
        Commodity_id.t ->
        Money.Diff.t ->
        'a * extra * Money.Diff.t) (w : world) : ('a * world) option =
    match aid with
    | [] -> None
    | aid0 :: subaid -> (
        let open Option.Let_syntax in
        match%bind Map.find w aid0 with
        | Asset acc ->
            let%map x, acc' =
              (unsafe_alter_aux subaid f).car acc (fun k -> Asset k)
            in
            (x, Map.set w ~key:aid0 ~data:(Asset acc'))
        | Expense acc ->
            let%map x, acc' =
              (unsafe_alter_aux subaid f).car acc (fun k -> Expense k)
            in
            (x, Map.set w ~key:aid0 ~data:(Expense acc'))
        | Income acc ->
            let%map x, acc' =
              (unsafe_alter_aux subaid f).car acc (fun k -> Income k)
            in
            (x, Map.set w ~key:aid0 ~data:(Income acc'))
        | Liability acc ->
            let%map x, acc' =
              (unsafe_alter_aux subaid f).car acc (fun k -> Liability k)
            in
            (x, Map.set w ~key:aid0 ~data:(Liability acc'))
        | Equity acc ->
            let%map x, acc' =
              (unsafe_alter_aux subaid f).car acc (fun k -> Equity k)
            in
            (x, Map.set w ~key:aid0 ~data:(Equity acc')))

  (** Update the balance (debit/credit ([dc])) of account [aid] [by_amount]
      (commodity: [in_comm]) in [world], giving the updated world and the pre
      and post balances for [aid] iff the account exists in [world]. *)
  let update_bal aid dc by_amount in_comm (w : world) :
      (Money.Diff.t * Money.Diff.t * world) option =
    let open Option.Let_syntax in
    let%bind mres, w' =
      unsafe_alter aid
        (fun acc_type acc_extra acc_comm acc_bal ->
          if not ([%equal: Commodity_id.t] acc_comm in_comm) then
            (None, acc_extra, acc_bal)
          else
            let on_debit =
              match acc_type with
              | Asset _ -> `Incr
              | Expense _ -> `Incr
              | Income _ -> `Decr
              | Liability _ -> `Decr
              | Equity _ -> `Decr
            in
            let acc_bal' =
              Money.Diff.(acc_bal + of_amount by_amount dc ~on_debit)
            in
            (Some (acc_bal, acc_bal'), acc_extra, acc_bal'))
        w
    in
    let%map pre_bal, post_bal = mres in
    (pre_bal, post_bal, w')

  let get_bal aid (w : world) : (Commodity_id.t * Money.Diff.t) option =
    let open Option.Let_syntax in
    let%map cb, _world' =
      unsafe_alter aid
        (fun _acc_type acc_extra acc_comm acc_bal ->
          ((acc_comm, acc_bal), acc_extra, acc_bal))
        w
    in
    cb

  let add_balance_maps m1 m2 : Money.Diff.t Commodity_id.Map.t =
    Map.merge m1 m2 ~f:(fun ~key:_comm -> function
      | `Both (b1, b2) -> Some Money.Diff.(b1 + b2)
      | `Left b | `Right b -> Some b)

  let rec collect_balances : type a.
      a f account -> Money.Diff.t Commodity_id.Map.t = function
    | _extra, Leaf (acc_comm, acc_bal) ->
        Commodity_id.Map.singleton acc_comm acc_bal
    | _extra, Ind subaccs ->
        Map.fold subaccs ~init:Commodity_id.Map.empty
          ~f:(fun ~key:_ ~data:subacc comm_bal_sums ->
            add_balance_maps comm_bal_sums (collect_balances subacc))
    | _extra, Node subaccs ->
        Map.fold subaccs ~init:Commodity_id.Map.empty
          ~f:(fun ~key:_ ~data:subacc comm_bal_sums ->
            add_balance_maps comm_bal_sums
              (Folder.fold { car = collect_balances } subacc))

  type delete_error = Not_found | Nonzero_balance

  (*
  let rec delete_aux : type a. (Account_path.t * a f account) -> (a f account, delete_error) result = function
    | [], (extra, Leaf (_acc_comm, acc_bal)) ->
       if Money.Diff.(acc_bal = ~$0) then 
  
  let delete (aid : Account_path.t) (w : world) =
   *)

  let world_inst : world =
    String.Map.of_alist_exn
      [
        ( "Assets",
          Asset
            ( { description = "assets" },
              Ind
                (String.Map.of_alist_exn
                   [
                     ( "Current",
                       ( { description = "current" },
                         Node
                           (String.Map.of_alist_exn
                              [
                                ( "Checking",
                                  Bank
                                    ( { description = "bnak accounts" },
                                      Ind
                                        (String.Map.of_alist_exn
                                           [
                                             ( "ING",
                                               ( { description = "ING bank" },
                                                 Leaf ("EUC", Money.Diff.(~$0))
                                               ) );
                                             ( "N26",
                                               ( { description = "ING bank" },
                                                 Leaf ("EUC", Money.Diff.(~$0))
                                               ) );
                                           ]) ) );
                              ]) ) );
                   ]) ) );
      ]
end

module Bal_assert = struct
  type t = { account : Account_path.t; labels : Labels.t; bal : Money.Diff.t }
  [@@deriving sexp_of]
end

module Account_decl = struct
  type t = {
    type_ : Account_type.t0;
    parent : Account_path.t;
    name : string;
    commodity : Commodity_id.t;
    extra : Account_hierarchy.extra;
  }
  [@@deriving sexp_of]
end

module Tx : sig
  type entry = {
    dc : Debit_credit.t;
    commodity : Commodity_id.t;
    amount : Money.Amount.t;
    assertion : Money.Diff.t option;
  }

  (* Private because we only want to allow constructing balanced transactions. *)
  type t = private {
    cleared : Date.t option;
    entries : entry Account_path.Map.t;
    labels : Labels.t;
  }

  type error = Unbalanced

  val make :
    cleared:Date.t option ->
    entries:entry Account_path.Map.t ->
    labels:Labels.t ->
    (t, error) result

  val sexp_of_t : t -> Sexp.t
end = struct
  type entry = {
    dc : Debit_credit.t;
    commodity : Commodity_id.t;
    amount : Money.Amount.t;
    assertion : Money.Diff.t option;
  }
  [@@deriving sexp_of]

  type t = {
    cleared : Date.t option;
    entries : entry Account_path.Map.t;
    labels : Labels.t;
  }
  [@@deriving sexp_of]

  type error = Unbalanced

  let is_balanced entries =
    Map.fold entries ~init:Commodity_id.Map.empty
      ~f:(fun ~key:_ ~data comm_balances ->
        Map.update comm_balances data.commodity ~f:(fun ocomm_bal ->
            let comm_bal = Option.value ocomm_bal ~default:Money.Diff.(~$0) in
            match data.dc with
            | Debit_credit.Debit -> Money.Diff.(comm_bal +% data.amount)
            | Debit_credit.Credit -> Money.Diff.(comm_bal -% data.amount)))
    |> Map.for_all ~f:(fun comm_bal -> Money.Diff.(comm_bal = ~$0))

  let make ~cleared ~entries ~labels =
    if not (is_balanced entries) then Error Unbalanced
    else Ok { cleared; entries; labels }
end

type item =
  | Tx_item of Tx.t
  | Bal_assert_item of Bal_assert.t
  | Account_decl_item of Account_decl.t
[@@deriving sexp_of]

type t = item list [@@deriving sexp_of]

module World = struct
  type t = Account_hierarchy.world

  let empty : t = String.Map.empty

  let apply_tx (tx : Tx.t) world : t option =
    Map.fold_option tx.entries ~init:world
      ~f:(fun ~key:aid ~(data : Tx.entry) world ->
        let open Option.Let_syntax in
        let%bind _old_bal, new_bal, world =
          Account_hierarchy.update_bal aid data.dc data.amount data.commodity
            world
        in
        match data.assertion with
        | None -> Some world
        | Some bal_ass ->
            if Money.Diff.(bal_ass = new_bal) then Some world else None)

  let apply_ba (ba : Bal_assert.t) world : t option =
    let open Option.Let_syntax in
    let%bind _comm, bal = Account_hierarchy.get_bal ba.account world in
    if not Money.Diff.(bal = ba.bal) then None else Some world

  let apply_ad (_ad : Account_decl.t) _world : t option = None

  let apply : item -> t -> t option = function
    | Tx_item tx -> apply_tx tx
    | Bal_assert_item ba -> apply_ba ba
    | Account_decl_item ad -> apply_ad ad
end

module Ctxd_item = struct end

let make = Fn.id