#!/usr/bin/env python3
"""Recover each player's maximum-income sum from the bankruptcy limit a save carries.

`ServerPlayer::UpdateBankruptcyLimits` stores

    BnkEl = max( trunc(maxIncome / -0.15000000596046448), -2000000000 )
    BnkPr = max( -trunc(BANKRUPTCY_PROTECTION_LIMIT_FACTOR * maxIncome), BnkEl )

where `maxIncome` is the sum, over the player's owned systems, of `max(ComputeMaxIncome(s), 0)` --
the per-system money output that blocks `ComputeBudget`, both savings phases and half the turn
record.  The first map has a slope of about 6.67, so it is injective: each stored `BnkEl` has at
most one integer preimage, and every save in the corpus therefore STATES the blocked term for
every player that owns anything.

That makes this a per-save oracle for a formula nobody has yet written: sum a candidate
per-system output over the owned systems and compare it with the number here.

Two caveats the tool reports rather than hides:
  * a player whose limit hit the -2,000,000,000 floor inverts to nothing, and is listed as such
    rather than given a wrong answer;
  * the divisor is the round-tripped float `-0.15000000596046448`, NOT `-0.15`.  Using `-0.15`
    changes the result for every `maxIncome` divisible by 3, and the tool flags the records where
    the two constants disagree.

  tools/max_income_oracle.py                       # every save in verify/results/saves
  tools/max_income_oracle.py PATH.sav ...          # named saves
  tools/max_income_oracle.py --json OUT
"""
import argparse
import glob
import json
import math
import os
import sys

RE_ROOT = os.path.abspath(os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))
SAVE_DIR = os.path.join(RE_ROOT, "verify", "results", "saves")
sys.path.insert(0, os.path.join(RE_ROOT, "verify", "save-reader"))
import save_reader as sr  # noqa: E402

DIVISOR = -0.15000000596046448   # the double in .rdata: (double)(float)-0.15f
NAIVE = -0.15                    # what a reimplementation writes by mistake
FLOOR = -2000000000


def kids(n):
    return getattr(n, "children", []) or []


def name(n):
    return getattr(n, "name", None)


def field(n, key):
    for c in kids(n):
        if name(c) == key:
            return getattr(c, "value", None)
    return None


def invert(bnkel):
    """Integer maxIncome values that map to this stored limit, [] when none do."""
    if bnkel == 0:
        return [0]
    if bnkel <= FLOOR:
        return []          # clamped: the preimage is unbounded, so refuse to answer
    centre = int(-bnkel * 0.15)
    return [m for m in range(max(centre - 8, 0), centre + 9)
            if math.trunc(m / DIVISOR) == bnkel]


def scan(path):
    res = sr.read_save(path)
    sim = [c for c in kids(res.tree) if name(c) == "Sim"][0]
    ch = kids(sim)
    rows = []
    for i, c in enumerate(ch):
        if name(c) != "Player":
            continue
        pid = getattr(ch[i - 1], "value", None)
        el, pr = field(c, "BnkEl"), field(c, "BnkPr")
        owned = sum(1 for k in kids(c) if name(k) == "OwnId")
        if not el:
            rows.append({"playerID": pid, "ownedSystems": owned, "BnkEl": el, "BnkPr": pr,
                         "maxIncome": None, "why": "limit is zero (no owned systems)"})
            continue
        cands = invert(el)
        if len(cands) != 1:
            rows.append({"playerID": pid, "ownedSystems": owned, "BnkEl": el, "BnkPr": pr,
                         "maxIncome": None,
                         "why": "clamped at the floor" if el <= FLOOR
                                else "%d preimages" % len(cands)})
            continue
        m = cands[0]
        rows.append({
            "playerID": pid,
            "ownedSystems": owned,
            "BnkEl": el,
            "BnkPr": pr,
            "maxIncome": m,
            "protectionFactor": (-pr / m) if m else None,
            "naiveDivisorDisagrees": math.trunc(m / NAIVE) != math.trunc(m / DIVISOR),
        })
    return rows


def main(argv=None):
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("saves", nargs="*")
    ap.add_argument("--json")
    args = ap.parse_args(argv)

    paths = args.saves or sorted(glob.glob(os.path.join(SAVE_DIR, "*.sav")))
    if not paths:
        print("no saves given and none in " + SAVE_DIR, file=sys.stderr)
        return 2

    out = {}
    factors, disagree, solved = [], 0, 0
    for p in paths:
        rows = scan(p)
        out[os.path.basename(p)] = rows
        print("== " + os.path.basename(p))
        for r in rows:
            if r["maxIncome"] is None:
                print("   pid=%-5s owns %-3d  BnkEl=%-12s  -- %s"
                      % (r["playerID"], r["ownedSystems"], r["BnkEl"], r["why"]))
                continue
            solved += 1
            factors.append(r["protectionFactor"])
            disagree += 1 if r["naiveDivisorDisagrees"] else 0
            print("   pid=%-5s owns %-3d  BnkEl=%-12d  maxIncome=%-10d  factor=%.9f%s"
                  % (r["playerID"], r["ownedSystems"], r["BnkEl"], r["maxIncome"],
                     r["protectionFactor"],
                     "   [-0.15 would differ]" if r["naiveDivisorDisagrees"] else ""))
    if factors:
        print("\n%d player-record(s) inverted; protection factor in [%.9f, %.9f]"
              % (solved, min(factors), max(factors)))
        print("%d of them would get a DIFFERENT stored limit from the naive -0.15 divisor"
              % disagree)
    if args.json:
        with open(args.json, "w") as f:
            json.dump(out, f, indent=2)
    return 0


if __name__ == "__main__":
    sys.exit(main())