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 Account_structure = struct
  module Categories = struct
    (* 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 type S = sig
    type 'a t [@@deriving sexp_of]

    type 'a f =
      | Accounts_payable : Categories.final f t -> Categories.liability f
      | Accounts_receivable : Categories.final f t -> Categories.asset f
      | Bank : Categories.bank f t -> Categories.asset f
      | Cash : Categories.final f t -> Categories.asset f
      | Credit : Categories.final f t -> Categories.liability f
      | Mutual_fund : Categories.final f t -> Categories.asset f
      | Stock : Categories.final f t -> Categories.asset f
      | Savings : Categories.final f t -> Categories.bank f
      | Checking : Categories.final f t -> Categories.bank f
    [@@deriving sexp_of]

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

  module Make (F : sig
    type 'a t [@@deriving sexp_of]
  end) : S with type 'a t = 'a F.t = struct
    include F

    type 'a f =
      | Accounts_payable : Categories.final f F.t -> Categories.liability f
      | Accounts_receivable : Categories.final f F.t -> Categories.asset f
      | Bank : Categories.bank f F.t -> Categories.asset f
      | Cash : Categories.final f F.t -> Categories.asset f
      | Credit : Categories.final f F.t -> Categories.liability f
      | Mutual_fund : Categories.final f F.t -> Categories.asset f
      | Stock : Categories.final f F.t -> Categories.asset f
      | Savings : Categories.final f F.t -> Categories.bank f
      | Checking : Categories.final f F.t -> Categories.bank f
    [@@deriving sexp_of]

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

  module Gen_f_cons = struct
    module Accounts_payable = struct
      type outer = Categories.liability
      type inner = Categories.final

      module Specialize (G : S) = struct
        let cons v = G.Accounts_payable v
      end
    end

    module type S = sig
      type inner
      type outer

      module Specialize : functor (G : S) -> sig
        val cons : inner G.f G.t -> outer G.f
      end
    end
  end

  module Visitor
      (G : S)
      (Acc : sig
        type 'a t
      end) =
  struct
    type nonrec 'outer t = {
      car :
        'inner.
        'inner G.f G.t ->
        (module Gen_f_cons.S with type inner = 'inner and type outer = 'outer) ->
        'outer Acc.t;
    }

    let visit (type a) (f : a t) : a G.f -> a Acc.t = function
      | G.Accounts_payable v -> f.car v (module Gen_f_cons.Accounts_payable)
      | _ -> failwith "kaas"
    (*
      | Accounts_receivable v ->
          f.car v (fun inner -> H.Accounts_receivable inner)
      | Bank v -> f.car v (fun inner -> H.Bank inner)
      | Cash v -> f.car v (fun inner -> H.Cash inner)
      | Credit v -> f.car v (fun inner -> H.Credit inner)
      | Mutual_fund v -> f.car v (fun inner -> H.Mutual_fund inner)
      | Stock v -> f.car v (fun inner -> H.Stock inner)
      | Savings v -> f.car v (fun inner -> H.Savings inner)
      | Checking v -> f.car v (fun inner -> H.Checking inner) *)
  end

  module Basic_visitor (G : S) = struct
    type nonrec 'b t = { car : 'a. 'a G.f G.t -> 'b }

    let visit (type b c) (f : c t) : b G.f -> c =
      let module Inst =
        Visitor
          (G)
          (struct
            type 'a t = c
          end)
      in
      Inst.visit { car = (fun v _cons -> f.car v) }
  end
end

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

  include Account_structure.Make (struct
    type 'a t = 'a elem [@@deriving sexp_of]
  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.Make (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

  module Structure = Account_structure.Make (struct
    type 'a t = 'a account [@@deriving sexp_of]
  end)

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

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

    let map (type b c) (f : c t) (mkt : b Account_type.f -> Account_type.t0) :
        b Structure.f -> (c * b Structure.f) option =
      let module Inst =
        Account_structure.Visitor
          (Structure)
          (struct
            type 'b t = (c * 'b Structure.f) option
          end)
      in
      Inst.visit
        {
          car =
            (fun (type inner)
              v
              (module Gen_cons : Account_structure.Gen_f_cons.S
                with type inner = inner
                 and type outer = b)
            ->
              let open Option.Let_syntax in
              let module Type_cons = Gen_cons.Specialize (Account_type) in
              let module Own_cons = Gen_cons.Specialize (Structure) in
              let%map c, v' = f.car v (fun el -> mkt (Type_cons.cons el)) in
              (c, Own_cons.cons v'));
        }
  end

  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 Structure.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 ->
            let module Visitor = Account_structure.Basic_visitor (Structure) in
            add_balance_maps comm_bal_sums
              (Visitor.visit { 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",
          Structure.Asset
            ( { description = "assets" },
              Ind
                (String.Map.of_alist_exn
                   [
                     ( "Current",
                       ( { description = "current" },
                         Node
                           (String.Map.of_alist_exn
                              [
                                ( "Checking",
                                  Structure.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