accounting/src/queries/balance.rs

use std::{collections::HashMap, time::Instant};

use crate::core::{Amount, Ledger};
use chrono::NaiveDate;
use rust_decimal_macros::dec;

use super::{
    base::{self, DataValue, Function},
    functions::ComparisonFunction,
    transaction::{PostingData, PostingField, TransactionField},
};

pub enum Query {
    StartDate(NaiveDate),
    EndDate(NaiveDate),
}

pub fn balance2(
    ledger: &Ledger,
    end_date: NaiveDate,
    convert_to_unit: Option<&str>,
) -> HashMap<u32, Vec<Amount>> {
    let q = ComparisonFunction::new(
        "<=",
        base::Query::from_field(PostingField::Transaction(TransactionField::Date)),
        base::Query::from(DataValue::from(end_date)),
    )
    .unwrap();

    let convert_to_unit = convert_to_unit.map(|u| ledger.get_unit_by_symbol(u).unwrap());

    let postings = ledger
        .get_transactions()
        .iter()
        .map(|t| {
            t.get_postings().iter().map(|p| PostingData {
                ledger,
                posting: p,
                parent_transaction: t,
            })
        })
        .flatten();

    let filtered_postings =
        postings.filter(|data| q.evaluate(data).map(|v| bool::from(v)).unwrap_or(false));

    let mut accounts = HashMap::new();
    for posting_data in filtered_postings {
        let posting = posting_data.posting;
        let mut amount = *posting.get_amount();
        if let Some(new_unit) = convert_to_unit {
            if amount.unit_id != new_unit.get_id() {
                let price = ledger.get_price_on_date(end_date, amount.unit_id, new_unit.get_id());
                if let Some(price) = price {
                    amount = Amount {
                        value: amount.value * price,
                        unit_id: new_unit.get_id(),
                    };
                }
            }
        }
        let account_vals = accounts
            .entry(posting.get_account_id())
            .or_insert(HashMap::new());
        let a = account_vals.entry(amount.unit_id).or_insert(dec!(0));
        *a += amount.value;
    }

    accounts
        .iter()
        .map(|(&k, v)| {
            (
                k,
                v.into_iter()
                    .map(|(&unit_id, &value)| Amount { value, unit_id })
                    .collect(),
            )
        })
        .collect()
}

pub fn balance3(ledger: &Ledger, query: &base::Query<PostingField>) -> HashMap<u32, Vec<Amount>> {
    // let q = ComparisonFunction::new(
    //     "<=",
    //     base::Query::from_field(PostingField::Transaction(TransactionField::Date)),
    //     base::Query::from(DataValue::from(end_date)),
    // )
    // .unwrap();

    let t0 = Instant::now();

    let postings = ledger
        .get_transactions()
        .iter()
        .map(|t| {
            t.get_postings().iter().map(|p| PostingData {
                ledger,
                posting: p,
                parent_transaction: t,
            })
        })
        .flatten();

    let t1 = Instant::now();

    let filtered_postings =
        postings.filter(|data| query.evaluate(data).map(|v| bool::from(v)).unwrap_or(false));

    let t2 = Instant::now();

    // println!("{:?} {:?}", t1-t0, t2-t1);

    let mut accounts = HashMap::new();
    for posting_data in filtered_postings {
        let posting = posting_data.posting;
        let amount = posting.get_amount();
        let account_vals = accounts
            .entry(posting.get_account_id())
            .or_insert(HashMap::new());
        let a = account_vals.entry(amount.unit_id).or_insert(dec!(0));
        *a += amount.value;
    }

    accounts
        .iter()
        .map(|(&k, v)| {
            (
                k,
                v.into_iter()
                    .map(|(&unit_id, &value)| Amount { value, unit_id })
                    .collect(),
            )
        })
        .collect()
}

pub fn balance(ledger: &Ledger, query: &[Query]) -> HashMap<u32, Vec<Amount>> {
    let relevant_transactions = ledger.get_transactions().iter().filter(|txn| {
        query.iter().all(|q| match q {
            Query::StartDate(date) => txn.get_date() >= *date,
            Query::EndDate(date) => txn.get_date() <= *date,
        })
    });

    let mut accounts = HashMap::new();

    for txn in relevant_transactions.clone() {
        for posting in txn.get_postings() {
            let amount = posting.get_amount();
            let account_vals = accounts
                .entry(posting.get_account_id())
                .or_insert(HashMap::new());
            let a = account_vals.entry(amount.unit_id).or_insert(dec!(0));
            *a += amount.value;
        }
    }

    accounts
        .iter()
        .map(|(&k, v)| {
            (
                k,
                v.into_iter()
                    .map(|(&unit_id, &value)| Amount { value, unit_id })
                    .collect(),
            )
        })
        .collect()
}