"""core.py - Common data classes for reporting functionality"""

# Copyright © 2020  Brett Smith

#

# This program is free software: you can redistribute it and/or modify

# it under the terms of the GNU Affero General Public License as published by

# the Free Software Foundation, either version 3 of the License, or

# (at your option) any later version.

#

# This program is distributed in the hope that it will be useful,

# but WITHOUT ANY WARRANTY; without even the implied warranty of

# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

# GNU Affero General Public License for more details.

#

# You should have received a copy of the GNU Affero General Public License

# along with this program.  If not, see <https://www.gnu.org/licenses/>.

import collections

from decimal import Decimal

from .. import data

from typing import (

    overload,

    Dict,

    Iterable,

    Iterator,

    List,

    Mapping,

    Optional,

    Sequence,

    Set,

    Tuple,

    Union,

)

from ..beancount_types import (

    MetaKey,

    MetaValue,

)

class Balance(Mapping[str, data.Amount]):

    """A collection of amounts mapped by currency

    Each key is a Beancount currency string, and each value represents the

    balance in that currency.

    """

    __slots__ = ('_currency_map',)

    def __init__(self,

                 source: Union[Iterable[Tuple[str, data.Amount]],

                               Mapping[str, data.Amount]]=(),

    ) -> None:

        if isinstance(source, Mapping):

            source = source.items()

        self._currency_map = {

            currency: amount.number for currency, amount in source

        }

    def __repr__(self) -> str:

        return f"{type(self).__name__}({self._currency_map!r})"

    def __getitem__(self, key: str) -> data.Amount:

        return data.Amount(self._currency_map[key], key)

    def __iter__(self) -> Iterator[str]:

        return iter(self._currency_map)

    def __len__(self) -> int:

        return len(self._currency_map)

    def is_zero(self) -> bool:

        return all(number == 0 for number in self._currency_map.values())

class MutableBalance(Balance):

    __slots__ = ()

    def add_amount(self, amount: data.Amount) -> None:

        try:

            self._currency_map[amount.currency] += amount.number

        except KeyError:

            self._currency_map[amount.currency] = amount.number

class RelatedPostings(Sequence[data.Posting]):

    """Collect and query related postings

    This class provides common functionality for collecting related postings

    and running queries on them: iterating over them, tallying their balance,

    etc.

    This class doesn't know anything about how the postings are related. That's

    entirely up to the caller.

    A common pattern is to use this class with collections.defaultdict

    to organize postings based on some key::

        report = collections.defaultdict(RelatedPostings)

        for txn in transactions:

            for post in Posting.from_txn(txn):

                if should_report(post):

                    key = post_key(post)

                    report[key].add(post)

    """

    def __init__(self) -> None:

        self._postings: List[data.Posting] = []

    @overload

    def __getitem__(self, index: int) -> data.Posting: ...

    @overload

    def __getitem__(self, s: slice) -> Sequence[data.Posting]: ...

    def __getitem__(self,

                    index: Union[int, slice],

    ) -> Union[data.Posting, Sequence[data.Posting]]:

        if isinstance(index, slice):

            raise NotImplementedError("RelatedPostings[slice]")

        else:

            return self._postings[index]

    def __len__(self) -> int:

        return len(self._postings)

    def add(self, post: data.Posting) -> None:

        self._postings.append(post)

    def iter_with_balance(self) -> Iterable[Tuple[data.Posting, Balance]]:

        balance = MutableBalance()

        for post in self:

            balance.add_amount(post.units)

            yield post, balance

    def balance(self) -> Balance:

        for _, balance in self.iter_with_balance():

            pass

        try:

            return balance

        except NameError:

            return Balance()

    def meta_values(self,

                    key: MetaKey,

                    default: Optional[MetaValue]=None,

    ) -> Set[Optional[MetaValue]]:

        return {post.meta.get(key, default) for post in self}

"""test_reports_related_postings - Unit tests for RelatedPostings"""

# Copyright © 2020  Brett Smith

#

# This program is free software: you can redistribute it and/or modify

# it under the terms of the GNU Affero General Public License as published by

# the Free Software Foundation, either version 3 of the License, or

# (at your option) any later version.

#

# This program is distributed in the hope that it will be useful,

# but WITHOUT ANY WARRANTY; without even the implied warranty of

# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

# GNU Affero General Public License for more details.

#

# You should have received a copy of the GNU Affero General Public License

# along with this program.  If not, see <https://www.gnu.org/licenses/>.

import collections

import datetime

import itertools

from decimal import Decimal

import pytest

from . import testutil

from conservancy_beancount import data

from conservancy_beancount.reports import core

def accruals_and_payments(acct, src_acct, dst_acct, start_date, *amounts):

    dates = testutil.date_seq(start_date)

    for amt, currency in amounts:

        yield testutil.Transaction(date=next(dates), postings=[

            (acct, amt, currency),

            (dst_acct if amt < 0 else src_acct, -amt, currency),

        ])

@pytest.fixture

def credit_card_cycle():

    return list(accruals_and_payments(

        'Liabilities:CreditCard',

        'Assets:Checking',

        'Expenses:Other',

        datetime.date(2020, 4, 1),

        (-110, 'USD'),

        (110, 'USD'),

        (-120, 'USD'),

        (120, 'USD'),

    ))

@pytest.fixture

def two_accruals_three_payments():

    return list(accruals_and_payments(

        'Assets:Receivable:Accounts',

        'Income:Donations',

        'Assets:Checking',

        datetime.date(2020, 4, 10),

        (440, 'USD'),

        (-230, 'USD'),

        (550, 'EUR'),

        (-210, 'USD'),

        (-550, 'EUR'),

    ))

def test_balance_empty():

    balance = core.RelatedPostings().balance()

    assert not balance

    assert balance.is_zero()

def test_balance_credit_card(credit_card_cycle):

    related = core.RelatedPostings()

    assert related.balance() == testutil.balance_map()

    expected = Decimal()

    for txn in credit_card_cycle:

        post = txn.postings[0]

        expected += post.units.number

        related.add(post)

        assert related.balance() == testutil.balance_map(USD=expected)

    assert expected == 0

def check_iter_with_balance(entries):

    expect_posts = [txn.postings[0] for txn in entries]

    expect_balances = []

    balance_tally = collections.defaultdict(Decimal)

    related = core.RelatedPostings()

    for post in expect_posts:

        number, currency = post.units

        balance_tally[currency] += number

        expect_balances.append(testutil.balance_map(balance_tally.items()))

        related.add(post)

    for (post, balance), exp_post, exp_balance in zip(

            related.iter_with_balance(),

            expect_posts,

            expect_balances,

    ):

        assert post is exp_post

        assert balance == exp_balance

    assert post is expect_posts[-1]

    assert related.balance() == expect_balances[-1]

def test_iter_with_balance_empty():

    assert not list(core.RelatedPostings().iter_with_balance())

def test_iter_with_balance_credit_card(credit_card_cycle):

    check_iter_with_balance(credit_card_cycle)

def test_iter_with_balance_two_acccruals(two_accruals_three_payments):

    check_iter_with_balance(two_accruals_three_payments)

def test_meta_values_empty():

    related = core.RelatedPostings()

    assert related.meta_values('key') == set()

def test_meta_values_no_match():

    related = core.RelatedPostings()

    related.add(testutil.Posting('Income:Donations', -1, metakey='metavalue'))

    assert related.meta_values('key') == {None}

def test_meta_values_no_match_default_given():

    related = core.RelatedPostings()

    related.add(testutil.Posting('Income:Donations', -1, metakey='metavalue'))

    assert related.meta_values('key', '') == {''}

def test_meta_values_one_match():

    related = core.RelatedPostings()

    related.add(testutil.Posting('Income:Donations', -1, key='metavalue'))

    assert related.meta_values('key') == {'metavalue'}

def test_meta_values_some_match():

    related = core.RelatedPostings()

    related.add(testutil.Posting('Income:Donations', -1, key='1'))

    related.add(testutil.Posting('Income:Donations', -2, metakey='2'))

    assert related.meta_values('key') == {'1', None}

def test_meta_values_some_match_default_given():

    related = core.RelatedPostings()

    related.add(testutil.Posting('Income:Donations', -1, key='1'))

    related.add(testutil.Posting('Income:Donations', -2, metakey='2'))

    assert related.meta_values('key', '') == {'1', ''}

def test_meta_values_all_match():

    related = core.RelatedPostings()

    related.add(testutil.Posting('Income:Donations', -1, key='1'))

    related.add(testutil.Posting('Income:Donations', -2, key='2'))

    assert related.meta_values('key') == {'1', '2'}

def test_meta_values_all_match_one_value():

    related = core.RelatedPostings()

    related.add(testutil.Posting('Income:Donations', -1, key='1'))

    related.add(testutil.Posting('Income:Donations', -2, key='1'))

    assert related.meta_values('key') == {'1'}

def test_meta_values_all_match_default_given():

    related = core.RelatedPostings()

    related.add(testutil.Posting('Income:Donations', -1, key='1'))

    related.add(testutil.Posting('Income:Donations', -2, key='2'))

    assert related.meta_values('key', '') == {'1', '2'}

def test_meta_values_many_types():

    expected = {

        datetime.date(2020, 4, 1),

        Decimal(42),

        testutil.Amount(5),

        'rt:42',

    }

    related = core.RelatedPostings()

    for index, value in enumerate(expected):

        related.add(testutil.Posting('Income:Donations', -index, key=value))

    assert related.meta_values('key') == expected