#!/usr/bin/env python3
"""Independently check water-handling histories and contradiction multipliers.

Uses Python Fraction and direct inventory simulation. Constraint coefficients
come from exact unit-basis evaluations, not the MoonBit affine compiler. Small
studies are independently decided by Gaussian vertex enumeration, not FM.
Only declared-model correctness is checked, never physical/contest originality.
"""
import argparse
from copy import deepcopy
from datetime import datetime, timezone
from fractions import Fraction as F
from itertools import combinations
import hashlib
import json
from pathlib import Path
import random
import re

from aqueous import ROOT, run_moon

FIELDS = ("calcium_nmol_l", "dic_nmol_l", "ta_neq_l", "sodium_nmol_l", "chloride_nmol_l")
ANALYTES = ("Ca", "DIC", "TA", "Na", "Cl")
CHEMICAL = {"co2", "calcite", "hcl", "naoh"}


def exact(value):
    assert set(value) == {"numerator", "denominator"}
    n, d = value["numerator"], value["denominator"]
    assert isinstance(n, str) and isinstance(d, str)
    assert re.fullmatch(r"-?[0-9]+", n) and re.fullmatch(r"[0-9]+", d)
    assert len(n) <= 700 and len(d) <= 700 and int(d) > 0
    x = F(int(n), int(d))
    assert str(x.numerator) == n and str(x.denominator) == d, "noncanonical output"
    return x


def encoded(value):
    x = F(value)
    return {"numerator": str(x.numerator), "denominator": str(x.denominator)}


def parameters(study):
    result = []
    for a, field in zip(ANALYTES, FIELDS):
        b = study["source"][field]
        if b["lower"] != b["upper"]:
            result.append(("source:" + a, b))
    for s in study["steps"]:
        b = s["amount_nmol"]
        if b["lower"] != b["upper"]:
            result.append(("step:" + s["id"], b))
    assert len(result) <= 10
    return result


def input_value(name, bounds, choices):
    return F(bounds["lower"]) if bounds["lower"] == bounds["upper"] else choices[name]


def simulate(study, choices):
    """Numeric inventory arithmetic only, preserving one value per shared dose."""
    source = study["source"]
    volume = source["volume_ul"]
    assert type(volume) is int and 0 < volume <= 10**9
    c = [input_value("source:" + a, source[f], choices) for a, f in zip(ANALYTES, FIELDS)]
    bottles = {source["id"]: (volume, [x * F(volume, 10**6) for x in c])}
    for s in study["steps"]:
        assert s["id"] not in bottles
        parent_volume, parent_amounts = bottles[s["parent"]]
        v, n = parent_volume, parent_amounts.copy()
        kind, dv = s["kind"], s["volume_ul"]
        dose = input_value("step:" + s["id"], s["amount_nmol"], choices)
        if kind == "take":
            assert 0 < dv <= v
            n = [x * F(dv, v) for x in n]
            v = dv
        elif kind in ("dilute", "hcl", "naoh"):
            v += dv
        elif kind == "evaporate":
            v -= dv
        if kind == "co2":
            n[1] += dose
        elif kind == "calcite":
            n[0] += dose
            n[1] += dose
            n[2] += 2 * dose
        elif kind == "hcl":
            n[2] -= dose
            n[4] += dose
        elif kind == "naoh":
            n[2] += dose
            n[3] += dose
        else:
            assert kind in ("take", "dilute", "evaporate")
        assert 0 < v <= 10**9
        bottles[s["id"]] = (v, n)
    # Actual allocation is checked independently of any parameter values.
    for name, (v, _) in bottles.items():
        children = [s for s in study["steps"] if s["parent"] == name]
        if len(children) > 1:
            assert all(s["kind"] == "take" for s in children)
        assert sum(s["volume_ul"] for s in children if s["kind"] == "take") <= v
    return bottles


def residuals(study, choices, zeros=()):
    """Every returned scalar must be <= 0, with stable chemical input IDs."""
    values = {}
    for name, b in parameters(study):
        values["bound:" + name + ":upper"] = choices[name] - b["upper"]
        values["bound:" + name + ":lower"] = b["lower"] - choices[name]
    bottles = simulate(study, choices)
    for name, (_, n) in bottles.items():
        for index in (0, 1, 3, 4):
            values[f"nonnegative:{name}:{ANALYTES[index]}"] = -n[index]
    n = bottles[study["source"]["id"]][1]
    closure = 2*n[0] + n[3] - n[4] - n[2]
    values["source:carbonate-charge:upper"] = closure
    values["source:carbonate-charge:lower"] = -closure
    for a in study["assays"]:
        v, n = bottles[a["node"]]
        concentration = n[ANALYTES.index(a["analyte"])] * F(10**6, v)
        values["assay:" + a["id"] + ":upper"] = concentration - a["bounds"]["upper"]
        values["assay:" + a["id"] + ":lower"] = a["bounds"]["lower"] - concentration
    assert len(set(zeros)) == len(zeros)
    for name in zeros:
        s = next(s for s in study["steps"] if s["id"] == name)
        assert s["kind"] in CHEMICAL
        dose = input_value("step:" + name, s["amount_nmol"], choices)
        values["claim:zero:" + name + ":upper"] = dose
        values["claim:zero:" + name + ":lower"] = -dose
    return values


def rows_for(study, zeros=()):
    names = [name for name, _ in parameters(study)]
    origin = {name: F(0) for name in names}
    constant = residuals(study, origin, zeros)
    basis = []
    for name in names:
        point = dict(origin)
        point[name] = F(1)
        basis.append(residuals(study, point, zeros))
    return {key: ([b[key] - value for b in basis], -value) for key, value in constant.items()}


def verify_history(study, history, zeros=()):
    required = [name for name, _ in parameters(study)]
    assert [x["parameter"] for x in history["parameters"]] == required
    choices = {x["parameter"]: exact(x["value"]) for x in history["parameters"]}
    assert all(x <= 0 for x in residuals(study, choices, zeros).values()), "history violates input"
    bottles = simulate(study, choices)
    assert [s["id"] for s in history["samples"]] == list(bottles), "wrong/missing bottle identity"
    for s in history["samples"]:
        v, n = bottles[s["id"]]
        assert s["volume_ul"] == v and type(s["volume_ul"]) is int
        for field, amount in zip(FIELDS, n):
            assert exact(s[field]) == amount * F(10**6, v), "wrong inventory/volume/counterion"


def verify_assessment(study, result, zeros=()):
    rows = rows_for(study, zeros)
    dimension = len(parameters(study))
    assert result["parameter_count"] == dimension and result["constraint_count"] == len(rows)
    status = result["status"]
    if status == "Compatible":
        assert result.get("history") is not None and result["conflict"] == []
        verify_history(study, result["history"], zeros)
    elif status == "Conflict":
        assert result.get("history") is None and result["conflict"]
        coefficients, upper, seen = [F(0)] * dimension, F(0), set()
        for term in result["conflict"]:
            key, weight = term["constraint"], exact(term["weight"])
            assert weight >= 0 and key not in seen
            seen.add(key)
            a, b = rows[key]
            coefficients = [x + weight*y for x, y in zip(coefficients, a)]
            upper += weight * b
        assert all(x == 0 for x in coefficients) and upper < 0, "not a contradiction"
    else:
        assert status == "ResourceLimit" and result.get("history") is None and not result["conflict"]


def solve_equalities(matrix, rhs):
    n = len(rhs)
    a = [list(row) + [b] for row, b in zip(matrix, rhs)]
    for col in range(n):
        pivot = next((i for i in range(col, n) if a[i][col]), None)
        if pivot is None:
            return None
        a[col], a[pivot] = a[pivot], a[col]
        divisor = a[col][col]
        a[col] = [x / divisor for x in a[col]]
        for i in range(n):
            if i != col:
                scale = a[i][col]
                a[i] = [x - scale*y for x, y in zip(a[i], a[col])]
    return [row[-1] for row in a]


def vertex_status(study, zeros=()):
    """Bounded <=3-dimensional polytope: a nonempty one has a vertex."""
    dim = len(parameters(study))
    assert dim <= 3, "independent enumeration deliberately limited to small studies"
    pruned = {}
    for a, b in rows_for(study, zeros).values():
        pivot = next((abs(x) for x in a if x), None)
        if pivot is None:
            if b < 0:
                return "Conflict"
            continue
        key, bound = tuple(x/pivot for x in a), b/pivot
        if key not in pruned or bound < pruned[key]:
            pruned[key] = bound
    constraints = list(pruned.items())
    for selection in combinations(constraints, dim):
        point = solve_equalities([a for a, _ in selection], [b for _, b in selection])
        if point is not None and all(sum(x*y for x,y in zip(a,point)) <= b for a,b in constraints):
            return "Compatible"
    return "Conflict"


def verify_claim(request, result):
    study, zeros = request["study"], request["zero_steps"]
    assert result["zero_steps"] == zeros
    verify_assessment(study, result["base"])
    baseline = result["base"]["status"]
    if baseline == "Conflict":
        assert result["status"] == "BaseConflict" and result.get("assumed_zero") is None
    elif baseline == "ResourceLimit":
        assert result["status"] == "ResourceLimit" and result.get("assumed_zero") is None
    else:
        tested = result["assumed_zero"]
        verify_assessment(study, tested, zeros)
        assert result["status"] == {"Compatible": "ZeroTransferPossible", "Conflict": "ZeroTransferRuledOut",
                                    "ResourceLimit": "ResourceLimit"}[tested["status"]]


def run_analysis(command, request, target):
    result = run_moon(command + "-json", request, target, package="cmd/main")
    assert result["ok"] is True, result
    return result["result"]


def verify_small(command, request, result):
    if command == "reconcile":
        verify_assessment(request, result)
        assert result["status"] == vertex_status(request), "disagrees with independent vertices"
    else:
        verify_claim(request, result)
        assert result["base"]["status"] == vertex_status(request["study"])
        if result.get("assumed_zero") is not None:
            assert result["assumed_zero"]["status"] == vertex_status(request["study"], request["zero_steps"])


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--target", choices=["wasm-gc", "js", "native"], default="wasm-gc")
    parser.add_argument("--compare-target", choices=["wasm-gc", "js", "native"])
    parser.add_argument("--output-dir", type=Path, default=ROOT/"audit/differentiation/verification/aliquots")
    args = parser.parse_args()
    out = args.output_dir.resolve()
    out.mkdir(parents=True, exist_ok=True)
    report = {"passed": False, "started_at_utc": datetime.now(timezone.utc).isoformat(),
              "target": args.target, "compare_target": args.compare_target, "cases": [],
              "negative_inputs": [], "rejected_mutations": [],
              "method": "Fraction inventory simulation + unit-basis constraints + Gaussian vertex enumeration <=3 unknowns",
              "limits": "synthetic declared five-inventory model; not experimental/causal or originality verification"}

    def save():
        (out/"report.json").write_text(json.dumps(report, ensure_ascii=False, indent=2)+"\n")

    def check(label, command, request, expected=None, compare=False):
        result = run_analysis(command, request, args.target)
        verify_small(command, request, result)
        if expected:
            assert result["status"] == expected, (label, result["status"], expected)
        if compare and args.compare_target:
            other = run_analysis(command, request, args.compare_target)
            verify_small(command, request, other)
            assert other == result, "cross-target exact output differs"
        (out/(label+".json")).write_text(json.dumps({"command": command, "request": request, "result": result}, indent=2)+"\n")
        report["cases"].append({"label": label, "command": command, "status": result["status"],
                                "cross_target_compared": bool(compare and args.compare_target)})
        save()
        return result

    try:
        expected = {"common-source-conflict": "Conflict", "acidified-split": "Compatible",
                    "masked-acid-base.claim": "ZeroTransferRuledOut", "mineral-acid.claim": "ZeroTransferRuledOut",
                    "underdetermined-gas.claim": "ZeroTransferPossible"}
        main_cases = {}
        for path in sorted((ROOT/"examples/aliquots").glob("*.json")):
            command = "challenge" if ".claim." in path.name else "reconcile"
            request = json.loads(path.read_text())
            main_cases[path.stem] = (request, check(path.stem, command, request, expected[path.stem], True))
        common = main_cases["common-source-conflict"][0]
        for i in range(2):
            marginal = deepcopy(common)
            marginal["assays"] = [marginal["assays"][i]]
            check("common-marginal-"+str(i), "reconcile", marginal, "Compatible")
        acid = main_cases["acidified-split"][0]
        check("shared-acid-required", "challenge", {"study": acid, "zero_steps": ["acid"]}, "ZeroTransferRuledOut")
        for key, ignored, zeros in [("masked-acid-base.claim", {"Na", "Cl"}, ["acid", "base"]),
                                     ("mineral-acid.claim", {"Ca", "Cl"}, ["mineral", "acid"])]:
            req = deepcopy(main_cases[key][0])
            req["study"]["assays"] = [a for a in req["study"]["assays"] if a["analyte"] not in ignored]
            req["zero_steps"] = zeros
            check(key+"-missing-counterevidence", "challenge", req, "ZeroTransferPossible")
        mineral = deepcopy(main_cases["mineral-acid.claim"][0])
        mineral["zero_steps"] = ["gas"]
        check("mineral-does-not-require-gas", "challenge", mineral, "ZeroTransferPossible")
        impossible = deepcopy(main_cases["underdetermined-gas.claim"][0])
        impossible["study"]["assays"].append({"id":"impossible-ca", "node":"gas", "analyte":"Ca",
                                              "bounds":{"lower":9000000,"upper":9000000}})
        check("baseline-not-causation", "challenge", impossible, "BaseConflict")
        # A grouped claim can be impossible even though either member may be zero.
        grouped = deepcopy(main_cases["underdetermined-gas.claim"][0])
        grouped["study"]["steps"].append({"id":"gas-b","parent":"gas","kind":"co2","volume_ul":0,
                                           "amount_nmol":{"lower":-1000000,"upper":0}})
        grouped["study"]["assays"][0].update(node="gas-b", bounds={"lower":2000000,"upper":2000000})
        grouped["zero_steps"] = ["gas", "gas-b"]
        check("joint-zero-impossible", "challenge", grouped, "ZeroTransferRuledOut")
        for name in grouped["zero_steps"]:
            check("one-zero-possible-"+name, "challenge", {**grouped, "zero_steps":[name]}, "ZeroTransferPossible")
        # Exact closed boundary: 2,000,000 / 3 is not rounded to 666,667.
        edge = deepcopy(common)
        edge["source"]["dic_nmol_l"] = {"lower":3000000,"upper":3000000}
        edge["steps"] = [{"id":"a","parent":"water","kind":"take","volume_ul":1,"amount_nmol":{"lower":0,"upper":0}},
                         {"id":"d","parent":"a","kind":"dilute","volume_ul":2,"amount_nmol":{"lower":0,"upper":0}}]
        edge["assays"] = [{"id":"edge","node":"d","analyte":"Ca","bounds":{"lower":666667,"upper":666667}}]
        check("one-third-not-rounded", "reconcile", edge, "Conflict")
        edge["assays"][0]["bounds"]["lower"] = 666666
        edge_result = check("one-third-enclosed", "reconcile", edge, "Compatible")
        # The identical upstream transfer cannot become separate descendant
        # choices. Moving the declared processes after the split changes this.
        shared = deepcopy(common)
        shared["source"]["dic_nmol_l"] = {"lower":3000000,"upper":3000000}
        shared["steps"] = [
            {"id":"gas","parent":"water","kind":"co2","volume_ul":0,"amount_nmol":{"lower":-1000000,"upper":0}},
            {"id":"a","parent":"gas","kind":"take","volume_ul":500000,"amount_nmol":{"lower":0,"upper":0}},
            {"id":"b","parent":"gas","kind":"take","volume_ul":500000,"amount_nmol":{"lower":0,"upper":0}}]
        check("shared-upstream-gas-conflict","reconcile",shared,"Conflict")
        separate = deepcopy(shared)
        separate["steps"] = [
            {"id":"take-a","parent":"water","kind":"take","volume_ul":500000,"amount_nmol":{"lower":0,"upper":0}},
            {"id":"take-b","parent":"water","kind":"take","volume_ul":500000,"amount_nmol":{"lower":0,"upper":0}},
            {"id":"a","parent":"take-a","kind":"co2","volume_ul":0,"amount_nmol":{"lower":-500000,"upper":0}},
            {"id":"b","parent":"take-b","kind":"co2","volume_ul":0,"amount_nmol":{"lower":-500000,"upper":0}}]
        check("separate-postsplit-gas-compatible","reconcile",separate,"Compatible")
        # A later compensating transfer cannot repair negative earlier inventory.
        interim = deepcopy(common)
        interim["source"]["dic_nmol_l"] = {"lower":3000000,"upper":3000000}
        interim["assays"] = []
        interim["steps"] = [
            {"id":"remove","parent":"water","kind":"calcite","volume_ul":0,"amount_nmol":{"lower":-3000000,"upper":-3000000}},
            {"id":"restore","parent":"remove","kind":"calcite","volume_ul":0,"amount_nmol":{"lower":3000000,"upper":3000000}}]
        check("negative-interim-cannot-be-repaired","reconcile",interim,"Conflict")
        reversed_steps = deepcopy(interim)
        reversed_steps["steps"][0]["amount_nmol"] = {"lower":3000000,"upper":3000000}
        reversed_steps["steps"][1]["amount_nmol"] = {"lower":-3000000,"upper":-3000000}
        check("reversed-order-physically-nonnegative","reconcile",reversed_steps,"Compatible")
        # Deterministic mixed exact-volume cases with three unknowns. The oracle
        # enumerates vertices of the independently reconstructed bounded model.
        rng = random.Random(20261027)
        for i in range(32):
            study = deepcopy(common)
            study["steps"] = [
                {"id":"a","parent":"water","kind":"take","volume_ul":123457,"amount_nmol":{"lower":0,"upper":0}},
                {"id":"acid","parent":"a","kind":"hcl","volume_ul":1000,"amount_nmol":{"lower":0,"upper":10000}},
                {"id":"base","parent":"acid","kind":"naoh","volume_ul":333,"amount_nmol":{"lower":0,"upper":10000}}]
            choices = {"source:DIC": F(rng.randrange(2200000,3800000)),
                       "step:acid": F(rng.randrange(500,9500)), "step:base": F(rng.randrange(500,9500))}
            volume, amounts = simulate(study, choices)["base"]
            study["assays"] = []
            for a, n in zip(ANALYTES, amounts):
                concentration = n * F(10**6, volume)
                lo, hi = concentration.numerator//concentration.denominator, -(-concentration.numerator//concentration.denominator)
                if i % 2 and a == "TA":
                    lo += 1000; hi += 1000
                study["assays"].append({"id":a,"node":"base","analyte":a,"bounds":{"lower":lo,"upper":hi}})
            check(f"mixed-{i:02}", "reconcile", study)
        # Wire/model rejections, including silent extra-ion and volume truncation.
        invalid = []
        s = deepcopy(acid); s["source"]["magnesium_nmol_l"] = {"lower":0,"upper":0}; invalid.append(("unknown-ion",s))
        s = deepcopy(acid); s["source"]["volume_ul"] = 1000000.5; invalid.append(("fractional-volume",s))
        s = deepcopy(acid); s["assays"][0]["analyte"] = "pH"; invalid.append(("nonadditive-pH",s))
        s = deepcopy(acid); s["steps"][0]["kind"] = "untracked-acid"; invalid.append(("unknown-reagent",s))
        s = deepcopy(acid); s["steps"][1]["volume_ul"] = 600000; invalid.append(("overallocated-parent",s))
        s = deepcopy(acid); s["steps"][1]["kind"] = "dilute"; invalid.append(("duplicated-parent",s))
        s = deepcopy(acid); s["steps"][0]["amount_nmol"]["lower"] = -1; invalid.append(("negative-HCl",s))
        invalid += [("malformed-json", "{"), ("wrong-schema", {**acid,"schema":"other"})]
        for name, value in invalid:
            output = run_moon("reconcile-json", value, args.target, package="cmd/main")
            assert output.get("ok") is False and "error" in output, name
            report["negative_inputs"].append({"label":name,"output":output})
        # Mutate real outputs; require the independent verifier to reject them.
        mutations = []
        req, original = main_cases["masked-acid-base.claim"]
        base = original["base"]
        bad = deepcopy(base); bad["history"]["samples"][-1]["chloride_nmol_l"] = encoded(1000000)
        mutations.append(("omitted-HCl-counterion", req["study"], bad))
        bad = deepcopy(base); bad["history"]["samples"][-1]["sodium_nmol_l"] = encoded(0)
        mutations.append(("omitted-NaOH-counterion", req["study"], bad))
        bad = deepcopy(base); bad["history"]["samples"][-1]["ta_neq_l"] = encoded(3000000)
        mutations.append(("ignored-reagent-volume", req["study"], bad))
        req_m, original_m = main_cases["mineral-acid.claim"]
        bad = deepcopy(original_m["base"]); bad["history"]["samples"][1]["ta_neq_l"] = encoded(3100000)
        mutations.append(("calcite-one-TA-instead-of-two",req_m["study"],bad))
        bad = deepcopy(edge_result); bad["history"]["samples"][-1]["calcium_nmol_l"] = encoded(666667)
        mutations.append(("rounded-exact-boundary",edge,bad))
        bad = deepcopy(base); bad["history"]["samples"][-1]["id"] = "water"
        mutations.append(("wrong-bottle-identity",req["study"],bad))
        bad = deepcopy(base); bad["history"]["parameters"][0]["value"] = encoded(0)
        mutations.append(("changed-shared-dose",req["study"],bad))
        bad = deepcopy(base); bad["history"]["samples"].pop()
        mutations.append(("truncated-history",req["study"],bad))
        conflict = main_cases["common-source-conflict"][1]
        bad = deepcopy(conflict); bad["conflict"][0]["weight"] = encoded(-1)
        mutations.append(("negative-proof-weight",common,bad))
        bad = deepcopy(conflict); bad["conflict"].pop()
        mutations.append(("incomplete-contradiction",common,bad))
        bad = deepcopy(conflict); bad["conflict"][0]["constraint"] = "assay:unknown:upper"
        mutations.append(("wrong-proof-input-binding",common,bad))
        for label, study, bad in mutations:
            try:
                verify_assessment(study,bad)
            except (AssertionError, KeyError, ValueError, ZeroDivisionError):
                report["rejected_mutations"].append(label)
            else:
                raise AssertionError("verifier accepted mutation: "+label)
        # The explicit MoonBit replay endpoint must also reject changed choices.
        replay = run_analysis("replay", {"study":req["study"],"parameters":base["history"]["parameters"]},args.target)
        verify_history(req["study"],replay)
        altered = deepcopy(base["history"]["parameters"]); altered[0]["value"] = encoded(0)
        replay_bad = run_moon("replay-json", {"study":req["study"],"parameters":altered},args.target,package="cmd/main")
        assert replay_bad["ok"] is False
        report["negative_inputs"].append({"label":"replay-changed-dose","output":replay_bad})
        paths = [p for p in ROOT.rglob('*') if p.is_file() and not any(
            x in {'.git','.moon','_build','target','__pycache__','audit'} for x in p.relative_to(ROOT).parts)]
        (out/'verified-files.json').write_text(json.dumps({str(p.relative_to(ROOT)):hashlib.sha256(p.read_bytes()).hexdigest()
                                                          for p in sorted(paths)},indent=2)+'\n')
        report.update(passed=True, completed_at_utc=datetime.now(timezone.utc).isoformat(),
                      case_count=len(report['cases']), negative_count=len(report['negative_inputs']),
                      rejected_mutation_count=len(report['rejected_mutations']))
        save()
        print(json.dumps({k:report[k] for k in ['passed','case_count','negative_count','rejected_mutation_count','target','compare_target']},indent=2))
    except Exception as error:
        report['error'] = repr(error)
        save()
        raise


if __name__ == '__main__':
    main()