#!/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` for the output half,
lane E1 `findings/subsystems/income-term.md` for the money half, both read off the
instruction stream):

    total = ComputeTotalOutput(SRoh = 0)                     # max mods -> trade rate 1, rest 0
    trade = roundHalfEven(total)
    money = trunc( TradePointsToMoney(trade) )
    maxIncome_s = max(money, 0)

with

    TradePointsToMoney(trade) =
          diffMod * ( f32(IncMod) * ( f32(speciesIncomeFactor)
              * ( Slaves + (PopIncome(1) + (((trade - trade mod 5) * 5 + 0) + PopIncome(0))) ) ) )
        - f32(speciesCostFactor) * ( CalcSuitMod * 10000 * 1.5 )

    diffMod = f32( f32(DifficultyMods(owner)[1]) * f32(server.IncMod) )

Population output per head is `typeOutputMod * 1.8 / 500000`; population INCOME per head is
`typeIncomeModifier / 14000`, with no 1.8 -- two different laws off two adjacent columns of the
same three-row table, and the income one truncates twice (once inside GroupIncome, once after
the morale/addiction product) PER SPECIES.

The difficulty table is built in code from .rdata float literals (0x005a3870); every corpus save
carries `aidf == 1`, whose AI income modifier is 1.1f.  Which players count as AI is NOT on the
wire -- see --ai-rule.
"""
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        # 0x009f8d60
RES_FACTOR = struct.unpack("<d", bytes.fromhex("cdccccccccccec3f"))[0]   # 0.9
BLOCK = 5.0                     # 0x009e2398, used as BOTH the modulus and the multiplier
SUIT_COST_A = 10000.0           # 0x009e9398
SUIT_COST_B = 1.5               # 0x009e90b8
UNOWNED_SUIT_MOD = 20.0         # 0x009e2c08
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], 2: None}

# The difficulty table, built in code by 0x005a3870 from .rdata float literals.
# row -> (aiMaintDiv, aiIncome, aiTrade, plMaintDiv, plIncome, plTrade)
DIFFICULTY_TABLE = {
    0: (1.0, 1.0, 1.0, 1.5, 1.5, 1.5),
    1: (3.0, 1.100000023841858, 1.5, 1.0, 1.0, 1.0),
    2: (1000000.0, 1.7000000476837158, 2.0, 1.0, 1.0, 1.0),
}
DIFFICULTY_DEFAULT = (1.0, 1.0, 1.0, 1.0, 1.0, 1.0)   # LoadDifficultyRow's memcpy'd default

# ---- values the data files supply, measured live on VM140 (2026-09-08, lane N) ---------------
# SpeciesDef +0x4c (base resource demand) and +0x50 (resource output factor) are PER SPECIES and
# are read off the SYSTEM OWNER's species, not the system's population species.  Taking them as
# one global pair is what cost lane N's predictor every Zuul record: a Zuul colony's harvest term
# is min(resAvail, 10) * 40 = 400 output points that a 0/10 pair scores as zero, and 400 points
# is 400/5 whole blocks worth 5 money each, i.e. exactly the 2000-per-colony shortfall observed.
SPECIES_BASE_DEMAND = {5: 10}          # SpeciesDef +0x4c; 0 for Human and Tarkas
SPECIES_RES_OUTPUT = {5: 40.0}         # SpeciesDef +0x50; 10 for Human and Tarkas
SPECIES_BASE_DEMAND_DEFAULT = 0
SPECIES_RES_OUTPUT_DEFAULT = 10.0
# SpeciesDef +0x18 income factor / +0x24 cost factor, per formula-gaps.md Q3 (NOT measured live).
SPECIES_INCOME_FACTOR = {5: 1.1, 6: 0.8}
SPECIES_COST_FACTOR = {5: 0.7}


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


def ftol(v):
    """_ftol2 -- truncation toward zero."""
    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:
        mo = morale_output_mod(morale, tuning, owned, independent)
    v = POPTYPE_OUT[group] * (sf * OUT_FACTOR) * mo * q
    return v if v > 0 else 0.0


def morale_output_mod(m, tuning, owned, independent):
    """0x00746910 -- 1.0 with no owner, on an independent system, or with cm == 0."""
    if not owned or independent or 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),
    "SLAVES_INCOME_MOD": f32(3.0),
    "ADDICTION_INCOME_MOD": f32(0.9),
}


# ---- 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 sparse_table(node, name, key_tag, val_tag):
    """`cm`/`nadct` on the wire: a count then (index, value) pairs."""
    out = {}
    holder = sub(node, name) if name else node
    if holder is None:
        return out
    pending = None
    for c in kids(holder):
        if c.name == key_tag:
            pending = c.value
        elif c.name == val_tag and pending is not None:
            out[pending] = c.value
            pending = None
    return out


def morale_table(node):
    return sparse_table(node, "cm", "msp", "mv")


def addiction_table(node):
    """The int[7] at ServerSystem+0x1e4: `nadct` then sparse (`ads`,`adt`) pairs."""
    return sparse_table(node, None, "ads", "adt")


def read_state(path):
    r = sr.read_save(path)
    systems, players, sim = [], [], None
    for _, n in walk(r.tree):
        if n.name == "Sim" and sim is None:
            sim = n
        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, sim


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"))


# ---- the income chain ------------------------------------------------------------------------
def group_income(t, count, tuning):
    """0x00535e80: _ftol2( f32(POPTYPE[t].income) * (count / 14000) )."""
    inc = POPTYPE_INC[t]
    if inc is None:
        inc = tuning["SLAVES_INCOME_MOD"]
    return ftol(f32(inc) * (float(count) / INCOME_DIVISOR))


def pop_income(t, counts, morale, addicted, tuning, owned, independent):
    """0x0074d760 / 0x0074b700.  `counts` is species -> population for this group type."""
    total = 0.0
    for sp in range(7):
        n = counts.get(sp, 0)
        if not n > 0:
            continue
        mo = morale_output_mod(morale.get(sp, 0), tuning, owned, independent) if t == 1 else 1.0
        ad = tuning["ADDICTION_INCOME_MOD"] if addicted.get(sp) else 1.0
        r = group_income(t, n, tuning)
        total += float(ftol(float(r) * mo * ad))
    return total


def calc_suit_mod(node, p, species, ideal_suit):
    """0x007484d0."""
    if field(node, "vnh"):
        return 0.0
    if not field(node, "PID"):
        return UNOWNED_SUIT_MOD
    if field(p, "RebAI"):
        return 0.0
    ideal = ideal_suit.get(species)
    if ideal is None:
        raise KeyError("no IdealSuit known for species %r" % (species,))
    tol = f32(field(p, "SuitTol") or 0.0)
    d = abs(f32(ideal) - f32(field(node, "Suit") or 0.0))
    return d if d <= tol else tol


def difficulty_income_mod(p, is_ai, server_inc_mod):
    """0x0080f470 + 0x0059b490 + 0x005a3990: float32 throughout."""
    m = f32(server_inc_mod)
    if p is None:
        return m
    level = field(p, "aidf")
    row = DIFFICULTY_TABLE.get(level, DIFFICULTY_DEFAULT)
    # Select(): the AI triple is f[0..2], the non-AI triple f[3..5]; ->+4 is index 1 of the triple.
    mods = row[0:3] if (is_ai and not field(p, "NPC")) else row[3:6]
    return f32(f32(mods[1]) * m)


def system_income(s, p, base_demand, res_output, ideal_suit, server_inc_mod, is_ai, tuning):
    owner_species = field(p, "Species")
    if base_demand is None:
        base_demand = SPECIES_BASE_DEMAND.get(owner_species, SPECIES_BASE_DEMAND_DEFAULT)
    if res_output is None:
        res_output = SPECIES_RES_OUTPUT.get(owner_species, SPECIES_RES_OUTPUT_DEFAULT)
    sp = system_species(s, owner_species)
    strip = bool(field(p, "AMine"))
    independent = is_independent(s)

    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)
    morale = morale_table(s)
    addicted = addiction_table(s)

    # ---- the output half (lane N, live-verified) --------------------------------------------
    if field(s, "rbfl"):
        total = 0.0
    else:
        civ_own = population(s, "Pop2", 1, sp) + population(s, "pbon2", 1, sp)
        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, tuning, True, independent)
        civilian = group_output(1, civ_own, morale.get(sp, 0), 0, tuning, True, independent)
        total = (civilian + (imperial + 0.0)) + (harvest + resource)
        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)

    # ---- the money half (lane E1) -------------------------------------------------------------
    trade = round_half_even(total)

    imperial_counts = {sp: pop} if pop > 0 else {}
    civ_counts = {}
    for q in range(7):
        c = population(s, "Pop2", 1, q) + population(s, "pbon2", 1, q)
        if c:
            civ_counts[q] = c

    t = pop_income(0, imperial_counts, morale, addicted, tuning, True, independent)
    t = ((trade - math.fmod(trade, BLOCK)) * BLOCK + 0.0) + t
    t = pop_income(1, civ_counts, morale, addicted, tuning, True, independent) + t
    t = 0.0 + t                                     # SlaveIncome(): no slaves in this corpus

    t = f32(SPECIES_INCOME_FACTOR.get(owner_species, 1.0)) * t
    t = difficulty_income_mod(p, is_ai, server_inc_mod) * (f32(field(p, "IncMod") or 1.0) * t)

    cost = f32(SPECIES_COST_FACTOR.get(owner_species, 1.0)) * (
        calc_suit_mod(s, p, sp, ideal_suit) * SUIT_COST_A * SUIT_COST_B)
    return ftol(t - cost)


AI_RULES = {
    # `ServerPlayer+0xf9` is a setup input, not a save field (income-term.md §3.1).
    "non-npc-not-first": lambda p, i: not field(p, "NPC") and (field(p, "PlyrIdx") or 0) != 0,
    "none": lambda p, i: False,
    "all-non-npc": lambda p, i: not field(p, "NPC"),
}


def predict(save, base_demand, res_output, ai_rule, 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, sim = read_state(save)
    server_inc_mod = field(sim, "IncMod", 1.0) if sim is not None else 1.0

    # `server->IdealSuit[species]` IS on the wire: the Sim block's `ISsp`/`ISsu` pairs are that
    # float[7], in species-index order, and they are randomised per game by the map generator.
    # Each ServerPlayer's own `IdealSuit` field carries the same value for its species, so the
    # two are cross-checked here rather than one being trusted blindly (rule 8: two checks that
    # share a hidden assumption are one check -- these two do not share a source).
    ideal_suit = {}
    if sim is not None:
        idx = 0
        for c in kids(sim):
            if c.name == "ISsu":
                ideal_suit[idx] = c.value
                idx += 1
    for p in players:
        s, v = field(p, "Species"), field(p, "IdealSuit")
        if s is None or v is None:
            continue
        if s in ideal_suit and ideal_suit[s] != v:
            print("   WARNING: species %d ISsu %r disagrees with a player's IdealSuit %r"
                  % (s, ideal_suit[s], v), file=sys.stderr)
        ideal_suit.setdefault(s, v)

    # 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 i, p in enumerate(players):
        h = handle_of.get(id(p))
        is_ai = AI_RULES[ai_rule](p, i)
        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, ideal_suit,
                                  server_inc_mod, is_ai, LIVE_TUNING)
                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 sp=%s ai=%s npc=%s (%d system(s))"
                  % (h, field(p, "Species"), is_ai, field(p, "NPC"), len(detail)))
            for sid, nm, m in detail:
                print("      sys %-4s %-16s money=%d" % (sid, nm, m))
    return out


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=None,
                    help="override SpeciesDef +0x4c for every species")
    ap.add_argument("--res-output", type=float, default=None,
                    help="override SpeciesDef +0x50 for every species")
    ap.add_argument("--ai-rule", choices=sorted(AI_RULES), default="non-npc-not-first",
                    help="which players count as AI (ServerPlayer+0xf9 is not on the wire)")
    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)
        if isinstance(raw, dict) and raw and isinstance(next(iter(raw.values())), list):
            for name, rows in raw.items():
                for row in rows:
                    oracle[(os.path.basename(name), row.get("BnkEl"))] = row.get("maxIncome")
        else:
            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.ai_rule, 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  (%.6f 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())