#!/usr/bin/env python3
"""Rewrite the absolute `graph_path` column in sample_labels_rich.csv.

Why this is needed
------------------
`src/data/04_save_dataset.py:237` writes graph_path as an ABSOLUTE cluster path
(<DATA_ROOT>/.../xxx.pt). After download those paths point
nowhere.  The failure is NON-FATAL, which is what makes it dangerous:
`GraphDataset` (src/train/utils.py) calls `filter_valid_graph_paths()`, which
drops unreadable paths and prints only
    WARNING: Filtered out <n> missing/empty graph files
before carrying on — so a stale column yields a SMALLER dataset instead of an
error.  If every path is stale you get an empty dataset and the crash surfaces
later and unhelpfully, as an IndexError on `train_ds[0]`.  Always read that
warning and check the resulting count.

Note that `generate_embs.py` passes the column to GraphDataset VERBATIM
(lines 430/439) — it never joins --sample_data_folder onto it — so paths are
resolved against the PROCESS WORKING DIRECTORY.  Use `--mode absolute --root
<abs dir>` unless you run from the directory the graphs sit in.

Usage
-----
  # RECOMMENDED: point paths at an absolute root (works with the shipped loader)
  python rebase_graph_paths.py --csv <csv> --mode absolute --root /data/BACH/graphs

  # relative mode: bare filenames. ONLY use if you patch the loader to resolve
  # them against the CSV's directory -- generate_embs.py does NOT.
  python rebase_graph_paths.py --csv <csv> --mode relative

  # check only, change nothing
  python rebase_graph_paths.py --csv <csv> --check

The original file is preserved as <csv>.orig unless --no-backup is given.
Streams row-by-row, so the 2.9 GB TCGA-BRCA CSV is fine.
"""
import argparse
import csv
import os
import shutil
import sys

csv.field_size_limit(10_000_000)


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--csv", required=True, help="path to sample_labels_rich.csv")
    ap.add_argument("--mode", choices=["relative", "absolute"], default="relative")
    ap.add_argument("--root", help="new root dir (required for --mode absolute)")
    ap.add_argument("--check", action="store_true", help="report only, do not modify")
    ap.add_argument("--no-backup", action="store_true")
    args = ap.parse_args()

    if args.mode == "absolute" and not args.root and not args.check:
        sys.exit("[ERR] --mode absolute requires --root")

    src = os.path.abspath(args.csv)
    base = os.path.dirname(src)
    if not os.path.isfile(src):
        sys.exit(f"[ERR] no such file: {src}")

    with open(src, newline="") as f:
        hdr = next(csv.reader(f))
    if "graph_path" not in hdr:
        sys.exit(f"[ERR] no graph_path column; header starts: {hdr[:6]}")
    col = hdr.index("graph_path")

    # --- inspect ---
    # `ok` counts paths that resolve on this machine at all (what --check reports).
    # `already` counts paths that ALSO already have the exact form --mode asks for.
    # The two differ: a bare filename sitting next to its CSV resolves fine, yet is
    # still wrong for `--mode absolute`, because generate_embs.py resolves the column
    # against the PROCESS CWD, not against the CSV's directory. Keying the early exit
    # on `ok` would silently no-op in exactly that case.
    n = ok = already = 0
    sample = None
    with open(src, newline="") as f:
        r = csv.reader(f)
        next(r)
        for row in r:
            if col >= len(row):
                continue
            p = row[col]
            n += 1
            if sample is None:
                sample = p
            cand = p if os.path.isabs(p) else os.path.join(base, p)
            if os.path.exists(cand):
                ok += 1
            if args.mode == "absolute":
                target = os.path.join(args.root, os.path.basename(p)) if args.root else None
                if target and p == target and os.path.exists(target):
                    already += 1
            elif p == os.path.basename(p) and os.path.exists(os.path.join(base, p)):
                already += 1
            if args.check and n >= 20000:
                break
    print(f"rows inspected : {n}")
    print(f"example path   : {sample}")
    print(f"resolvable now : {ok}/{n}" + ("  (first 20k rows only)" if args.check and n >= 20000 else ""))
    if args.check:
        print("check-only; nothing written")
        return
    if already == n and n:
        print(f"all paths already in {args.mode} form and resolvable — nothing to do")
        return

    # --- rewrite ---
    if not args.no_backup and not os.path.exists(src + ".orig"):
        shutil.copy2(src, src + ".orig")
        print(f"backup written : {src}.orig")

    tmp = src + ".tmp"
    changed = 0
    with open(src, newline="") as fin, open(tmp, "w", newline="") as fout:
        r, w = csv.reader(fin), csv.writer(fout)
        w.writerow(next(r))
        for row in r:
            if col < len(row) and row[col]:
                name = os.path.basename(row[col])
                row[col] = name if args.mode == "relative" else os.path.join(args.root, name)
                changed += 1
            w.writerow(row)
    os.replace(tmp, src)
    print(f"rows rewritten : {changed}  (mode={args.mode}"
          + (f", root={args.root}" if args.mode == "absolute" else ", relative to the CSV's directory") + ")")
    print("NOTE: with --mode relative the loader must resolve paths against the CSV's directory;\n"
          "      pass --mode absolute --root <dir> if you cannot patch the loader.")


if __name__ == "__main__":
    main()