"""Shared summation rules for the budget model (used by the parser for residuals and by the verifier).

A printed number in the law is reproduced by summing Allocation records that match its selector:
fund, flow, year, nature, untilEnd, purpose subtree, block, commitment kind, and classification codes.
Consolidation: a node's total includes an intra-fund transfer only if that node itself prints the transfer's code;
otherwise intra-fund transfers are eliminated (they cancel out inside the node).
"""
import collections


def kind_stem(k):
    """Commitment-kind codes are hierarchical in 2-digit groups: 0101100000 -> 010110."""
    s = k or ''
    while len(s) > 2 and s.endswith('00'):
        s = s[:-2]
    return s


class Calc:
    def __init__(self, allocs, purpose_parent, cls_parent, intra_codes, blockcodes, purpose_holder=None):
        self.purpose_parent = purpose_parent
        self.purpose_holder = purpose_holder or {}
        self.cls_parent = cls_parent
        self.intra = intra_codes
        self.blockcodes = {k: {f: set(v) for f, v in d.items()} for k, d in blockcodes.items()}
        self.idx = collections.defaultdict(list)
        for a in allocs:
            self.add(a)

    def chain(self, p):
        out = set()
        while p and p not in out:
            out.add(p)
            p = self.purpose_parent.get(p)
        return out

    def code_chain(self, k):
        out = set()
        while k and k not in out:
            out.add(k)
            k = self.cls_parent.get(k)
        return out

    def add(self, a):
        a['_pchain'] = self.chain(a.get('purpose'))
        a['_ck'] = f"{a['scheme']}:{a['code']}"
        a['_cchain'] = self.code_chain(a['_ck'])
        a['_intra'] = bool(a['_cchain'] & self.intra)
        self.idx[(a['fund'], a['flow'], int(a['year']), a['nature'])].append(a)

    def total(self, fund, flow, year, nature, purpose=None, block=None, codes=None, ckind=None, untilEnd=False,
              blk=None, ekk_prefix=None, consolidated=None, holder_view=False):
        """holder_view: also count budget-level amounts (no purpose) whose holder administers the selected purpose."""
        printed = self.blockcodes.get(blk, {}).get(flow) if blk else None
        if consolidated is None:
            consolidated = purpose is None
        codes = set(codes) if codes else None
        s = 0.0
        for f in ([fund] if fund else ['basic', 'special']):
            for a in self.idx.get((f, flow, int(year), nature), ()):
                if a.get('partOf'):
                    continue
                if bool(a.get('untilEnd')) != bool(untilEnd):
                    continue
                if purpose and purpose not in a['_pchain']:
                    if not (holder_view and not a.get('purpose') and a.get('holder') and a.get('holder') == self.purpose_holder.get(purpose)):
                        continue
                if block and a.get('block') != block:
                    continue
                if ckind and not (a.get('commitmentKind') or '').startswith(kind_stem(ckind)):
                    continue
                if ekk_prefix and not (a['scheme'] == 'ekk' and a['code'].startswith(ekk_prefix)):
                    continue
                ck, cc = a['_ck'], a['_cchain']
                if printed is not None:
                    shown = ck in printed
                    if not shown and a['_intra']:
                        continue
                    if codes is not None:
                        if not (ck in codes or (not shown and cc & codes)):
                            continue
                elif codes is not None:
                    if not (cc & codes):
                        continue
                    if consolidated and a['_intra'] and ck not in codes:
                        continue
                elif consolidated and a['_intra']:
                    continue
                s += float(a['amount'])
        return s