#!/usr/bin/env python3
"""scan_experiments.py - inventory of the machine-learning experiments in a project (stdlib only).

    python scan_experiments.py [ROOT] [--out paper/inventory]

Walks ROOT (default '.'), skipping paper/, .git, venvs, node_modules, caches, and classifies:
  code     .py / .ipynb files that import an ML framework; entry points (__main__, argparse,
           hydra, click), CLI arguments with defaults, seeds, datasets, models, optimizers, metrics
  configs  .yaml/.yml/.json/.toml files that look like hyperparameter configs
  results  .csv/.tsv/.json/.jsonl result tables, *.log files with metric lines, TensorBoard
           event files, wandb/ and mlruns/ folders, checkpoints (only counted)
  docs     README*, *.md, requirements.txt, environment.yml, pyproject.toml (dependencies)
Writes <out>.json (full detail) and <out>.md (human-readable), prints a compact summary.
The script only reads; nothing is executed.
"""
import ast, json, re, sys
from collections import Counter
from pathlib import Path

SKIP = {"paper", ".git", "__pycache__", ".venv", "venv", "env", "node_modules", ".ipynb_checkpoints",
        ".pytest_cache", ".mypy_cache", ".ruff_cache", "dist", "build", "site-packages", ".idea", ".vscode"}
ML_LIBS = {"torch": "PyTorch", "tensorflow": "TensorFlow", "keras": "Keras", "jax": "JAX", "flax": "Flax",
           "sklearn": "scikit-learn", "xgboost": "XGBoost", "lightgbm": "LightGBM", "catboost": "CatBoost",
           "transformers": "Hugging Face Transformers", "datasets": "Hugging Face Datasets",
           "lightning": "PyTorch Lightning", "pytorch_lightning": "PyTorch Lightning",
           "torchvision": "torchvision", "timm": "timm", "peft": "PEFT", "trl": "TRL", "accelerate": "Accelerate",
           "optuna": "Optuna", "statsmodels": "statsmodels", "gymnasium": "Gymnasium", "gym": "Gym",
           "stable_baselines3": "Stable-Baselines3", "torch_geometric": "PyG", "dgl": "DGL",
           "sentence_transformers": "Sentence-Transformers", "wandb": "Weights & Biases", "mlflow": "MLflow",
           "numpy": "NumPy", "pandas": "pandas", "scipy": "SciPy"}
CORE_ML = set(ML_LIBS) - {"numpy", "pandas", "scipy", "wandb", "mlflow", "optuna"}
DATASETS = r"""MNIST|Fashion-?MNIST|KMNIST|EMNIST|CIFAR-?10|CIFAR-?100|SVHN|STL-?10|ImageNet(?:-1k|-21k)?|Tiny-?ImageNet|
CelebA|LSUN|COCO|MS-?COCO|Pascal\s?VOC|Cityscapes|ADE20K|Places365|Caltech-?101|Caltech-?256|Oxford-?Pets|Flowers-?102|Food-?101|
IMDB|SST-?2|GLUE|SuperGLUE|SQuAD|MNLI|QNLI|QQP|RTE|CoLA|AG-?News|Yelp|Amazon|WikiText-?\d*|Penn\s?Treebank|PTB|OpenWebText|
C4|The\s?Pile|LAMBADA|HellaSwag|MMLU|GSM8K|HumanEval|MBPP|ARC|TruthfulQA|WMT\d*|Multi30k|LibriSpeech|Common\s?Voice|
UCI|Adult|Iris|Wine|Breast\s?Cancer|California\s?Housing|Boston|Covertype|Higgs|Titanic|OpenML|
Cora|CiteSeer|PubMed|OGB[\w-]*|QM9|ZINC|MovieLens|Criteo|ETTh\d|ETTm\d|Electricity|Traffic|Weather|M4|
CartPole|MountainCar|Atari|MuJoCo|HalfCheetah|Hopper|Walker2d|Ant|Humanoid|Procgen|D4RL"""
DATASET_RE = re.compile(r"(?<![\w-])(" + DATASETS.replace("\n", "") + r")(?![\w-])", re.I)
HF_DS_RE = re.compile(r"load_dataset\(\s*['\"]([^'\"]+)['\"](?:\s*,\s*['\"]([^'\"]+)['\"])?")
MODEL_RE = re.compile(r"(?<![\w])(ResNet-?\d+|ResNeXt-?\d*|WideResNet[\w-]*|WRN-?\d+-?\d*|VGG-?\d+|DenseNet-?\d+|"
                      r"MobileNet(?:V\d)?|EfficientNet-?\w*|ViT-?[\w/]*|DeiT[\w-]*|Swin[\w-]*|ConvNeXt[\w-]*|U-?Net|"
                      r"BERT[\w-]*|RoBERTa[\w-]*|DistilBERT|ALBERT|DeBERTa[\w-]*|ELECTRA|T5[\w-]*|BART|GPT-?\w*|"
                      r"Llama[\w.-]*|Mistral[\w-]*|Qwen[\w.-]*|Gemma[\w-]*|Phi-?\d[\w-]*|CLIP|LSTM|GRU|Transformer|"
                      r"MLP|CNN|GCN|GAT|GraphSAGE|GIN|VAE|GAN|DCGAN|StyleGAN\d?|DDPM|DDIM|Diffusion|"
                      r"RandomForest\w*|GradientBoosting\w*|LogisticRegression|SVC|SVM|KNeighbors\w*|XGB\w*|LGBM\w*|"
                      r"PPO|SAC|DQN|A2C|TD3|DDPG)(?![\w])", re.I)
OPTIM_RE = re.compile(r"\b(SGD|Adam|AdamW|RMSprop|Adagrad|Adafactor|LAMB|LARS|Lion)\b")
METRIC_WORDS = r"acc(?:uracy)?|top-?[15]|loss|f1|precision|recall|auc|auroc|auprc|mse|rmse|mae|r2|bleu|rouge[\w-]*|" \
               r"perplexity|ppl|map|miou|iou|dice|psnr|ssim|fid|is_score|reward|return|ece|nll|wer|cer|em|exact_match"
METRIC_RE = re.compile(r"(?<![a-z])(" + METRIC_WORDS + r")(?![a-z])", re.I)
SEED_RE = re.compile(r"(?:manual_seed|random\.seed|np\.random\.seed|seed_everything|set_seed|random_state)\s*[(=]\s*([\w.]+)")
LOG_METRIC_RE = re.compile(r"(?i)\b(" + METRIC_WORDS + r"|val_\w+|test_\w+|train_\w+)\s*[:=]\s*(-?\d+\.?\d*(?:e-?\d+)?)")
HP_KEYS = {"lr", "learning_rate", "batch_size", "epochs", "num_epochs", "max_epochs", "weight_decay", "momentum",
           "dropout", "hidden_size", "hidden_dim", "num_layers", "seed", "optimizer", "warmup", "max_steps",
           "n_estimators", "max_depth", "model", "dataset", "scheduler", "temperature", "gamma", "beta1", "beta2"}
RESULT_HINT = re.compile(r"(?i)result|metric|score|eval|summary|log|history|progress|benchmark|runs?|ablation|output")


def walk(root):
    for p in sorted(root.rglob("*")):
        rel = p.relative_to(root).parts
        if any(part in SKIP or (part.startswith(".") and part != ".") for part in rel[:-1]):
            continue
        if p.is_file():
            yield p


def read(p, limit=2_000_000):
    try:
        return p.read_text(encoding="utf-8", errors="replace")[:limit]
    except OSError:
        return ""


def py_info(src, name):
    info = {"imports": set(), "entry_point": False, "cli_args": [], "seeds": [], "functions": [], "classes": []}
    try:
        tree = ast.parse(src)
    except SyntaxError:
        tree = None
    if tree:
        for node in ast.walk(tree):
            if isinstance(node, ast.Import):
                info["imports"] |= {a.name.split(".")[0] for a in node.names}
            elif isinstance(node, ast.ImportFrom) and node.module and node.level == 0:
                info["imports"].add(node.module.split(".")[0])
            elif isinstance(node, ast.FunctionDef) and node.col_offset == 0:
                info["functions"].append(node.name)
            elif isinstance(node, ast.ClassDef) and node.col_offset == 0:
                bases = [ast.unparse(b) for b in node.bases]
                info["classes"].append(f"{node.name}({', '.join(bases)})" if bases else node.name)
            elif isinstance(node, ast.Call) and getattr(node.func, "attr", "") == "add_argument":
                flags = [a.value for a in node.args if isinstance(a, ast.Constant) and isinstance(a.value, str)]
                kw = {k.arg: ast.unparse(k.value) for k in node.keywords if k.arg in ("default", "type", "help", "choices")}
                if flags:
                    info["cli_args"].append({"flag": flags[-1], **{k: v[:80] for k, v in kw.items()}})
    info["entry_point"] = bool(re.search(r"__name__\s*==\s*['\"]__main__['\"]|@hydra\.main|@click\.command|typer\.run", src))
    info["seeds"] = sorted(set(SEED_RE.findall(src)))[:10]
    return info


def nb_source(text):
    try:
        nb = json.loads(text)
    except json.JSONDecodeError:
        return "", []
    code, outs = [], []
    for c in nb.get("cells", []):
        s = "".join(c.get("source", []))
        if c.get("cell_type") == "code":
            code.append(s)
            for o in c.get("outputs", []):
                t = "".join(o.get("text", []) or o.get("data", {}).get("text/plain", []))
                if t and LOG_METRIC_RE.search(t):
                    outs.append(t.strip()[:300])
        elif c.get("cell_type") == "markdown":
            code.append("# " + s.replace("\n", "\n# "))
    return "\n".join(code), outs[:10]


def table_info(p, text):
    if p.suffix in (".csv", ".tsv"):
        lines = [l for l in text.splitlines() if l.strip()]
        if not lines:
            return None
        sep = "\t" if p.suffix == ".tsv" or lines[0].count("\t") > lines[0].count(",") else ","
        return {"kind": "table", "columns": [c.strip() for c in lines[0].split(sep)][:40], "rows": len(lines) - 1,
                "head": lines[1:4]}
    if p.suffix == ".jsonl":
        lines = text.splitlines()
        try:
            first = json.loads(lines[0])
        except (json.JSONDecodeError, IndexError):
            return None
        return {"kind": "jsonl", "keys": list(first)[:40] if isinstance(first, dict) else None, "rows": len(lines),
                "head": lines[:2]}
    if p.suffix == ".json":
        try:
            js = json.loads(text)
        except json.JSONDecodeError:
            return None
        def shape(x, d=0):
            if d > 2: return "..."
            if isinstance(x, dict): return {k: shape(v, d + 1) for k, v in list(x.items())[:15]}
            if isinstance(x, list): return [f"list[{len(x)}]", shape(x[0], d + 1) if x else None]
            return type(x).__name__ if not isinstance(x, (int, float)) else x
        return {"kind": "json", "shape": shape(js)}
    return None


def is_config(text, suffix):
    keys = set(re.findall(r"(?m)^\s*[\"']?([A-Za-z_][\w.-]*)[\"']?\s*[:=]", text))
    return len({k.lower().split(".")[-1] for k in keys} & HP_KEYS) >= 2


def scan(root):
    inv = {"root": str(root.resolve()), "code": [], "notebooks": [], "configs": [], "results": [], "logs": [],
           "tensorboard": [], "trackers": [], "checkpoints": Counter(), "docs": [], "figures": [],
           "libraries": Counter(), "datasets": Counter(), "models": Counter(), "optimizers": Counter(),
           "metrics": Counter()}
    for d in ("wandb", "mlruns", "lightning_logs", "runs", "outputs", "multirun"):
        for p in root.rglob(d):
            if p.is_dir() and not (set(p.relative_to(root).parts) & SKIP):
                inv["trackers"].append(str(p.relative_to(root)))
    for p in walk(root):
        rel, suf, name = str(p.relative_to(root)), p.suffix.lower(), p.name.lower()
        if suf in (".pt", ".pth", ".ckpt", ".safetensors", ".h5", ".keras", ".pkl", ".joblib", ".onnx", ".bin", ".npz", ".npy"):
            inv["checkpoints"][suf] += 1; continue
        if "tfevents" in name:
            inv["tensorboard"].append(rel); continue
        if suf in (".png", ".jpg", ".jpeg", ".pdf", ".svg"):
            inv["figures"].append(rel); continue
        if name.startswith("readme") or name in ("requirements.txt", "environment.yml", "pyproject.toml", "setup.py", "agents.md") or suf == ".md":
            text = read(p, 20_000)
            inv["docs"].append({"path": rel, "first_lines": [l for l in text.splitlines() if l.strip()][:6]})
            inv["datasets"].update(m.group(1) for m in DATASET_RE.finditer(text))
            inv["models"].update(MODEL_RE.findall(text))
            continue
        if suf in (".py", ".ipynb"):
            text = read(p)
            outs = []
            if suf == ".ipynb":
                text, outs = nb_source(text)
            info = py_info(text, rel)
            libs = sorted(i for i in info["imports"] if i in ML_LIBS)
            inv["libraries"].update(ML_LIBS[l] for l in libs)
            ds = Counter(m.group(1) for m in DATASET_RE.finditer(text))
            ds.update(":".join(x for x in m.groups() if x) for m in HF_DS_RE.finditer(text))
            models = Counter(MODEL_RE.findall(text))
            optims = Counter(OPTIM_RE.findall(text))
            metrics = Counter(m.lower() for m in METRIC_RE.findall(text))
            inv["datasets"].update(ds); inv["models"].update(models); inv["optimizers"].update(optims)
            if not (set(libs) & CORE_ML) and not ds and not info["entry_point"]:
                continue  # plain helper without ML content
            inv["metrics"].update(metrics)
            entry = {"path": rel, "lines": text.count("\n") + 1, "ml_libraries": [ML_LIBS[l] for l in libs],
                     "entry_point": info["entry_point"], "cli_args": info["cli_args"][:40], "seeds": info["seeds"],
                     "datasets": sorted(ds), "models": sorted(models), "optimizers": sorted(optims),
                     "metrics": sorted(metrics), "functions": info["functions"][:30], "classes": info["classes"][:20]}
            if outs:
                entry["metric_outputs"] = outs
            inv["notebooks" if suf == ".ipynb" else "code"].append(entry)
            continue
        if suf in (".yaml", ".yml", ".toml", ".cfg", ".ini") or (suf == ".json" and not RESULT_HINT.search(rel)):
            text = read(p, 200_000)
            if is_config(text, suf):
                inv["configs"].append({"path": rel, "preview": [l for l in text.splitlines() if l.strip()][:25]})
                inv["datasets"].update(m.group(1) for m in DATASET_RE.finditer(text))
                continue
            if suf != ".json":
                continue
        if suf in (".csv", ".tsv", ".json", ".jsonl"):
            text = read(p)
            ti = table_info(p, text)
            if ti:
                cols = " ".join(ti.get("columns") or ti.get("keys") or [str(ti.get("shape"))[:500]])
                inv["metrics"].update(m.lower() for m in METRIC_RE.findall(cols))
                inv["results"].append({"path": rel, "bytes": p.stat().st_size, **ti})
            continue
        if suf in (".log", ".txt", ".out", ".err"):
            text = read(p)
            hits = LOG_METRIC_RE.findall(text)
            if hits:
                inv["logs"].append({"path": rel, "metric_lines": len(hits),
                                    "metrics": sorted({h[0].lower() for h in hits})[:20],
                                    "last": [l for l in text.splitlines() if LOG_METRIC_RE.search(l)][-3:]})
    for k in ("libraries", "datasets", "models", "optimizers", "metrics", "checkpoints"):
        inv[k] = dict(inv[k].most_common(40))
    return inv


def to_md(inv):
    L = [f"# Experiment inventory\n\nRoot: `{inv['root']}`\n"]
    for k in ("libraries", "datasets", "models", "optimizers", "metrics"):
        L.append(f"- **{k}**: " + (", ".join(f"{a} ({n})" for a, n in inv[k].items()) or "-"))
    L.append(f"- **checkpoints**: {inv['checkpoints'] or '-'}  |  **tensorboard files**: {len(inv['tensorboard'])}"
             f"  |  **trackers**: {', '.join(inv['trackers']) or '-'}\n")
    L.append("## Code (entry points first)\n")
    for c in sorted(inv["code"] + inv["notebooks"], key=lambda c: (not c["entry_point"], c["path"])):
        L.append(f"### `{c['path']}` ({c['lines']} lines{', ENTRY POINT' if c['entry_point'] else ''})")
        L.append(f"libs: {', '.join(c['ml_libraries']) or '-'}; datasets: {', '.join(c['datasets']) or '-'}; "
                 f"models: {', '.join(c['models']) or '-'}; optimizers: {', '.join(c['optimizers']) or '-'}; "
                 f"metrics: {', '.join(c['metrics']) or '-'}; seeds: {', '.join(c['seeds']) or '-'}")
        if c["cli_args"]:
            L.append("args: " + "; ".join(f"{a['flag']}={a.get('default', '?')}" for a in c["cli_args"]))
        if c["classes"]:
            L.append("classes: " + ", ".join(c["classes"]))
        if c["functions"]:
            L.append("functions: " + ", ".join(c["functions"]))
        for o in c.get("metric_outputs", []):
            L.append(f"    output> {o[:200]}")
        L.append("")
    L.append("## Configs\n")
    for c in inv["configs"]:
        L.append(f"### `{c['path']}`\n```\n" + "\n".join(c["preview"]) + "\n```")
    L.append("\n## Result files\n")
    for r in inv["results"]:
        desc = (f"columns={r['columns']}, rows={r['rows']}" if r["kind"] == "table" else
                f"keys={r['keys']}, rows={r['rows']}" if r["kind"] == "jsonl" else f"shape={json.dumps(r['shape'])[:400]}")
        L.append(f"- `{r['path']}` ({r['kind']}, {r['bytes']} B): {desc}")
    L.append("\n## Logs with metrics\n")
    for g in inv["logs"]:
        L.append(f"- `{g['path']}`: {g['metric_lines']} metric lines, {g['metrics']}; last: {g['last'][-1][:150] if g['last'] else ''}")
    if inv["tensorboard"]:
        L.append("\n## TensorBoard event files\n" + "\n".join(f"- `{t}`" for t in inv["tensorboard"][:50]))
    if inv["figures"]:
        L.append("\n## Existing figures (not reused blindly - regenerate from data)\n" + "\n".join(f"- `{f}`" for f in inv["figures"][:50]))
    L.append("\n## Docs\n")
    for d in inv["docs"]:
        L.append(f"- `{d['path']}`: " + " / ".join(d["first_lines"])[:300])
    return "\n".join(L) + "\n"


def main(argv):
    if "-h" in argv or "--help" in argv:
        print(__doc__); return 0
    out = Path(argv[argv.index("--out") + 1]) if "--out" in argv else Path("paper/inventory")
    pos = [a for i, a in enumerate(argv) if not a.startswith("--") and (i == 0 or argv[i - 1] != "--out")]
    root = Path(pos[0] if pos else ".")
    inv = scan(root)
    out.parent.mkdir(parents=True, exist_ok=True)
    out.with_suffix(".json").write_text(json.dumps(inv, indent=1, ensure_ascii=False, default=list), encoding="utf-8")
    out.with_suffix(".md").write_text(to_md(inv), encoding="utf-8")
    entries = [c["path"] for c in inv["code"] + inv["notebooks"] if c["entry_point"]]
    print(f"wrote {out}.md and {out}.json")
    print(f"code files with ML content: {len(inv['code'])} (+{len(inv['notebooks'])} notebooks); entry points: {', '.join(entries[:12]) or '-'}")
    print(f"configs: {len(inv['configs'])}; result files: {len(inv['results'])}; logs with metrics: {len(inv['logs'])}; "
          f"tensorboard: {len(inv['tensorboard'])}; trackers: {', '.join(inv['trackers']) or '-'}")
    for k in ("libraries", "datasets", "models", "metrics"):
        print(f"{k}: " + (", ".join(list(inv[k])[:12]) or "-"))
    if not inv["results"] and not inv["logs"] and not inv["tensorboard"]:
        print("NOTE: no result files found - the experiments may not have been run yet.")
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))