#!/usr/bin/env python
"""
emergence_audit — four controls for any claim of the form
    "X (a geometry, a law, a metric, a symmetry, a world model) EMERGES from a network trained on D".

The audit does not know what X is. It asks the user's pipeline to fit data from several
GENERATORS under several PRIORS, summarises each output with the user's FEATURE function, and
reports four differentials:

  C1  STATELESS generator   — data with the claimed physics removed (a closed-form curve, a
                              shuffled/random process). If feature(stateless) ≈ feature(real),
                              X is a property of the pipeline or of low-order data shape.
  C2  NON-SOLUTION generator— data that VIOLATE the law claimed to emerge. If it fits as well and
                              the violation is reproduced, the pipeline does not forbid the
                              alternative — it cannot have "found" the law.
  C3  PRIOR SWAP at fixed fit— among outputs with (near-)equal fit, swap the regulariser. The
                              part of X that moves is the prior's, not the data's.
  C4  FAMILY DISTANCE + conditioning — distance of the output to the nearest member of the
                              claimed solution family vs to the generator; and the condition
                              number of the parametrisation (function-space MSE hides parameter
                              error when it is large).

Plugin interface (a Python module or object with these attributes):
    generators : dict  name -> callable(seed) -> data          (must include 'real';
                                                                optional 'stateless', 'nonsolution')
    priors     : dict  name -> prior spec (opaque; passed to fit)   (first entry = the paper's)
    fit        : callable(data, prior, seed) -> (output, fit_quality: float, lower is better)
    feature    : callable(output) -> 1-D np.ndarray  (the summary of X)
    family_distance : callable(output) -> float        (optional)
    generator_distance : callable(output, data) -> float (optional)
    condition_number : float or callable() -> float     (optional)
    fit_tolerance : float   relative tolerance for "equal fit" in C3 (default 0.25)

Usage:
    python emergence_audit.py --plugin path/to/plugin.py --seeds 5 --out audit.json
    python emergence_audit.py --demo            # AdS/DL reproduction (Hashimoto et al. 2018)
"""
from __future__ import annotations

import argparse
import importlib.util
import json
import os
import sys
import numpy as np


def _load_plugin(path):
    spec = importlib.util.spec_from_file_location("audit_plugin", path)
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    return mod


def _stat(xs):
    xs = np.asarray(xs, float)
    return float(np.mean(xs)), float(np.std(xs))


def run_audit(plugin, seeds=5, verbose=True):
    gens = plugin.generators
    priors = plugin.priors
    paper_prior = next(iter(priors))
    tol = getattr(plugin, "fit_tolerance", 0.25)
    has_fam = hasattr(plugin, "family_distance")
    has_gen = hasattr(plugin, "generator_distance")

    runs = {}   # (gen, prior) -> list of dict(fit, feature, fam, gen)
    for g in gens:
        for p in priors:
            if g != "real" and p != paper_prior and g != "nonsolution":
                continue          # stateless only under the paper's prior
            rows = []
            for s in range(seeds):
                data = gens[g](s)
                out, fit = plugin.fit(data, priors[p], s)
                row = dict(fit=float(fit), feature=np.asarray(plugin.feature(out), float))
                if has_fam:
                    row["fam"] = float(plugin.family_distance(out))
                if has_gen:
                    row["gen"] = float(plugin.generator_distance(out, data))
                rows.append(row)
            runs[(g, p)] = rows
            if verbose:
                f = _stat([r["fit"] for r in rows])
                extra = ""
                if has_fam:
                    extra += f"  family {_stat([r['fam'] for r in rows])[0]:.3f}"
                if has_gen:
                    extra += f"  generator {_stat([r['gen'] for r in rows])[0]:.3f}"
                print(f"  [{g:11s} | {p:12s}] fit {f[0]:.4g}±{f[1]:.2g}{extra}")

    def feat(g, p):
        F = np.array([r["feature"] for r in runs[(g, p)]])
        return F.mean(0), F.std(0).mean() + 1e-12

    report = {"seeds": seeds, "paper_prior": paper_prior, "controls": {}}
    base_f, base_sd = feat("real", paper_prior)
    base_fit = _stat([r["fit"] for r in runs[("real", paper_prior)]])[0]

    # C1 stateless
    if "stateless" in gens:
        sf, _ = feat("stateless", paper_prior)
        d = float(np.linalg.norm(sf - base_f) / np.sqrt(len(base_f)))
        z = d / base_sd
        report["controls"]["C1_stateless"] = dict(feature_rms_gap=d, gap_over_seed_sd=z,
            verdict=("SAME X from stateless data — X is not evidence of the physics" if z < 2
                     else "stateless data give a different X (differential exists)"))

    # C2 non-solution
    if "nonsolution" in gens:
        ns_fit = _stat([r["fit"] for r in runs[("nonsolution", paper_prior)]])[0]
        c2 = dict(fit_ratio_nonsolution_over_real=float(ns_fit / base_fit) if base_fit > 0 else float("nan"))
        if has_fam:
            fr = _stat([r["fam"] for r in runs[("real", paper_prior)]])[0]
            fn = _stat([r["fam"] for r in runs[("nonsolution", paper_prior)]])[0]
            c2.update(family_distance_real=fr, family_distance_nonsolution=fn,
                      violation_reproduced=bool(fn > 2 * fr))
        fits_as_well = c2["fit_ratio_nonsolution_over_real"] <= 1 + tol
        c2["verdict"] = ("non-solution data fit as well" + (" and the violation is reproduced" if c2.get("violation_reproduced") else "")
                         + " — the pipeline does not forbid the alternative") if fits_as_well else "non-solution data fit worse (pipeline discriminates)"
        report["controls"]["C2_nonsolution"] = c2

    # C3 prior swap at fixed fit
    c3 = {}
    for p in priors:
        if p == paper_prior:
            continue
        pf, _ = feat("real", p)
        pfit = _stat([r["fit"] for r in runs[("real", p)]])[0]
        equal = abs(pfit - base_fit) <= tol * base_fit
        d = float(np.linalg.norm(pf - base_f) / np.sqrt(len(base_f)))
        entry = dict(fit=pfit, equal_fit=bool(equal), feature_rms_shift=d, shift_over_seed_sd=d / base_sd)
        if has_fam:
            fam_p = _stat([r["fam"] for r in runs[("real", p)]])[0]
            fam_0 = _stat([r["fam"] for r in runs[("real", paper_prior)]])[0]
            entry["family_distance"] = fam_p
            entry["family_distance_ratio_vs_paper_prior"] = float(fam_p / fam_0) if fam_0 > 0 else float("nan")
        if has_gen:
            entry["generator_distance"] = _stat([r["gen"] for r in runs[("real", p)]])[0]
            entry["generator_distance_paper_prior"] = _stat([r["gen"] for r in runs[("real", paper_prior)]])[0]
        if "nonsolution" in gens and ("nonsolution", p) in runs and has_fam:
            entry["nonsolution_family_distance_under_this_prior"] = _stat([r["fam"] for r in runs[("nonsolution", p)]])[0]
        c3[p] = entry
    if c3:
        moved = [p for p, e in c3.items() if e["equal_fit"] and (e["shift_over_seed_sd"] > 2 or
                 not (0.5 <= e.get("family_distance_ratio_vs_paper_prior", 1.0) <= 2.0))]
        report["controls"]["C3_prior_swap"] = dict(priors=c3,
            verdict=("X moves under a prior swap at equal fit — the moved part is the prior's, not the data's: " + ", ".join(moved))
                    if moved else "X is stable under prior swap at equal fit")

    # C4 family + conditioning
    c4 = {}
    if has_fam:
        c4["family_distance_real"] = _stat([r["fam"] for r in runs[("real", paper_prior)]])
    if has_gen:
        c4["generator_distance_real"] = _stat([r["gen"] for r in runs[("real", paper_prior)]])
    cn = getattr(plugin, "condition_number", None)
    if cn is not None:
        c4["condition_number"] = float(cn() if callable(cn) else cn)
    if c4:
        notes = []
        if "family_distance_real" in c4 and "generator_distance_real" in c4 and c4["generator_distance_real"][0] > 2 * c4["family_distance_real"][0]:
            notes.append("output is closer to ANOTHER member of the family than to its generator")
        if c4.get("condition_number", 0) > 30:
            notes.append(f"ill-conditioned parametrisation (cond {c4['condition_number']:.0f}): function-space fit hides parameter error")
        c4["verdict"] = "; ".join(notes) if notes else "no family/conditioning flags"
        report["controls"]["C4_family_conditioning"] = c4

    # overall
    flags = [k for k, v in report["controls"].items() if any(w in v["verdict"] for w in
             ("not evidence", "does not forbid", "prior's", "ANOTHER member", "ill-conditioned"))]
    report["overall"] = ("EMERGENCE CLAIM NOT SUPPORTED by these controls: " + ", ".join(flags)) if flags else \
                        "no control flagged — the differential the claim needs is present"
    return report, runs


def print_report(report):
    print("\n=== emergence audit ===")
    for k, v in report["controls"].items():
        print(f"{k:26s} {v['verdict']}")
    print("-" * 26)
    print(report["overall"])


def main(argv=None):
    ap = argparse.ArgumentParser(description="Four controls for 'X emerges from a network' claims.")
    ap.add_argument("--plugin", help="path to a plugin .py exposing generators/priors/fit/feature[/family_distance/...]")
    ap.add_argument("--seeds", type=int, default=5)
    ap.add_argument("--out", help="write JSON report here")
    ap.add_argument("--demo", action="store_true", help="run the AdS/DL (Hashimoto et al. 2018) plugin, 3 seeds, full recipe (~4 min)")
    a = ap.parse_args(argv)
    if a.demo:
        here = os.path.dirname(os.path.abspath(__file__))
        plugin = _load_plugin(os.path.join(here, "plugins", "adsdl_plugin.py"))
        seeds = min(a.seeds, 3)      # full recipe (60 epochs); reduced epochs flip verdicts
    elif a.plugin:
        plugin = _load_plugin(a.plugin); seeds = a.seeds
    else:
        ap.error("give --plugin or --demo")
    report, _ = run_audit(plugin, seeds=seeds)
    print_report(report)
    if a.out:
        json.dump(report, open(a.out, "w"), indent=1, default=float)
        print(f"\nwrote {a.out}")
    return 0


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