"""q_rules.py: the machine-checkable side of the pre-registration.

Every prediction in inputs/plan.json carries boolean expressions over value names:
  testable            over population values measured in Phase 3 (values3.json) and re-measured in Phase 4
  testable_runtime    over Phase 4 results that do not reveal the outcome (undefined bootstrap draws,
                      a degenerate null); evaluated before the Holm step
  direction           the estimate lies on the predicted side of the null
  magnitude           the estimate reaches the practical threshold
  precision           the Holm-adjusted one-sided p-value is below the family-wise level
  rule                the decision rule: magnitude and precision (magnitude implies direction, see the check)
and robustness variants, each with an applicability expression (population values) and a contradiction
expression (Phase 4 values): a variant contradicts its prediction only when its whole interval lies below
the threshold, so that a variant's lost power never counts as evidence against the prediction.

testable() and classify() apply the failure taxonomy fixed in Phase 3, in this order, and tally() counts the
classes. A value that is missing or not finite where a decision needs it makes the prediction NOT_TESTED,
never FAILED. Expressions may use value names, comparisons, and, or, not, + - * /, abs, min and max. Numeric
literals are refused: every threshold is a named value. A name ending in __lo or __hi reads that bound of an
interval record (for example h1_v3_drop_ci__hi).
"""
import ast, math

CLASSES = ("CONFIRMED", "CONFIRMED_NOT_ROBUST", "FAILED_DIRECTION", "FAILED_MAGNITUDE", "FAILED_PRECISION", "NOT_TESTED")
FAILED = ("FAILED_DIRECTION", "FAILED_MAGNITUDE", "FAILED_PRECISION")
_NODES = (ast.Expression, ast.BoolOp, ast.And, ast.Or, ast.UnaryOp, ast.Not, ast.USub, ast.Compare, ast.Lt, ast.LtE,
          ast.Gt, ast.GtE, ast.BinOp, ast.Add, ast.Sub, ast.Mult, ast.Div, ast.Name, ast.Load, ast.Call)
_CALLS = {"abs": abs, "min": min, "max": max}


def parse(expr):
    """Parse and validate an expression; returns the AST. Raises ValueError on anything outside the grammar."""
    tree = ast.parse(expr, mode="eval")
    for node in ast.walk(tree):
        if not isinstance(node, _NODES):
            raise ValueError(f"{type(node).__name__} not allowed in {expr!r}")
        if isinstance(node, ast.Call) and not (isinstance(node.func, ast.Name) and node.func.id in _CALLS):
            raise ValueError(f"call not allowed in {expr!r}")
    return tree


def names(expr):
    """Value names an expression reads (function names excluded)."""
    tree = parse(expr)
    calls = {id(n.func) for n in ast.walk(tree) if isinstance(n, ast.Call)}
    return {n.id for n in ast.walk(tree) if isinstance(n, ast.Name) and id(n) not in calls}


def lookup(name, values):
    """The number a name stands for: a plain number, a record's value, or an interval record's bound."""
    base, field = name, "value"
    if name.endswith(("__lo", "__hi")):
        base, field = name[:-4], name[-2:]
        if name in values and not isinstance(values[name], dict):
            return values[name]
    if base not in values:
        raise KeyError(f"{name} is not in the values")
    v = values[base]
    x = v.get(field) if isinstance(v, dict) else v
    if x is None:
        raise ValueError(f"{name} has no value")
    return x


def evaluate(expr, values):
    """Evaluate over a mapping name -> number (or name -> values-file record)."""
    env = {k: lookup(k, values) for k in names(expr)}
    return eval(compile(parse(expr), "<rule>", "eval"), {"__builtins__": {}}, dict(env, **_CALLS))


def finite(expr, values):
    """True when every name the expression reads exists and is a finite number."""
    try:
        return all(isinstance(x, (int, float)) and math.isfinite(x) for x in (lookup(k, values) for k in names(expr)))
    except (KeyError, ValueError):
        return False


def testable(pred, values, population):
    """(True, '') when both testability conditions hold; otherwise (False, reason)."""
    if not finite(pred["testable"], population) or not evaluate(pred["testable"], population):
        return False, "population condition false"
    rt = pred.get("testable_runtime")
    if rt and (not finite(rt, values) or not evaluate(rt, values)):
        return False, "runtime condition false"
    return True, ""


def classify(pred, values, population):
    """Class of one prediction and the ids of the applicable variants that contradict it."""
    ok, _ = testable(pred, values, population)
    if not ok:
        return "NOT_TESTED", []
    for part in ("direction", "magnitude", "precision"):
        if not finite(pred[part], values):
            return "NOT_TESTED", []
    if not evaluate(pred["direction"], values):
        return "FAILED_DIRECTION", []
    if not evaluate(pred["magnitude"], values):
        return "FAILED_MAGNITUDE", []
    if not evaluate(pred["precision"], values):
        return "FAILED_PRECISION", []
    against = [v["id"] for v in pred["variants"] if evaluate(v["applicable"], population) and _contradicts(v, values)]
    return ("CONFIRMED" if not against else "CONFIRMED_NOT_ROBUST"), against


def _contradicts(variant, values):
    """An applicable variant contradicts when its contradiction expression holds. A variant that could not be
    computed does not contradict; the Phase 4 check reports it as missing and the deviations log records why."""
    if not finite(variant["contradicts"], values):
        return False
    return bool(evaluate(variant["contradicts"], values))


def tally(preds, values, population):
    """Counts of every class, the failed total and the number of predictions."""
    cls = {p["id"]: classify(p, values, population)[0] for p in preds}
    out = {c: sum(1 for x in cls.values() if x == c) for c in CLASSES}
    out["FAILED"] = sum(out[c] for c in FAILED)
    out["TOTAL"] = len(preds)
    return out, cls