#!/usr/bin/env python3
"""Predict each player's maxIncome from a save and check it against the bankruptcy oracle.

`tools/max_income_oracle.py` (lane Y) inverts the stored `BnkEl` to recover

    maxIncome = sum over owned systems of max(ComputeMaxIncome(s), 0)

which is the per-system money output that blocks `ComputeBudget` and four other phases.
This script goes the other way: it computes the same number from the colony state and
compares, so the formula is falsified per player-record rather than per lane.

The chain it implements (lane N, findings/subsystems/output-term.md, live-verified on VM140):

    total = ComputeTotalOutput(SRoh = 0)                     # max mods -> trade rate 1, rest 0
          = ( overHarvestDemand x speciesResourceOutput
            + (TRes + resAvail) x stripMineFraction x 0.9
            + populationOutput )
            x OutMod x sysOutMod x setupOutMod x RebOutMod x ScOutMod
    trade = roundHalfEven(total)
    money = ftol( TradePointsToMoney(trade) )

Population output per head is `typeOutputMod x 1.8 / 500000`; population INCOME per head is
`typeIncomeModifier / 14000`, with no 1.8 -- two different laws off the same table.

Everything the executable carries is hard-coded here as such.  The two per-species fields the
data files supply (SpeciesDef +0x4c base resource demand, +0x50 resource output factor) are
options, defaulting to the values measured live on VM140 for the species this corpus contains.
"""
import argparse
import json
import math
import os
import struct
import sys

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

# ---- constants the executable carries itself ------------------------------------------------
OUT_FACTOR = struct.unpack("<d", bytes.fromhex("ccccccccccccfc3f"))[0]   # 1.8, as the image holds it
OUT_DIVISOR = 500000.0
INCOME_DIVISOR = 14000.0
RES_FACTOR = struct.unpack("<d", bytes.fromhex("cdccccccccccec3f"))[0]   # 0.9
POPTYPE_OUT = {0: 1.0, 1: struct.unpack("<f", struct.pack("<f", 0.33))[0]}
POPTYPE_INC = {0: 1.0, 1: struct.unpack("<f", struct.pack("<f", 0.33))[0]}

# ---- values the data files supply, measured live on VM140 (2026-09-08) ----------------------
SPECIES_BASE_DEMAND = 0        # SpeciesDef +0x4c
SPECIES_RES_OUTPUT = 10.0      # SpeciesDef +0x50


def f32(v):
    return struct.unpack("<f", struct.pack("<f", v))[0]


def ftol(v):
    return math.trunc(v)


def round_half_even(v):
    # `fistp`/`fild` with the x87 in round-to-nearest: ties go to the even neighbour.
    f = math.floor(v)
    d = v - f
    if d > 0.5:
        return f + 1.0
    if d < 0.5:
        return float(f)
    return float(f if f % 2 == 0 else f + 1)


def clamp01(v):
    return 0.0 if v < 0 else (1.0 if v > 1 else v)


def signed_cbrt(x):
    return math.pow(x, 1.0 / 3.0) if x >= 0 else -math.pow(-x, 1.0 / 3.0)


def strip_mine_fraction(pop, infra, ibon):
    fi = f32(ibon + infra)
    r = clamp01(signed_cbrt(pop / 100.0) * 0.01)
    if 0.0001 + r >= 1.0:
        r = fi
    return f32(r if r < fi else fi)


def over_harvest_demand(rate, avail, pop, base_demand):
    b = 0.0
    if rate > 0.0:
        v = rate * float(avail) * clamp01(float(pop) * 1e-05)
        b = v if v > 1.0 else 1.0
    t = float(base_demand) + b
    lo = t if t > 0.0 else 0.0
    return float(avail) if float(avail) < lo else lo


def group_output(group, count, morale, stations, tuning, owned, independent):
    if not count > 0:
        return 0.0
    q = float(count) / OUT_DIVISOR
    sf = 1.0
    if owned and group == 0:
        b = tuning["STATION_BONUS_IMPERIAL_OUTPUT"]
        sf = 1.0 + stations * (b if b > 0 else 0.0)
    mo = 1.0
    if group == 1 and owned and not independent:
        mo = morale_output_mod(morale, tuning)
    v = POPTYPE_OUT[group] * (sf * OUT_FACTOR) * mo * q
    return v if v > 0 else 0.0


def morale_output_mod(m, tuning):
    if m == 0:
        return 1.0
    if m >= tuning["MORALE_INCREASE_OUTPUT"]:
        x = tuning["MORALE_INCREASE_OUTPUT_MOD"]
        return x if x > 0 else 1.0
    if m <= tuning["MORALE_DECREASE_OUTPUT"]:
        x = tuning["MORALE_DECREASE_OUTPUT_MOD"]
        return x if x > 0 else 1.0
    return 1.0


# The tuning values this corpus runs with, read out of the live process on VM140.
LIVE_TUNING = {
    "STATION_BONUS_IMPERIAL_OUTPUT": f32(0.1),
    "MORALE_INCREASE_OUTPUT": 85,
    "MORALE_INCREASE_OUTPUT_MOD": f32(1.5),
    "MORALE_DECREASE_OUTPUT": 20,
    "MORALE_DECREASE_OUTPUT_MOD": f32(0.5),
    "SLAVES_OUTPUT_MOD": f32(3.0),
}


# ---- save reading ---------------------------------------------------------------------------
def kids(n):
    return n.children or []


def walk(node, path=""):
    yield path, node
    for c in kids(node):
        yield from walk(c, path + "/" + (c.name or "?"))


def field(node, name, default=None):
    for c in kids(node):
        if c.name == name:
            return c.value
    return default


def sub(node, name):
    for c in kids(node):
        if c.name == name:
            return c
    return None


def population(node, name, group, species):
    p = sub(node, name)
    if p is None:
        return 0
    total = 0
    for g in kids(p):
        d = {c.name: c.value for c in kids(g)}
        if d.get("PopT") == group and d.get("PopS") == species:
            total += d.get("PopC") or 0
    return total


def morale_table(node):
    """`cm` on the wire is a count then (species, value) pairs."""
    cm = sub(node, "cm")
    out = {}
    if cm is None:
        return out
    pending = None
    for c in kids(cm):
        if c.name == "msp":
            pending = c.value
        elif c.name == "mv" and pending is not None:
            out[pending] = c.value
            pending = None
    return out


def read_state(path):
    r = sr.read_save(path)
    systems, players = [], []
    for _, n in walk(r.tree):
        names = {c.name for c in kids(n)}
        if {"Pop", "Rts", "Infra", "pbon"} <= names:
            systems.append(n)
        elif {"OutMod", "IncMod", "ScOutMod", "BnkEl", "PlyrIdx"} <= names:
            players.append(n)
    return systems, players


def system_species(node, owner_species):
    """The species the system's population is credited to.

    `indi` is written unconditionally and reads back as a zero-filled record on an ordinary
    colony; the gate is the separate `hindi` bool. Reading `indsp` without checking `hindi`
    silently credits every colony to species 0 -- which agrees with the oracle on a HUMAN
    empire and disagrees on every other, the exact failure mode rule 8 warns about.
    """
    if field(node, "hindi"):
        indi = sub(node, "indi")
        if indi is not None:
            s = field(indi, "indsp")
            if s is not None:
                return s
    return owner_species


def is_independent(node):
    return bool(field(node, "hindi"))


def predict(save, base_demand, res_output, verbose=False):
    """Return {BnkEl: (predicted maxIncome, per-system detail)}.

    The join key is `BnkEl` rather than any id: a system stores its owner as a HANDLE id
    (`PID`), a player stores an ordinal (`PlyrIdx`), and the two are different numbering
    schemes.  The oracle inverts the same `BnkEl` the player node carries, so keying on it
    needs no id mapping at all and cannot silently pair the wrong two records.
    """
    systems, players = read_state(save)
    # A player's handle id is not a named field, but every system names its owner's, so the
    # set of distinct non-zero `PID` values is the set of owning players -- in the same order
    # the player records appear.  Pair them by position among the players that own anything.
    owner_handles = []
    for s in systems:
        h = field(s, "PID")
        if h and h not in owner_handles:
            owner_handles.append(h)
    owner_handles.sort()

    owners = [p for p in players if (field(p, "NumOwn") or 0) > 0]
    handle_of = {}
    if len(owners) == len(owner_handles):
        for p, h in zip(owners, owner_handles):
            handle_of[id(p)] = h

    out = {}
    for p in players:
        h = handle_of.get(id(p))
        total = 0
        detail = []
        if h is not None:
            for s in systems:
                if field(s, "PID") != h:
                    continue
                m = system_income(s, p, base_demand, res_output)
                detail.append((field(s, "Idx"), field(s, "Name"), m))
                total += max(m, 0)
        out[field(p, "BnkEl")] = (total, detail)
        if verbose and detail:
            print("   player handle %s (%d system(s))" % (h, len(detail)))
            for sid, nm, m in detail:
                print("      sys %-4s %-16s money=%d" % (sid, nm, m))
    return out


def system_income(s, p, base_demand, res_output):
    rts = sub(s, "Rts")
    owner_species = field(p, "Species")
    sp = system_species(s, owner_species)
    strip = bool(field(p, "AMine"))

    res = field(s, "Res") or 0
    avail = res + ((field(s, "MRes") or 0) + (field(s, "ARes2") or 0) if strip else 0)
    pop = (field(s, "Pop") or 0) + (field(s, "pbon") or 0)
    civ = population(s, "Pop2", 1, sp) + population(s, "pbon2", 1, sp)
    morale = morale_table(s).get(sp, 0)
    independent = is_independent(s)

    harvest = over_harvest_demand(0.0, avail, pop, base_demand) * f32(res_output)
    resource = (float((field(s, "TRes") or 0) + avail)
                * float(strip_mine_fraction(pop, field(s, "Infra") or 0.0, field(s, "ibon") or 0.0))
                * RES_FACTOR)
    imperial = group_output(0, pop, 0, 0, LIVE_TUNING, True, independent)
    civilian = group_output(1, civ, morale, 0, LIVE_TUNING, True, independent)
    base = (civilian + (imperial + 0.0)) + (harvest + resource)

    if field(s, "rbfl"):
        return 0
    total = base
    total *= f32(field(p, "OutMod") or 1.0)
    total *= f32(field(s, "OutMod") or 1.0)
    total *= f32(field(p, "RebOutMod") or 1.0)
    total *= f32(field(p, "ScOutMod") or 1.0)

    trade = round_half_even(total)
    t = (trade - math.fmod(trade, 5.0)) * 5.0
    t += ftol(POPTYPE_INC[0] * (float(pop) / INCOME_DIVISOR))
    t += ftol(POPTYPE_INC[1] * (float(civ) / INCOME_DIVISOR))
    t *= f32(field(p, "IncMod") or 1.0)
    return ftol(t)


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("saves", nargs="*")
    ap.add_argument("--oracle", help="JSON from max_income_oracle.py --json")
    ap.add_argument("--base-demand", type=int, default=SPECIES_BASE_DEMAND)
    ap.add_argument("--res-output", type=float, default=SPECIES_RES_OUTPUT)
    ap.add_argument("-v", "--verbose", action="store_true")
    a = ap.parse_args()

    oracle = {}
    if a.oracle:
        with open(a.oracle) as fh:
            raw = json.load(fh)
        for row in (raw if isinstance(raw, list) else raw.get("records", [])):
            oracle[(os.path.basename(row.get("save", "")), row.get("BnkEl"))] = row.get("maxIncome")

    hits = misses = unknown = 0
    for save in a.saves:
        name = os.path.basename(save)
        print("==", name)
        for bnkel, (total, detail) in sorted(predict(save, a.base_demand, a.res_output,
                                                     a.verbose).items()):
            if not detail:
                continue
            want = oracle.get((name, bnkel))
            if want is None:
                unknown += 1
                print("   BnkEl=%-12s predicted=%-12d (no oracle record)" % (bnkel, total))
            elif want == total:
                hits += 1
                print("   BnkEl=%-12s predicted=%-12d MATCH" % (bnkel, total))
            else:
                misses += 1
                print("   BnkEl=%-12s predicted=%-12d oracle=%-12d delta=%+d  (%.4f x)"
                      % (bnkel, total, want, total - want,
                         (total / want) if want else float("nan")))
    print("\n%d match, %d differ, %d with no oracle record" % (hits, misses, unknown))
    return 1 if misses else 0


if __name__ == "__main__":
    raise SystemExit(main())