#!/usr/bin/env python3
"""Assemble the anonymized supplementary archive for the paper.

The archive contains the experiment code, pinned dependency versions, every
row-level CSV referenced by the paper's artifact map, the generated LaTeX
table fragments, the figure files, a machine readable seed manifest, and a
SHA256 manifest over every packaged file. Research notes, draft papers, raw
datasets, and git history are excluded.
"""

from __future__ import annotations

import argparse
import csv
import glob
import hashlib
import json
import shutil
import zipfile
from pathlib import Path

CODE_FILES = ["scripts", "requirements.txt", "requirements-lock.txt"]

OUTPUT_GLOBS = [
    "outputs/aaai_depth_experiments/*.csv",
    "outputs/aaai_depth_experiments/summary.json",
    "outputs/cnn_initialization_validation/*",
    "outputs/real_data_validation_mnist/*.csv",
    "outputs/real_data_validation_mnist/summary.json",
    "outputs/soft_erosion_theory_vs_empirical/*.csv",
    "outputs/learning_rate_compensation/*",
    "outputs/parameter_matched_depth_control/*",
    "outputs/teacher_test_gap_sgd_T8000/runs.csv",
    "outputs/teacher_test_gap_sgd_T8000/width_summary.csv",
    "outputs/teacher_test_gap_sgd_T8000/summary.json",
    "outputs/operator_stress_grid_summary/*.csv",
    "outputs/compressed_operator_s20_256traj_T50_width64_N*/*",
    "outputs/compressed_operator_retangent_T50_width64_N*/*",
    "outputs/compressed_operator_T50_width64_N*/*",
    "outputs/early_directional_predictors_T50_width64_N*/*",
    "outputs/early_kernel_predictors_T50_N128_32runs/*",
    "outputs/finite_time_kernel_probe_T50_N128_8runs/*.csv",
    "outputs/finite_time_kernel_probe_T50_N128_8runs/summary.json",
    "outputs/fa_tangent_hierarchy_derivative_probe_eps005/*",
    "outputs/aaai_main_figures/*",
    "outputs/aaai_supplementary/*",
    "outputs/tables/*.tex",
]

SEED_SOURCES = [
    ("depth_initialization", "outputs/aaai_depth_experiments/initialization_rows.csv"),
    ("depth_finite_time", "outputs/aaai_depth_experiments/finite_time_rows.csv"),
    ("cnn_initialization", "outputs/cnn_initialization_validation/cnn_initialization_rows.csv"),
    ("mnist_initialization", "outputs/real_data_validation_mnist/e0_mnist_rows.csv"),
    ("mnist_estimator", "outputs/real_data_validation_mnist/estimator_mnist_rows.csv"),
    ("learning_rate_evaluation", "outputs/learning_rate_compensation/evaluation_rows.csv"),
    ("learning_rate_tuning", "outputs/learning_rate_compensation/tuning_rows.csv"),
    ("parameter_matched", "outputs/parameter_matched_depth_control/parameter_matched_rows.csv"),
    ("teacher_task", "outputs/teacher_test_gap_sgd_T8000/runs.csv"),
    ("stress_grid", "outputs/operator_stress_grid_summary/stress_grid_rows.csv"),
]

README = """\
# Supplementary code and data

This archive accompanies the anonymous submission "The Optimization Cost of
Fixed Random Feedback". It contains the experiment code, the row-level CSV
outputs behind every reported number, the generated LaTeX table fragments,
the figure files, a seed manifest, and a SHA256 manifest.

## Environment

Python 3.13, CPU, double precision. Pinned versions in
`requirements-lock.txt`. Install with:

    pip install -r requirements-lock.txt

## Reproduction

Appendix K of the paper lists the exact command for every experiment.
The shared self tests are:

    python scripts/test_feedback_rules.py
    python scripts/cnn_initialization_validation.py --self-test

## Figure and table map

| Paper object | Generating script | Source data |
|---|---|---|
| Figures 1-2 | scripts/plot_aaai_main_figures.py | outputs/aaai_depth_experiments/, outputs/soft_erosion_theory_vs_empirical/, outputs/cnn_initialization_validation/, outputs/real_data_validation_mnist/ |
| Figure 3 | scripts/cnn_initialization_validation.py | outputs/cnn_initialization_validation/ |
| Figure 4 | scripts/finite_time_supplement.py | outputs/aaai_depth_experiments/finite_time_rows.csv |
| Figure 5 | scripts/learning_rate_compensation.py | outputs/learning_rate_compensation/ |
| Figure 6 | scripts/parameter_matched_depth_control.py | outputs/parameter_matched_depth_control/ |
| Figure 7 | scripts/downstream_capacity_sweep.py + scripts/plot_teacher_test_gap.py | outputs/teacher_test_gap_sgd_T8000/ |
| Generated appendix tables | scripts/generate_paper_tables.py, scripts/finite_time_supplement.py | outputs/tables/*.tex |

## Manifests

- `artifacts/seed_manifest.csv`: every (experiment, init seed, feedback seed)
  pair present in the row-level outputs, with the shared data seed per
  experiment where recorded.
- `artifacts/manifest.json`: byte size and SHA256 for every file in this
  archive (excluding the manifest itself).

The CSV files are the authoritative source for all reported experimental
numbers; no table value in the paper is transcribed manually.
"""


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--outdir", type=Path, default=Path("supplement_build"))
    return parser.parse_args()


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1 << 20), b""):
            digest.update(chunk)
    return digest.hexdigest()


def seed_manifest(root: Path) -> list[dict[str, str]]:
    rows: list[dict[str, str]] = []
    for experiment, rel in SEED_SOURCES:
        path = Path(rel)
        if not path.exists():
            continue
        data_seed = ""
        summary = path.parent / "summary.json"
        if summary.exists():
            config = json.load(summary.open()).get("config", {})
            data_seed = str(config.get("data_seed", ""))
        seen: set[tuple[str, str]] = set()
        for row in csv.DictReader(path.open()):
            init_seed = row.get("init_seed", "")
            feedback_seed = row.get("feedback_seed", "")
            key = (init_seed, feedback_seed)
            if key in seen:
                continue
            seen.add(key)
            rows.append(
                {
                    "experiment": experiment,
                    "data_seed": data_seed,
                    "init_seed": init_seed,
                    "feedback_seed": feedback_seed,
                }
            )
    return rows


def main() -> None:
    args = parse_args()
    stage = args.outdir / "supplement"
    if stage.exists():
        shutil.rmtree(stage)
    stage.mkdir(parents=True)

    for item in CODE_FILES:
        src = Path(item)
        if src.is_dir():
            shutil.copytree(src, stage / src.name)
        elif src.exists():
            shutil.copy2(src, stage / src.name)

    for pattern in OUTPUT_GLOBS:
        for match in sorted(glob.glob(pattern)):
            src = Path(match)
            if not src.is_file():
                continue
            dest = stage / src
            dest.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(src, dest)

    (stage / "README.md").write_text(README)

    artifacts = stage / "artifacts"
    artifacts.mkdir(exist_ok=True)
    seeds = seed_manifest(Path("."))
    with (artifacts / "seed_manifest.csv").open("w", newline="") as handle:
        writer = csv.DictWriter(
            handle, fieldnames=["experiment", "data_seed", "init_seed", "feedback_seed"]
        )
        writer.writeheader()
        writer.writerows(seeds)

    manifest = {
        str(path.relative_to(stage)): {"bytes": path.stat().st_size, "sha256": sha256(path)}
        for path in sorted(stage.rglob("*"))
        if path.is_file() and path.name != "manifest.json"
    }
    (artifacts / "manifest.json").write_text(json.dumps(manifest, indent=1))

    zip_path = args.outdir / "supplement.zip"
    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as archive:
        for path in sorted(stage.rglob("*")):
            if path.is_file():
                archive.write(path, path.relative_to(args.outdir))

    total = sum(p.stat().st_size for p in stage.rglob("*") if p.is_file())
    print(f"files: {len(manifest) + 1}")
    print(f"seed rows: {len(seeds)}")
    print(f"staged bytes: {total/1e6:.1f} MB")
    print(f"zip: {zip_path} ({zip_path.stat().st_size/1e6:.1f} MB)")


if __name__ == "__main__":
    main()