hparam_search.py

#!/usr/bin/env python3
"""
hparam_search.py  –  Hyperparameter search for AF3CRAFT gradient binder design
================================================================================

Two evaluation modes
--------------------
PROXY mode (default, fast)
    Objective = gradient-phase loss + interface contact.
    No diffusion; ~1 min/trial after JIT warm-up.  Good for broad exploration.

FULL mode (--full_eval, accurate)
    Every trial runs: gradient design → sequence → full AF3 diffusion → ipTM + pLDDT.
    Score = ipTM + 0.5 * (binder_pLDDT / 100)   (both higher = better)
    Averaged across --n_seeds seeds per config.
    Slower (~3–5 min/trial) but directly optimises the metrics you care about.
    Recommended for Bayesian search with 20–30 trials.

Usage examples
--------------
# Fast proxy search (exploration):
python hparam_search.py \\
    --target_pdb    examples/pdl1.pdb \\
    --target_chain  A \\
    --binder_length 20 \\
    --search_mode   random --n_trials 20 --n_seeds 3 \\
    --output_dir    hparam_results/proxy

# Full ipTM+pLDDT Bayesian search (recommended):
python hparam_search.py \\
    --target_pdb    examples/pdl1.pdb \\
    --target_chain  A \\
    --binder_length 20 \\
    --search_mode   bayesian --n_trials 25 --n_seeds 2 \\
    --full_eval \\
    --output_dir    hparam_results/full

# Ligand target:
python hparam_search.py \\
    --ligand_ccd    SAM \\
    --binder_length 20 \\
    --search_mode   bayesian --n_trials 25 --n_seeds 2 \\
    --full_eval \\
    --output_dir    hparam_results/sam

Notes
-----
- AF3 gradient function is JIT-compiled once and reused across all trials.
- All results → {output_dir}/results.csv
- Best config   → {output_dir}/best_config.json  + ready-to-paste CLI flags
- --full_eval results also include per-trial .cif structure files.
"""

import sys
import os

os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false")
os.environ.setdefault("XLA_PYTHON_CLIENT_ALLOCATOR", "platform")
os.environ.setdefault("CUDA_VISIBLE_DEVICES", "1")   # use GPU cuda:1

import alphafold3.cpp  # must be first AF3 import
sys.path.insert(0, os.path.expanduser("~/alphafold3/src"))

import argparse
import csv
import itertools
import json
import pathlib
import random
import time
from copy import deepcopy
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Any

import jax
import jax.numpy as jnp
import numpy as np

from gradient_utils.af3_gradient import (
    build_evoformer_distogram_fn,
    design_3stage,
    logits_to_sequence,
    prepare_batch_dict,
)

# ─────────────────────────────────────────────────────────────────────────────
# Hyperparameter search space
# ─────────────────────────────────────────────────────────────────────────────

# Each entry: (name, type, grid_values, random_low, random_high, log_scale)
# log_scale=True → sample uniformly in log space for random/Bayesian search.
HPARAM_SPACE: List[Dict] = [
    # ── Stage iterations ─────────────────────────────────────────────────────
    dict(name="soft_iters",    type="int",   grid=[80, 100, 120, 150],
         low=80,   high=150, log=False),
    dict(name="temp_iters",    type="int",   grid=[40, 60, 80, 100],
         low=40,   high=100, log=False),
    dict(name="hard_iters",    type="int",   grid=[0, 5, 10],
         low=0,    high=10,  log=False),
    dict(name="warmup_iters",  type="int",   grid=[0, 10, 20, 30],
         low=0,    high=30,  log=False),

    # ── Optimizer & LR  (SGD only) ────────────────────────────────────────────
    dict(name="optimizer_type",type="cat",   grid=["sgd"],
         choices=["sgd"]),
    dict(name="learning_rate", type="float", grid=[0.05, 0.1, 0.2, 0.5],
         low=0.05, high=0.5, log=True),

    # ── Soft / temperature schedule ──────────────────────────────────────────
    dict(name="e_soft",        type="float", grid=[0.8, 0.9, 1.0],
         low=0.8,  high=1.0, log=False),
    dict(name="warmup_alpha",  type="float", grid=[1.0, 1.5, 2.0],
         low=1.0,  high=2.0, log=False),

    # ── Loss weights ─────────────────────────────────────────────────────────
    dict(name="helix_weight",  type="float", grid=[-0.3, -0.2, -0.1, 0.0],
         low=-0.3, high=0.0, log=False),
    dict(name="conf_weight",   type="float", grid=[0.0, 0.05, 0.1, 0.2],
         low=0.0,  high=0.3, log=False),

    # ── Progressive num_pos ────────────────────────────────────────────────
    dict(name="num_pos_start", type="int",   grid=[1, 4, None],
         low=1,    high=8,   log=False, allow_none=True),
]

# Reduced grid for --search_mode grid (combinatorial explosion guard):
# Only vary the most impactful params; others stay at their first grid value.
GRID_FOCUS_PARAMS = [
    "soft_iters", "temp_iters", "learning_rate",
    "optimizer_type", "helix_weight", "conf_weight",
]


# ─────────────────────────────────────────────────────────────────────────────
# Objective function
# ─────────────────────────────────────────────────────────────────────────────

def _run_design_3stage(config, batch_dict, binder_start_idx, binder_length,
                       af3_fn, seed, verbose):
    """Run design_3stage for one seed. Returns (final_logits, history)."""
    return design_3stage(
        batch_dict=batch_dict,
        binder_start_idx=binder_start_idx,
        binder_length=binder_length,
        af3_fn=af3_fn,
        soft_iters=config.get("soft_iters", 80),
        temp_iters=config.get("temp_iters", 0),
        hard_iters=config.get("hard_iters", 0),
        e_soft=config.get("e_soft", 1.0),
        warmup_iters=config.get("warmup_iters", 0),
        warmup_alpha=config.get("warmup_alpha", 2.0),
        learning_rate=config.get("learning_rate", 0.2),
        optimizer_type=config.get("optimizer_type", "sgd"),
        helix_weight=config.get("helix_weight", -0.2),
        conf_weight=config.get("conf_weight", 0.1),
        num_pos_start=config.get("num_pos_start", 1),
        rng_seed=seed,
        verbose=verbose,
        log_every=999,  # suppress per-step output during search
    )


def evaluate_config_proxy(
    config: Dict[str, Any],
    batch_dict: Dict,
    binder_start_idx: int,
    binder_length: int,
    af3_fn,
    seeds: List[int],
    verbose: bool = False,
) -> Dict[str, float]:
    """
    PROXY objective — gradient phase only, no diffusion.

    Score (higher = better):
        score = -mean_loss  +  0.5 * mean_contact  -  0.2 * std_loss

    Fast (~1 min/trial after JIT warm-up).  Good for broad exploration.
    """
    all_losses:   List[float] = []
    all_contacts: List[float] = []
    all_con_loss: List[float] = []
    all_i_con:    List[float] = []

    for seed in seeds:
        try:
            final_logits, history = _run_design_3stage(
                config, batch_dict, binder_start_idx, binder_length,
                af3_fn, seed, verbose,
            )
            # Use BEST step in history (not last — design can overshoot)
            best_step = min(history, key=lambda s: s.get("total_loss", float("inf")))
            all_losses.append(float(best_step.get("total_loss", float("inf"))))
            all_contacts.append(float(best_step.get("mean_interface_contact", 0.0)))
            all_con_loss.append(float(best_step.get("con_loss", float("inf"))))
            all_i_con.append(float(best_step.get("i_con_loss", float("inf"))))
        except Exception as exc:
            print(f"    [EVAL] seed={seed} FAILED: {exc}")
            all_losses.append(float("inf"))
            all_contacts.append(0.0)
            all_con_loss.append(float("inf"))
            all_i_con.append(float("inf"))

    mean_loss    = float(np.mean(all_losses))
    mean_contact = float(np.mean(all_contacts))
    std_loss     = float(np.std(all_losses))
    best_loss    = float(np.min(all_losses))

    # Combined proxy score (higher = better):
    #  -mean_loss:    lower gradient loss → higher score
    #  +0.5*contact:  more interface contact → higher score
    #  -0.2*std_loss: penalise variance across seeds (stability)
    score = -mean_loss + 0.5 * mean_contact - 0.2 * std_loss

    return {
        "eval_mode":    "proxy",
        "score":        score,
        "mean_loss":    mean_loss,
        "best_loss":    best_loss,
        "std_loss":     std_loss,
        "mean_contact": mean_contact,
        "mean_con":     float(np.mean(all_con_loss)),
        "mean_i_con":   float(np.mean(all_i_con)),
        # ipTM/pLDDT are not available in proxy mode — add --full_eval to get them
    }


def evaluate_config_full(
    config: Dict[str, Any],
    batch_dict: Dict,
    binder_start_idx: int,
    binder_length: int,
    af3_fn,
    seeds: List[int],
    model_runner,
    args,
    trial_id: int,
    output_dir: Path,
    verbose: bool = False,
) -> Dict[str, float]:
    """
    FULL objective — gradient phase + full AF3 diffusion per seed.

    Score (higher = better):
        score = mean_ipTM  +  0.5 * (mean_pLDDT / 100)
                            -  0.3 * std_ipTM   (penalise instability)

    Both ipTM and pLDDT are from the actual AF3 diffusion prediction,
    so this is the ground-truth objective for binder quality.
    Slower (~3–5 min/trial) but directly optimises what you care about.
    """
    from bindcraft_af3_gradient import run_full_af3, _ensure_pipeline_imported
    _ensure_pipeline_imported()

    all_iptm:   List[float] = []
    all_plddt:  List[float] = []
    all_losses: List[float] = []
    all_contacts: List[float] = []

    for s_idx, seed in enumerate(seeds):
        try:
            # ── gradient phase ───────────────────────────────────────────────
            final_logits, history = _run_design_3stage(
                config, batch_dict, binder_start_idx, binder_length,
                af3_fn, seed, verbose,
            )
            best_step = min(history, key=lambda s: s.get("total_loss", float("inf")))
            all_losses.append(float(best_step.get("total_loss", float("inf"))))
            all_contacts.append(float(best_step.get("mean_interface_contact", 0.0)))

            sequence = logits_to_sequence(
                final_logits[binder_start_idx:binder_start_idx + binder_length, :20]
            )
            print(f"    [FULL seed={seed}] seq={sequence}  running AF3...")

            # ── full AF3 diffusion ────────────────────────────────────────────
            val_dir = output_dir / f"trial{trial_id:03d}_seed{seed}"
            val_metrics = run_full_af3(
                binder_sequence=sequence,
                target_pdb=args.target_pdb if args.target_pdb else None,
                target_chain_id=args.target_chain if args.target_chain else None,
                model_runner=model_runner,
                output_dir=val_dir,
                iteration=0,
                buckets=args.buckets,
                conformer_max_iterations=args.conformer_max_iterations,
                rng_seed=seed,
                ligand_ccd=args.ligand_ccd if args.ligand_ccd else None,
            )
            iptm  = float(val_metrics.get("iptm", 0.0))
            plddt = float(val_metrics.get("binder_plddt_mean", 0.0))
            all_iptm.append(iptm)
            all_plddt.append(plddt)
            print(
                f"    [FULL seed={seed}]  ipTM={iptm:.3f}  "
                f"pLDDT={plddt:.1f}  pTM={val_metrics.get('ptm',0):.3f}"
            )

        except Exception as exc:
            print(f"    [FULL seed={seed}] FAILED: {exc}")
            all_iptm.append(0.0)
            all_plddt.append(0.0)
            all_losses.append(float("inf"))
            all_contacts.append(0.0)

    mean_iptm  = float(np.mean(all_iptm))
    mean_plddt = float(np.mean(all_plddt))
    std_iptm   = float(np.std(all_iptm))
    best_iptm  = float(np.max(all_iptm))
    best_plddt = float(np.max(all_plddt))

    # Primary score:
    #   ipTM in [0,1] — interface quality (most important)
    #   pLDDT in [0,100] / 100 → [0,1] — binder folding confidence
    #   std_iptm penalty: reward configs that are consistently good
    score = mean_iptm + 0.5 * (mean_plddt / 100.0) - 0.3 * std_iptm

    return {
        "eval_mode":    "full",
        "score":        score,
        "mean_iptm":    mean_iptm,
        "mean_plddt":   mean_plddt,
        "best_iptm":    best_iptm,
        "best_plddt":   best_plddt,
        "std_iptm":     std_iptm,
        "mean_loss":    float(np.mean(all_losses)),
        "best_loss":    float(np.min(all_losses)) if all_losses else float("inf"),
        "std_loss":     float(np.std(all_losses)),
        "mean_contact": float(np.mean(all_contacts)),
        "mean_con":     float("nan"),
        "mean_i_con":   float("nan"),
    }


def evaluate_config(
    config: Dict[str, Any],
    batch_dict: Dict,
    binder_start_idx: int,
    binder_length: int,
    af3_fn,
    seeds: List[int],
    verbose: bool = False,
    # full eval
    full_eval: bool = False,
    model_runner=None,
    args=None,
    trial_id: int = 0,
    output_dir: Optional[Path] = None,
) -> Dict[str, float]:
    """Dispatch to proxy or full evaluator based on full_eval flag."""
    if full_eval:
        return evaluate_config_full(
            config, batch_dict, binder_start_idx, binder_length,
            af3_fn, seeds, model_runner, args, trial_id, output_dir, verbose,
        )
    else:
        return evaluate_config_proxy(
            config, batch_dict, binder_start_idx, binder_length,
            af3_fn, seeds, verbose,
        )


# ─────────────────────────────────────────────────────────────────────────────
# Search strategies
# ─────────────────────────────────────────────────────────────────────────────

def _random_sample_config(rng: random.Random) -> Dict[str, Any]:
    """Sample one random configuration from HPARAM_SPACE."""
    config = {}
    for hp in HPARAM_SPACE:
        if hp["type"] == "cat":
            config[hp["name"]] = rng.choice(hp["choices"])
        elif hp.get("allow_none") and rng.random() < 0.15:
            config[hp["name"]] = None
        elif hp["type"] == "int":
            if hp.get("log"):
                v = int(round(np.exp(rng.uniform(np.log(max(hp["low"], 1)),
                                                  np.log(max(hp["high"], 1))))))
            else:
                v = rng.randint(int(hp["low"]), int(hp["high"]))
            config[hp["name"]] = v
        elif hp["type"] == "float":
            if hp.get("log"):
                lo, hi = hp["low"], hp["high"]
                lo = max(lo, 1e-8)  # avoid log(0)
                v = float(np.exp(rng.uniform(np.log(lo), np.log(hi))))
            else:
                v = rng.uniform(hp["low"], hp["high"])
            config[hp["name"]] = v
    return config


def _grid_configs() -> List[Dict[str, Any]]:
    """Build a grid over GRID_FOCUS_PARAMS; other params use first grid value."""
    # Base config: first grid value for every param
    base = {}
    for hp in HPARAM_SPACE:
        base[hp["name"]] = hp["grid"][0]

    # Build axes for focused params only
    axes = {}
    for hp in HPARAM_SPACE:
        if hp["name"] in GRID_FOCUS_PARAMS:
            axes[hp["name"]] = hp["grid"]

    keys   = list(axes.keys())
    values = list(axes.values())
    configs = []
    for combo in itertools.product(*values):
        cfg = deepcopy(base)
        for k, v in zip(keys, combo):
            cfg[k] = v
        configs.append(cfg)
    return configs


def _bayesian_search(
    n_trials: int,
    batch_dict,
    binder_start_idx,
    binder_length,
    af3_fn,
    seeds,
    results: List[Dict],
    csv_path: Path,
    verbose: bool,
    full_eval: bool = False,
    model_runner=None,
    args=None,
    output_dir: Optional[Path] = None,
):
    """Optuna-based Bayesian search (TPE sampler)."""
    try:
        import optuna
        optuna.logging.set_verbosity(optuna.logging.WARNING)
    except ImportError:
        raise ImportError(
            "Optuna is required for Bayesian search.  "
            "Install with:  pip install optuna"
        )

    def objective(trial: "optuna.Trial") -> float:
        config = {}
        for hp in HPARAM_SPACE:
            n = hp["name"]
            if hp["type"] == "cat":
                config[n] = trial.suggest_categorical(n, hp["choices"])
            elif hp.get("allow_none"):
                use_none = trial.suggest_categorical(f"{n}_is_none", [True, False])
                if use_none:
                    config[n] = None
                else:
                    config[n] = trial.suggest_int(n, int(hp["low"]), int(hp["high"]))
            elif hp["type"] == "int":
                config[n] = trial.suggest_int(n, int(hp["low"]), int(hp["high"]),
                                               log=hp.get("log", False))
            elif hp["type"] == "float":
                lo = max(hp["low"], 1e-8) if hp.get("log") else hp["low"]
                config[n] = trial.suggest_float(n, lo, hp["high"],
                                                log=hp.get("log", False))

        t0 = time.time()
        metrics = evaluate_config(
            config, batch_dict, binder_start_idx, binder_length,
            af3_fn, seeds, verbose=verbose,
            full_eval=full_eval, model_runner=model_runner,
            args=args, trial_id=trial.number, output_dir=output_dir,
        )
        elapsed = time.time() - t0
        row = {**config, **metrics, "trial_id": trial.number, "elapsed_s": elapsed}
        results.append(row)
        _append_csv(csv_path, row)
        _print_trial(trial.number, config, metrics, elapsed)
        # Optuna minimises; negate our score
        return -metrics["score"]

    study = optuna.create_study(
        direction="minimize",
        sampler=optuna.samplers.TPESampler(seed=42),
    )
    study.optimize(objective, n_trials=n_trials)


# ─────────────────────────────────────────────────────────────────────────────
# CSV helpers
# ─────────────────────────────────────────────────────────────────────────────

_CSV_FIELDNAMES: Optional[List[str]] = None

# Canonical metric columns for each eval mode — order matters for readability.
_PROXY_METRIC_COLS = [
    "eval_mode", "score",
    "mean_loss", "best_loss", "std_loss",
    "mean_contact", "mean_con", "mean_i_con",
]
_FULL_METRIC_COLS = [
    "eval_mode", "score",
    "mean_iptm", "std_iptm", "best_iptm",
    "mean_plddt", "best_plddt",
    "mean_loss", "best_loss", "std_loss",
    "mean_contact", "mean_con", "mean_i_con",
]

def _init_csv_fieldnames(full_eval: bool) -> None:
    """Call once at startup so the CSV header is always correct regardless of
    which eval mode runs first.  Must be called before any _append_csv call."""
    global _CSV_FIELDNAMES
    hparam_cols = [hp["name"] for hp in HPARAM_SPACE]
    metric_cols = _FULL_METRIC_COLS if full_eval else _PROXY_METRIC_COLS
    _CSV_FIELDNAMES = hparam_cols + metric_cols + ["trial_id", "elapsed_s"]


def _append_csv(csv_path: Path, row: Dict):
    global _CSV_FIELDNAMES
    write_header = not csv_path.exists()
    # Fallback: if _init_csv_fieldnames was never called, build from first row.
    if _CSV_FIELDNAMES is None:
        _CSV_FIELDNAMES = list(row.keys())
    with open(csv_path, "a", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=_CSV_FIELDNAMES,
                                extrasaction="ignore", restval=float("nan"))
        if write_header:
            writer.writeheader()
        writer.writerow(row)


def _print_trial(trial_id: int, config: Dict, metrics: Dict, elapsed: float):
    mode = metrics.get("eval_mode", "proxy")
    if mode == "full":
        quality = (
            f"ipTM={metrics['mean_iptm']:.3f}±{metrics['std_iptm']:.3f}  "
            f"pLDDT={metrics['mean_plddt']:.1f}  "
            f"best_ipTM={metrics['best_iptm']:.3f}"
        )
    else:
        quality = (
            f"loss={metrics['mean_loss']:.4f}±{metrics['std_loss']:.4f}  "
            f"contact={metrics['mean_contact']:.3f}"
        )
    print(
        f"  [trial {trial_id:3d}]  "
        f"score={metrics['score']:+.4f}  "
        + quality +
        f"  t={elapsed:.0f}s  "
        + "  ".join(
            f"{k}={v:.3g}" if isinstance(v, float) else f"{k}={v}"
            for k, v in list(config.items())[:5]
        )
    )


# ─────────────────────────────────────────────────────────────────────────────
# Optional: full AF3 validation on top-K configs
# ─────────────────────────────────────────────────────────────────────────────

def _run_topk_validation(
    results: List[Dict],
    k: int,
    args,
    batch_dict,
    binder_start_idx: int,
    binder_length: int,
    af3_fn,
    output_dir: Path,
):
    """Re-run the top-K configs with full AF3 diffusion and record ipTM."""
    # lazy import to avoid GLIBC issues at module load
    from pipeline import Input, make_model_config, ModelRunner, predict_structure, write_outputs
    from bindcraft_af3_gradient import run_full_af3, _ensure_pipeline_imported

    sorted_results = sorted(results, key=lambda r: r.get("score", float("-inf")),
                            reverse=True)
    topk = sorted_results[:k]

    print(f"\n{'='*70}")
    print(f"  TOP-{k} VALIDATION  –  Full AF3 (diffusion)")
    print(f"{'='*70}")

    gpu_devices = jax.devices("gpu")
    jax_device = gpu_devices[args.gpu] if gpu_devices else jax.devices()[0]
    model_runner = ModelRunner(
        config=make_model_config(
            flash_attention_implementation=args.flash_attention_implementation,
            num_diffusion_samples=args.num_diffusion_samples,
            num_recycles=args.num_hard_recycles,
            return_embeddings=False,
        ),
        device=jax_device,
        model_dir=pathlib.Path(args.model_dir),
    )
    _ensure_pipeline_imported()

    validation_rows = []
    for rank, cfg_row in enumerate(topk):
        # Re-run gradient design with this config to get the sequence
        config = {hp["name"]: cfg_row.get(hp["name"]) for hp in HPARAM_SPACE}
        seed = args.seed if args.seed is not None else 42
        try:
            final_logits, _ = design_3stage(
                batch_dict=batch_dict,
                binder_start_idx=binder_start_idx,
                binder_length=binder_length,
                af3_fn=af3_fn,
                rng_seed=seed,
                verbose=False,
                log_every=999,
                **config,
            )
            sequence = logits_to_sequence(
                final_logits[binder_start_idx:binder_start_idx + binder_length, :20]
            )
            val_dir = output_dir / f"topk_rank{rank+1:02d}"
            val_metrics = run_full_af3(
                binder_sequence=sequence,
                target_pdb=args.target_pdb if args.target_pdb else None,
                target_chain_id=args.target_chain if args.target_chain else None,
                model_runner=model_runner,
                output_dir=val_dir,
                iteration=0,
                buckets=args.buckets,
                conformer_max_iterations=args.conformer_max_iterations,
                rng_seed=seed,
                ligand_ccd=args.ligand_ccd if args.ligand_ccd else None,
            )
            print(
                f"  [rank {rank+1}]  seq={sequence}  "
                f"ipTM={val_metrics['iptm']:.3f}  "
                f"pTM={val_metrics['ptm']:.3f}  "
                f"score={val_metrics['ranking_score']:.3f}"
            )
            validation_rows.append({
                "rank": rank + 1,
                "sequence": sequence,
                **val_metrics,
                **config,
            })
        except Exception as exc:
            print(f"  [rank {rank+1}]  FAILED: {exc}")

    # Save validation CSV
    if validation_rows:
        val_csv = output_dir / "topk_validation.csv"
        with open(val_csv, "w", newline="") as f:
            writer = csv.DictWriter(f, fieldnames=list(validation_rows[0].keys()),
                                    extrasaction="ignore")
            writer.writeheader()
            writer.writerows(validation_rows)
        print(f"\n  Validation results → {val_csv}")


# ─────────────────────────────────────────────────────────────────────────────
# Argument parser
# ─────────────────────────────────────────────────────────────────────────────

def parse_arguments():
    p = argparse.ArgumentParser(
        description="Hyperparameter search for AF3CRAFT gradient binder design"
    )
    # ── Target ────────────────────────────────────────────────────────────────
    p.add_argument("--target_pdb",    default="",
                   help="Target PDB file (optional if --ligand_ccd given)")
    p.add_argument("--target_chain",  default="",
                   help="Target chain ID")
    p.add_argument("--binder_length", type=int, required=True,
                   help="Number of residues to design")
    p.add_argument("--ligand_ccd",    default="",
                   help="Ligand CCD code (optional)")

    # ── Search settings ───────────────────────────────────────────────────────
    p.add_argument("--search_mode",   choices=["random", "grid", "bayesian"],
                   default="random",
                   help="Search strategy: random, grid, or bayesian (Optuna)")
    p.add_argument("--n_trials",      type=int, default=20,
                   help="Number of trials (random/Bayesian only)")
    p.add_argument("--n_seeds",       type=int, default=3,
                   help="Seeds per config (higher = more reliable estimate)")
    p.add_argument("--base_seed",     type=int, default=0,
                   help="Base seed; n_seeds seeds = base_seed … base_seed+n_seeds-1")

    # ── Evaluation mode ───────────────────────────────────────────────────────
    p.add_argument("--full_eval", action="store_true", default=True,
                   help=(
                       "Run full AF3 diffusion for EVERY trial and optimise directly "
                       "on ipTM + pLDDT.  Slower but uses the real objective. "
                       "Recommended for Bayesian search with 20-30 trials. (default: True)"
                   ))
    p.add_argument("--proxy_eval", dest="full_eval", action="store_false",
                   help="Use fast proxy objective (gradient loss only). Overrides --full_eval.")

    # ── Output ────────────────────────────────────────────────────────────────
    p.add_argument("--output_dir",    default="./hparam_results")
    p.add_argument("--validate_topk", type=int, default=0,
                   help="After search, run full AF3 validation on top-K configs (0=skip)")

    # ── AF3 model ─────────────────────────────────────────────────────────────
    p.add_argument("--model_dir",
                   default=os.path.expanduser("/home/jupyter-yehlin/alphafold3/models"))
    p.add_argument("--num_recycles",  type=int, default=0,
                   help="Evoformer recycles for gradient phase")
    p.add_argument("--flash_attention_implementation",
                   choices=["triton", "cudnn", "xla"], default="cudnn")
    p.add_argument("--num_diffusion_samples", type=int, default=1)
    p.add_argument("--num_hard_recycles",     type=int, default=3)
    p.add_argument("--gpu",  type=int, default=0,
                   help="JAX GPU index after CUDA_VISIBLE_DEVICES filtering (default: 0 = cuda:1)")
    p.add_argument("--seed", type=int, default=None)
    p.add_argument("--buckets", nargs="+", type=int,
                   default=[256, 512, 768, 1024, 1280, 1536, 2048, 2560, 3072, 3584, 4096])
    p.add_argument("--conformer_max_iterations", type=int, default=None)

    p.add_argument("--verbose", action="store_true",
                   help="Print per-step gradient output for each trial")
    return p.parse_args()


# ─────────────────────────────────────────────────────────────────────────────
# Main
# ─────────────────────────────────────────────────────────────────────────────

def main():
    args = parse_arguments()

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    csv_path = output_dir / "results.csv"

    # Initialise CSV column layout BEFORE any trials run so the header is always
    # correct even if the first few trials fail or are from a different eval mode.
    # If an old CSV with incompatible columns exists, rename it out of the way.
    _init_csv_fieldnames(args.full_eval)
    if csv_path.exists():
        import csv as _csv_mod
        try:
            with open(csv_path, newline="") as _f:
                existing_cols = next(_csv_mod.reader(_f), [])
            if set(existing_cols) != set(_CSV_FIELDNAMES):
                stale = csv_path.with_suffix(".stale.csv")
                csv_path.rename(stale)
                print(
                    f"  [CSV] Existing results.csv has different columns "
                    f"→ renamed to {stale.name} to avoid column misalignment."
                )
        except Exception:
            pass  # can't inspect → leave file as-is

    seeds = list(range(args.base_seed, args.base_seed + args.n_seeds))

    print("\n" + "=" * 70)
    print("  AF3CRAFT  –  HYPERPARAMETER SEARCH")
    print("=" * 70)
    print(f"  Mode        : {args.search_mode}")
    print(f"  Eval        : {'FULL (ipTM + pLDDT)' if args.full_eval else 'PROXY (gradient loss)  ← pass --no-full_eval to use fast proxy'}")
    if args.search_mode in ("random", "bayesian"):
        print(f"  Trials      : {args.n_trials}")
    print(f"  Seeds/trial : {args.n_seeds}  {seeds}")
    if args.target_pdb:
        print(f"  Target      : {args.target_pdb} (chain {args.target_chain})")
    if args.ligand_ccd:
        print(f"  Ligand      : {args.ligand_ccd}")
    print(f"  Binder len  : {args.binder_length}")
    print(f"  Output      : {csv_path}")
    print("=" * 70 + "\n")

    # ── Step 1: Featurise (shared across all trials) ──────────────────────────
    print("[SETUP] Featurising input...")
    batch_dict, binder_start_idx, binder_length, target_length = prepare_batch_dict(
        target_pdb=args.target_pdb if args.target_pdb else None,
        target_chain=args.target_chain if args.target_chain else None,
        binder_length=args.binder_length,
        ligand_ccd=args.ligand_ccd if args.ligand_ccd else None,
    )
    print(f"[SETUP] ✓  binder_start={binder_start_idx}  binder_len={binder_length}  "
          f"target_len={target_length}")

    # ── Step 2: Build AF3 gradient fn ─────────────────────────────────────────
    print("[SETUP] Building differentiable Evoformer + Distogram function...")
    af3_fn = build_evoformer_distogram_fn(
        model_dir=args.model_dir,
        num_recycles=args.num_recycles,
    )
    print("[SETUP] ✓  AF3 differentiable function ready")

    # ── Step 3: Build full AF3 model runner (only needed for --full_eval) ─────
    model_runner = None
    if args.full_eval:
        from pipeline import Input, make_model_config, ModelRunner
        print("[SETUP] Building full AF3 ModelRunner for ipTM/pLDDT evaluation...")
        gpu_devices = jax.devices("gpu")
        jax_device = (gpu_devices[args.gpu]
                      if gpu_devices and 0 <= args.gpu < len(gpu_devices)
                      else (gpu_devices[0] if gpu_devices else jax.devices()[0]))
        model_runner = ModelRunner(
            config=make_model_config(
                flash_attention_implementation=args.flash_attention_implementation,
                num_diffusion_samples=args.num_diffusion_samples,
                num_recycles=args.num_hard_recycles,
                return_embeddings=False,
            ),
            device=jax_device,
            model_dir=pathlib.Path(args.model_dir),
        )
        print(f"[SETUP] ✓  ModelRunner ready on {jax_device}")
    print()

    results: List[Dict] = []

    # shared kwargs for evaluate_config
    eval_kwargs = dict(
        full_eval=args.full_eval,
        model_runner=model_runner,
        args=args,
        output_dir=output_dir,
        verbose=args.verbose,
    )

    # ── Step 4: Run search ────────────────────────────────────────────────────
    t_search_start = time.time()

    if args.search_mode == "random":
        rng = random.Random(42)
        for trial_id in range(args.n_trials):
            config = _random_sample_config(rng)
            t0 = time.time()
            metrics = evaluate_config(
                config, batch_dict, binder_start_idx, binder_length,
                af3_fn, seeds, trial_id=trial_id, **eval_kwargs,
            )
            elapsed = time.time() - t0
            row = {**config, **metrics, "trial_id": trial_id, "elapsed_s": elapsed}
            results.append(row)
            _append_csv(csv_path, row)
            _print_trial(trial_id, config, metrics, elapsed)

    elif args.search_mode == "grid":
        configs = _grid_configs()
        print(f"[GRID] {len(configs)} configurations to evaluate")
        for trial_id, config in enumerate(configs):
            t0 = time.time()
            metrics = evaluate_config(
                config, batch_dict, binder_start_idx, binder_length,
                af3_fn, seeds, trial_id=trial_id, **eval_kwargs,
            )
            elapsed = time.time() - t0
            row = {**config, **metrics, "trial_id": trial_id, "elapsed_s": elapsed}
            results.append(row)
            _append_csv(csv_path, row)
            _print_trial(trial_id, config, metrics, elapsed)

    elif args.search_mode == "bayesian":
        _bayesian_search(
            n_trials=args.n_trials,
            batch_dict=batch_dict,
            binder_start_idx=binder_start_idx,
            binder_length=binder_length,
            af3_fn=af3_fn,
            seeds=seeds,
            results=results,
            csv_path=csv_path,
            full_eval=args.full_eval,
            model_runner=model_runner,
            args=args,
            output_dir=output_dir,
            verbose=args.verbose,
        )

    t_search_total = time.time() - t_search_start

    # ── Step 5: Report ────────────────────────────────────────────────────────
    if results:
        best = max(results, key=lambda r: r.get("score", float("-inf")))
        print("\n" + "=" * 70)
        print("  SEARCH COMPLETE")
        print("=" * 70)
        print(f"  Trials evaluated : {len(results)}")
        print(f"  Total time       : {t_search_total/60:.1f} min")
        print(f"  Results CSV      : {csv_path}")
        print()
        print("  ─── BEST CONFIG ───────────────────────────────────────────")
        for hp in HPARAM_SPACE:
            n = hp["name"]
            print(f"    {n:20s} = {best.get(n)}")
        print()
        print(f"  score        = {best['score']:+.4f}")
        if args.full_eval:
            print(f"  mean_ipTM    = {best['mean_iptm']:.3f} ± {best['std_iptm']:.3f}")
            print(f"  best_ipTM    = {best['best_iptm']:.3f}")
            print(f"  mean_pLDDT   = {best['mean_plddt']:.1f}")
            print(f"  best_pLDDT   = {best['best_plddt']:.1f}")
        else:
            print(f"  mean_loss    = {best['mean_loss']:.4f} ± {best['std_loss']:.4f}")
            print(f"  mean_contact = {best['mean_contact']:.3f}")
        print("=" * 70 + "\n")

        # Save best config as JSON
        best_config = {hp["name"]: best.get(hp["name"]) for hp in HPARAM_SPACE}
        best_json = output_dir / "best_config.json"
        with open(best_json, "w") as f:
            json.dump(best_config, f, indent=2)
        print(f"  Best config saved → {best_json}")

        # Print as ready-to-use CLI flags
        print("\n  ─── CLI FLAGS FOR BEST CONFIG ─────────────────────────────")
        flags = []
        for hp in HPARAM_SPACE:
            n = hp["name"]
            v = best.get(n)
            if v is None:
                flags.append(f"--{n} 0")
            else:
                flags.append(f"--{n} {v}")
        print("  python bindcraft_af3_gradient.py \\\n    " +
              " \\\n    ".join(flags))

    # ── Step 6: Optional top-K full validation (proxy mode only) ─────────────
    if args.validate_topk > 0 and results and not args.full_eval:
        _run_topk_validation(
            results=results,
            k=args.validate_topk,
            args=args,
            batch_dict=batch_dict,
            binder_start_idx=binder_start_idx,
            binder_length=binder_length,
            af3_fn=af3_fn,
            output_dir=output_dir,
        )


if __name__ == "__main__":
    main()