#!/usr/bin/env python3
"""Qelly plan checker: an independent, single-file replay of one answer.

    python3 qelly_check.py instance.json plan.json
    python3 qelly_check.py instance.json plan.json --solution solution.json
    python3 qelly_check.py instance.json sample.json --scenario late_1 --solver my-annealer

``instance.json`` follows schema ``qelly.liquidity.v1``. ``plan.json`` is either

* a plan, ``{"item_id": "option_id", ...}``, or
* a binary sample as returned by a QUBO/Ising/CQM solver,
  ``{"x[item_id|option_id]": 0 or 1, ...}``. Slack variables (``s[...]``) are
  ignored. Every item must have exactly one selected option.

The checker replays the plan week by week with integer arithmetic: closing
cash = opening cash + committed flows + the flows of every selected option. A
plan is feasible when every account stays at or above its reserve in every
week. It prints a hash-linked certificate as JSON on standard output and
diagnostics on standard error.

A certificate replays feasibility and cost. Optimality is only relative to a
supplied reference optimum: with ``--solution``, ``reference_optimum`` is the
optimum published in that file, ``gap`` is cost minus that value and
``matches_reference_optimum`` says whether they are equal. The checker does not
prove the reference optimum (for Qelly Bench it comes from an exact
contraction cross-checked by HiGHS), and it does not establish that the
instance's inputs are true.

Exit status: 0 feasible, 1 infeasible or invalid answer, 2 usage or input error.

Standard library only; Python 3.8 or newer. This file shares no code with any
solver. For the same inputs it produces the same certificate digest as the
reference implementation (``qelly_dtn.verify.certificate``).
"""

import argparse
import hashlib
import json
import sys

SCHEMA = "qelly.liquidity.v1"
CHECKER_VERSION = "qelly-plan-checker/1.1"


class InputError(Exception):
    """The instance, plan or solution file cannot be used."""


# ---- instance normalization ---------------------------------------------------


def _is_int(value):
    return isinstance(value, int) and not isinstance(value, bool)


def _int(data, key):
    value = data.get(key) if isinstance(data, dict) else None
    if not _is_int(value):
        raise InputError("%s must be an integer" % key)
    return value


def _ints(values, name):
    if not isinstance(values, list) or not all(_is_int(v) for v in values):
        raise InputError("%s must be a list of integers" % name)
    return values


def _list(data, key):
    value = data.get(key) if isinstance(data, dict) else None
    if not isinstance(value, list):
        raise InputError("%s must be a list" % key)
    return value


def canonical_instance(data):
    """Validate an instance and return its canonical dictionary.

    Optional fields get their defaults, option flows are sorted by
    (week, account position, amount) and every scenario lists every account, so
    two files describing the same instance hash to the same digest.
    """
    if not isinstance(data, dict) or data.get("schema") != SCHEMA:
        raise InputError("unsupported schema: %r" % (data.get("schema") if isinstance(data, dict) else None,))
    horizon = _int(data, "horizon")
    if horizon < 1:
        raise InputError("horizon must be >= 1")

    accounts = []
    for raw in _list(data, "accounts"):
        if not isinstance(raw, dict) or "id" not in raw:
            raise InputError("every account needs an id")
        fixed = _ints(raw.get("fixed_flows"), "fixed_flows")
        if len(fixed) != horizon:
            raise InputError("fixed_flows must have one entry per week")
        accounts.append(
            {
                "id": str(raw["id"]),
                "opening_cash": _int(raw, "opening_cash"),
                "reserve": _int(raw, "reserve"),
                "fixed_flows": list(fixed),
            }
        )
    if not accounts:
        raise InputError("at least one account is required")
    position = {account["id"]: i for i, account in enumerate(accounts)}
    if len(position) != len(accounts):
        raise InputError("duplicate account id")

    items = []
    for raw in _list(data, "items"):
        if not isinstance(raw, dict) or "id" not in raw:
            raise InputError("every item needs an id")
        options = []
        for opt in _list(raw, "options"):
            if not isinstance(opt, dict) or "id" not in opt:
                raise InputError("every option needs an id")
            flows = []
            for flow in opt.get("flows", []):
                week = flow.get("week")
                account = flow.get("account", accounts[0]["id"])
                amount = flow.get("amount")
                if account not in position:
                    raise InputError("unknown account %r" % (account,))
                if not (_is_int(week) and 1 <= week <= horizon):
                    raise InputError("flow week outside the horizon")
                if not _is_int(amount):
                    raise InputError("flow amount must be an integer")
                flows.append((week, position[account], amount))
            options.append(
                {
                    "id": str(opt["id"]),
                    "label": str(opt.get("label", "")),
                    "cost": _int(opt, "cost"),
                    "flows": [
                        {"week": w, "account": accounts[a]["id"], "amount": amount} for w, a, amount in sorted(flows)
                    ],
                }
            )
        if not options:
            raise InputError("item %r has no options" % (raw.get("id"),))
        if len({o["id"] for o in options}) != len(options):
            raise InputError("item %r has duplicate option ids" % (raw.get("id"),))
        items.append(
            {
                "id": str(raw["id"]),
                "kind": str(raw.get("kind", "choice")),
                "label": str(raw.get("label", "")),
                "options": options,
            }
        )
    if len({item["id"] for item in items}) != len(items):
        raise InputError("duplicate item id")

    scenarios = {}
    for name, per_account in (data.get("scenarios") or {}).items():
        if not isinstance(per_account, dict):
            raise InputError("scenario %r must map account ids to weekly flows" % (name,))
        vectors = {}
        for account in accounts:
            vector = _ints(per_account.get(account["id"], list(account["fixed_flows"])), "scenario flows")
            if len(vector) != horizon:
                raise InputError("scenario fixed_flows must have one entry per week")
            vectors[account["id"]] = list(vector)
        scenarios[str(name)] = vectors

    if "id" not in data:
        raise InputError("the instance needs an id")
    return {
        "schema": SCHEMA,
        "id": str(data["id"]),
        "title": str(data.get("title", "")),
        "currency": str(data.get("currency", "USD")),
        "horizon": horizon,
        "accounts": accounts,
        "items": items,
        "scenarios": scenarios,
        "notes": str(data.get("notes", "")),
    }


def _digest(value):
    canonical = json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False)
    return "sha256:" + hashlib.sha256(canonical.encode("utf-8")).hexdigest()


def instance_digest(instance):
    return _digest(instance)


def fixed_flows(instance, scenario):
    """Committed weekly flows per account id for ``scenario`` ("base" = the instance itself)."""
    if scenario in (None, "base"):
        return {account["id"]: account["fixed_flows"] for account in instance["accounts"]}
    if scenario not in instance["scenarios"]:
        raise InputError(
            "unknown scenario %r (available: base%s)"
            % (scenario, "".join(", " + name for name in sorted(instance["scenarios"])))
        )
    return instance["scenarios"][scenario]


# ---- answers ------------------------------------------------------------------


def is_binary_sample(answer):
    return isinstance(answer, dict) and any(isinstance(key, str) and key.startswith("x[") for key in answer)


def decode_sample(instance, sample):
    """Binary sample -> (plan, problems). Items that are not one-hot are left out of the plan.

    A sample that cannot belong to this instance (unknown decision variables,
    values other than 0/1) is an input error, not a solver failure.
    """
    problems = []
    expected = {}
    for item in instance["items"]:
        for option in item["options"]:
            expected["x[%s|%s]" % (item["id"], option["id"])] = (item["id"], option["id"])
    selected = {}
    for key, value in sample.items():
        if not isinstance(key, str) or not key.startswith("x["):
            continue  # slack and auxiliary variables carry no decision
        if key not in expected:
            raise InputError("decision variable %r does not belong to this instance" % (key,))
        if isinstance(value, bool):
            value = int(value)
        if value not in (0, 1):
            raise InputError(
                "variable %r has value %r; samples must be binary (convert spins with x = (s + 1) / 2)" % (key, value)
            )
        if value == 1:
            item_id, option_id = expected[key]
            selected.setdefault(item_id, []).append(option_id)
    plan = {}
    for item in instance["items"]:
        chosen = selected.get(item["id"], [])
        if len(chosen) == 1:
            plan[item["id"]] = chosen[0]
        else:
            problems.append("item %r: %d options selected, exactly one is required" % (item["id"], len(chosen)))
    return plan, problems


def evaluate(instance, plan, scenario="base"):
    """Replay ``plan`` ({item_id: option_id}) and report feasibility and cost."""
    fixed = fixed_flows(instance, scenario)
    problems = []
    chosen = {}
    for item in instance["items"]:
        option_id = plan.get(item["id"])
        match = [option for option in item["options"] if option["id"] == option_id]
        if not match:
            problems.append("item %r: missing or unknown option %r" % (item["id"], option_id))
        else:
            chosen[item["id"]] = match[0]
    unknown = sorted(set(plan) - {item["id"] for item in instance["items"]})
    if unknown:
        problems.append("unknown items: %s" % (unknown,))
    if problems:
        return {"valid": False, "feasible": False, "problems": problems}

    cost = sum(option["cost"] for option in chosen.values())
    paths, breaches = {}, []
    for account in instance["accounts"]:
        cash, path = account["opening_cash"], []
        for week in range(1, instance["horizon"] + 1):
            cash += fixed[account["id"]][week - 1]
            for option in chosen.values():
                for flow in option["flows"]:
                    if flow["week"] == week and flow["account"] == account["id"]:
                        cash += flow["amount"]
            path.append(cash)
            if cash < account["reserve"]:
                breaches.append(
                    {"account": account["id"], "week": week, "cash": cash, "reserve": account["reserve"]}
                )
        paths[account["id"]] = path
    reserves = {account["id"]: account["reserve"] for account in instance["accounts"]}
    return {
        "valid": True,
        "feasible": not breaches,
        "cost": cost,
        "cash_paths": paths,
        "min_headroom": min(cash - reserves[name] for name, path in paths.items() for cash in path),
        "breaches": breaches,
        "problems": [],
    }


def certificate(instance, plan, solver, reference_optimum=None, scenario="base"):
    """A portable, hash-linked record of one checked answer.

    Feasibility and cost are replayed here; ``reference_optimum`` is supplied by
    the caller and is not verified.
    """
    result = evaluate(instance, plan, scenario)
    body = {
        "checker": CHECKER_VERSION,
        "instance_id": instance["id"],
        "instance_digest": instance_digest(instance),
        "scenario": scenario,
        "solver": solver,
        "plan": dict(sorted(plan.items())),
        "valid": result["valid"],
        "feasible": result["feasible"],
        "cost": result.get("cost"),
        "reference_optimum": reference_optimum,
        "gap": None if reference_optimum is None or not result["feasible"] else result["cost"] - reference_optimum,
        "matches_reference_optimum": bool(
            result["feasible"] and reference_optimum is not None and result["cost"] == reference_optimum
        ),
        "breaches": result.get("breaches", []),
    }
    body["certificate_digest"] = _digest(body)
    return body, result


def reference_optimum_from(solution, instance, scenario):
    """Read the published reference optimum for ``scenario``; refuse a solution file for another instance."""
    if not isinstance(solution, dict):
        raise InputError("the solution file must contain a JSON object")
    if solution.get("instance_digest") != instance_digest(instance):
        raise InputError("the solution file belongs to a different instance (digest mismatch)")
    block = solution
    if scenario not in (None, "base"):
        block = (solution.get("scenarios") or {}).get(scenario)
        if block is None:
            raise InputError("the solution file has no entry for scenario %r" % (scenario,))
    optimum = block.get("optimum")
    if optimum is not None and not _is_int(optimum):
        raise InputError("the published optimum must be an integer or null")
    return optimum


# ---- command line -------------------------------------------------------------


def _load(path, what):
    try:
        with open(path, encoding="utf-8") as handle:
            return json.load(handle)
    except OSError as error:
        raise InputError("cannot read %s %r: %s" % (what, path, error.strerror or error))
    except ValueError as error:
        raise InputError("%s %r is not valid JSON: %s" % (what, path, error))


def main(argv=None):
    parser = argparse.ArgumentParser(
        description=(
            "Independently replay one plan for a qelly.liquidity.v1 instance and print a certificate. "
            "The certificate replays feasibility and cost; optimality is only relative to a supplied "
            "reference optimum (--solution), which this checker does not prove."
        ),
    )
    parser.add_argument("instance", help="instance.json")
    parser.add_argument("plan", help="plan.json: {item: option} or a binary sample {\"x[item|option]\": 0/1}")
    parser.add_argument(
        "--solution",
        help=(
            "solution.json with a published reference optimum; adds reference_optimum, gap and "
            "matches_reference_optimum (relative to that reference, not a proof of optimality)"
        ),
    )
    parser.add_argument("--scenario", default="base", help="scenario name (default: base)")
    parser.add_argument("--solver", default="unspecified", help="free-text name of the solver that produced the plan")
    args = parser.parse_args(argv)

    try:
        instance = canonical_instance(_load(args.instance, "instance"))
        answer = _load(args.plan, "plan")
        if not isinstance(answer, dict):
            raise InputError("the plan file must contain a JSON object")
        fixed_flows(instance, args.scenario)  # validates the scenario name
        optimum = None
        if args.solution:
            optimum = reference_optimum_from(_load(args.solution, "solution"), instance, args.scenario)
    except InputError as error:
        print("qelly_check: %s" % error, file=sys.stderr)
        return 2

    problems = []
    if is_binary_sample(answer):
        try:
            plan, problems = decode_sample(instance, answer)
        except InputError as error:
            print("qelly_check: %s" % error, file=sys.stderr)
            return 2
    else:
        plan = {str(key): value for key, value in answer.items()}
    record, result = certificate(instance, plan, args.solver, optimum, args.scenario)
    print(json.dumps(record, indent=1, sort_keys=True, ensure_ascii=False))

    for problem in problems + result.get("problems", []):
        print("qelly_check: %s" % problem, file=sys.stderr)
    if result["valid"]:
        print(
            "qelly_check: cost %s, minimum headroom above reserve %s, %s"
            % (result["cost"], result["min_headroom"], "feasible" if result["feasible"] else "NOT feasible"),
            file=sys.stderr,
        )
    return 0 if record["feasible"] else 1


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