#!/usr/bin/env python3
"""
hf_train_1kb.py
------------------
HF Jobs version of modal_pipeline_1kb.py. Combines prepare_data() and
train_model() into one script, since HF Jobs runs a single command to
completion rather than exposing persistent, independently-callable
Functions the way Modal does -- there's no direct equivalent of
`train_model.remote()` to call repeatedly against a warm environment.

All the actual logic (blacklist handling, bedpe parsing with the chr1
fix, span-fit window placement, anchor-weighted masked loss) is copied
verbatim from modal_pipeline_1kb.py, which was already tested against
synthetic data -- this file only changes HOW it's invoked (plain
function calls + argparse, not Modal decorators) and WHERE data lives
(a mounted bucket path instead of a Modal Volume path).

Expects two mounted volumes (see submit_hf_job.py):
  /code  -- read-only dataset repo containing src/ and configs/
  /data  -- read-write bucket, plays the role all three Modal volumes
            (hic-data, hic-cache, hic-runs) played: raw/ downloads,
            cache/windows_1kb/ per-window .npz cache, runs/ checkpoints.

Run (inside the job, via submit_hf_job.py):
    python hf_train_1kb.py --model bilstm
    python hf_train_1kb.py --model bilstm --exclude-track RAD21
    python hf_train_1kb.py --prepare-only
"""
from __future__ import annotations

import argparse
import csv
import gzip
import hashlib
import json
import math
import random
import sys
import time
from pathlib import Path

import numpy as np
import pandas as pd
import yaml


def _cfg(config_dir: Path, config_name: str):
    with open(config_dir / config_name) as f:
        return yaml.safe_load(f)


def _track_fingerprint(track_names: list[str]) -> str:
    h = hashlib.sha1("|".join(sorted(track_names)).encode()).hexdigest()
    return h[:8]


# ── Blacklist -- verbatim from modal_pipeline_1kb.py ────────────────────────

def load_blacklist(bed_gz_path: Path) -> dict:
    regions: dict = {}
    with gzip.open(bed_gz_path, "rt") as fh:
        for line in fh:
            if line.startswith("#") or line.startswith("track"):
                continue
            p = line.rstrip("\n").split("\t")
            if len(p) < 3:
                continue
            chrom = p[0] if p[0].startswith("chr") else f"chr{p[0]}"
            try:
                start, end = int(p[1]), int(p[2])
            except ValueError:
                continue
            regions.setdefault(chrom, []).append((start, end))
    for chrom in regions:
        regions[chrom].sort()
    return regions


def overlaps_blacklist(chrom: str, start: int, end: int, blacklist: dict) -> bool:
    for (bl_start, bl_end) in blacklist.get(chrom, []):
        if start < bl_end and end > bl_start:
            return True
    return False


# ── Loop parsing -- verbatim, includes the chr1 header-detection fix ───────

def parse_encode_bedpe_gz(gz_path) -> list[dict]:
    loops = []
    with gzip.open(gz_path, "rt") as fh:
        for line in fh:
            if line.startswith("#"):
                continue
            p = line.strip().split("\t")
            if len(p) < 6:
                continue
            c1, c2 = p[0], p[3]
            try:
                x1, x2, y1, y2 = int(p[1]), int(p[2]), int(p[4]), int(p[5])
            except ValueError:
                continue
            if c1 != c2:
                continue
            chrom = c1 if c1.startswith("chr") else f"chr{c1}"
            loops.append({"chrom": chrom, "mid_x": (x1 + x2) // 2, "mid_y": (y1 + y2) // 2})
    return loops


# ── Window placement -- verbatim ────────────────────────────────────────────

def place_window_loop_centered(loop, window_size, flank_margin, jitter, rng, sizes, blacklist,
                                max_attempts: int = 5):
    chrom = loop["chrom"]
    mid_x, mid_y = loop["mid_x"], loop["mid_y"]
    span = abs(mid_y - mid_x)
    max_span = window_size - 2 * flank_margin
    if span > max_span:
        return None, "span_too_large"
    clen = sizes.get(chrom)
    if clen is None:
        return None, "unknown_chrom"
    true_center = (mid_x + mid_y) // 2
    for _ in range(max_attempts):
        center = true_center + rng.randint(-jitter, jitter)
        start = center - window_size // 2
        end = start + window_size
        if start < 0 or end > clen:
            continue
        if not (start + flank_margin <= mid_x <= end - flank_margin):
            continue
        if not (start + flank_margin <= mid_y <= end - flank_margin):
            continue
        if overlaps_blacklist(chrom, start, end, blacklist):
            continue
        return (start, end), None
    return None, "placement_failed"


def sample_background_window(chroms, window_size, sizes, blacklist, rng, max_attempts=20):
    chrom_list = list(chroms)
    for _ in range(max_attempts):
        chrom = rng.choice(chrom_list)
        clen = sizes.get(chrom)
        if clen is None or clen <= window_size:
            continue
        start = rng.randint(0, clen - window_size)
        end = start + window_size
        if overlaps_blacklist(chrom, start, end, blacklist):
            continue
        return chrom, start, end
    return None


# ── Losses -- verbatim ───────────────────────────────────────────────────────

def masked_huber_loss(pred, target, valid_mask, delta: float = 0.5):
    import torch
    diff = pred - target
    abs_diff = diff.abs()
    quadratic = torch.clamp(abs_diff, max=delta)
    linear = abs_diff - quadratic
    per_pixel = (0.5 * quadratic ** 2 + delta * linear) * valid_mask
    denom = valid_mask.sum().clamp(min=1.0)
    return per_pixel.sum() / denom


def masked_mse_loss(pred, target, valid_mask):
    per_pixel = (pred - target) ** 2 * valid_mask
    denom = valid_mask.sum().clamp(min=1.0)
    return per_pixel.sum() / denom


def build_anchor_weight_map(n_bins: int, anchor_r, anchor_c, boost_radius: int = 5,
                              boost_factor: float = 10.0):
    w = np.ones((n_bins, n_bins), dtype=np.float32)
    if anchor_r is not None and anchor_c is not None and anchor_r >= 0 and anchor_c >= 0:
        r0, r1 = max(0, anchor_r - boost_radius), min(n_bins, anchor_r + boost_radius + 1)
        c0, c1 = max(0, anchor_c - boost_radius), min(n_bins, anchor_c + boost_radius + 1)
        w[r0:r1, c0:c1] = boost_factor
    return w


def masked_weighted_huber_loss(pred, target, valid_mask, weight_map, delta: float = 0.5):
    import torch
    diff = pred - target
    abs_diff = diff.abs()
    quadratic = torch.clamp(abs_diff, max=delta)
    linear = abs_diff - quadratic
    per_pixel = (0.5 * quadratic ** 2 + delta * linear) * valid_mask * weight_map
    denom = (valid_mask * weight_map).sum().clamp(min=1.0)
    return per_pixel.sum() / denom


# ── Data prep ────────────────────────────────────────────────────────────────

def resolve_file(rel_path: str, shared_raw, local_raw: Path, url, genomic_io) -> Path:
    """Returns a usable local path for `rel_path`, preferring a file that
    already exists over downloading it.

    Lookup order:
      1. shared_raw/rel_path -- the read-only data bucket (e.g. the existing
         3d-genomics-prediction bucket, which already holds every bigWig and
         both .hic files)
      2. local_raw/rel_path  -- downloaded by an earlier run
      3. download from `url` into local_raw/rel_path

    This exists because the first HF Jobs run re-downloaded ~90GB from
    ENCODE into a fresh bucket while identical files sat unused in an
    existing one. Checking before fetching turns that into a no-op.
    Raises if missing everywhere AND no url is given, rather than returning
    a path that doesn't exist and failing later inside pyBigWig/hicstraw
    with a much less obvious message.
    """
    if shared_raw is not None:
        candidate = Path(shared_raw) / rel_path
        if candidate.exists() and candidate.stat().st_size > 0:
            print(f"  [shared] {rel_path}")
            return candidate
    candidate = local_raw / rel_path
    if candidate.exists() and candidate.stat().st_size > 0:
        print(f"  [cache] {rel_path}")
        return candidate
    if not url:
        raise FileNotFoundError(
            f"'{rel_path}' not found in shared raw dir ({shared_raw}) or local raw dir "
            f"({local_raw}), and no download url was configured for it.")
    genomic_io.download_file(url, candidate)
    return candidate


def prepare_data(cfg: dict, data_dir: Path, genomic_io, genome_constants,
                  shared_raw=None) -> None:
    exp = cfg["experiment"]
    raw = data_dir / "raw"
    raw.mkdir(parents=True, exist_ok=True)
    manifest_path = data_dir / "loop_manifest_1kb.csv"

    track_name_sets = {cl: sorted(cl_cfg["tracks"].keys()) for cl, cl_cfg in cfg["cell_lines"].items()}
    reference_tracks = next(iter(track_name_sets.values()))
    for cl, names in track_name_sets.items():
        if names != reference_tracks:
            raise ValueError(f"Track set mismatch: '{cl}' has {names}, expected {reference_tracks}.")
    track_fp = _track_fingerprint(reference_tracks)

    fingerprint_path = data_dir / "manifest_fingerprint_1kb.txt"
    needed_fp = (f"{sorted(cfg['cell_lines'].keys())}|{track_fp}|"
                 f"{exp['bin_size']}|{exp['window_size']}|{exp['flank_margin']}")

    blacklist = {}
    if cfg.get("blacklist_filename"):
        bl_path = resolve_file(cfg["blacklist_filename"], shared_raw, raw,
                                cfg.get("blacklist_url"), genomic_io)
        blacklist = load_blacklist(bl_path)
        n_regions = sum(len(v) for v in blacklist.values())
        print(f"[prepare] Blacklist: {n_regions} regions across {len(blacklist)} chroms.")

    for cl_name, cl_cfg in cfg["cell_lines"].items():
        print(f"\n[prepare] {cl_name}: resolving {len(cl_cfg['tracks'])} tracks, loop calls, "
              f"and the intact .hic file...")
        for track_name, entry in cl_cfg["tracks"].items():
            resolve_file(entry["path"], shared_raw, raw, entry.get("url"), genomic_io)
        resolve_file(cl_cfg["loop_filename"], shared_raw, raw, cl_cfg.get("loop_url"), genomic_io)
        # hic_path is config-driven now -- the bucket names .hic files by
        # accession (HepG2/ENCFF807IRK.hic), so the old hardcoded
        # f"{cell_line}.hic" convention doesn't match it. Falls back to that
        # convention when a config doesn't specify hic_path.
        resolve_file(cl_cfg.get("hic_path", f"{cl_name}.hic"), shared_raw, raw,
                      cl_cfg.get("hic_url"), genomic_io)
        print(f"[prepare] {cl_name}: done.")

    if manifest_path.exists():
        old_fp = fingerprint_path.read_text().strip() if fingerprint_path.exists() else None
        if old_fp == needed_fp:
            print(f"\n[prepare] Manifest exists for {needed_fp}. Downloads verified above; "
                  f"skipping window rebuild.")
            return
        print(f"\n[prepare] Config changed ({old_fp} -> {needed_fp}). Rebuilding manifest.")
        manifest_path.unlink()

    sizes = genome_constants.chrom_sizes_hg38()
    bin_size = exp["bin_size"]; window_size = exp["window_size"]
    jitter = exp["jitter"]; flank_margin = exp["flank_margin"]
    n_bins = window_size // bin_size
    train_chroms = set(exp["train_chroms"]); test_chroms = set(exp["test_chroms"])
    rng = random.Random(exp["seed"])

    rows = []
    n_span_rejected = 0; n_placement_failed = 0
    for cl_name, cl_cfg in cfg["cell_lines"].items():
        loop_file = resolve_file(cl_cfg["loop_filename"], shared_raw, raw,
                                  cl_cfg.get("loop_url"), genomic_io)
        loops = parse_encode_bedpe_gz(loop_file)
        print(f"\n[prepare] {cl_name}: {len(loops)} intra-chromosomal loops parsed.")

        for split, chroms, n_target in [
            ("train", train_chroms, exp["train_loops_per_cell_line"]),
            ("test", test_chroms, exp["test_loops_per_cell_line"]),
        ]:
            eligible = [l for l in loops if l["chrom"] in chroms]
            rng.shuffle(eligible)
            placed = 0
            for loop in eligible:
                if placed >= n_target:
                    break
                result, reason = place_window_loop_centered(
                    loop, window_size, flank_margin, jitter, rng, sizes, blacklist)
                if result is None:
                    if reason == "span_too_large":
                        n_span_rejected += 1
                    else:
                        n_placement_failed += 1
                    continue
                start, end = result
                rows.append({
                    "split": split, "cell_line": cl_name, "window_type": "loop",
                    "chrom": loop["chrom"], "start": start, "end": end,
                    "bin_size": bin_size, "n_bins": n_bins,
                    "loop_mid_x": loop["mid_x"], "loop_mid_y": loop["mid_y"],
                    "loop_span_bp": abs(loop["mid_y"] - loop["mid_x"]),
                    "track_fingerprint": track_fp,
                })
                placed += 1
            status = "WARNING: only" if placed < n_target else " "
            print(f"  {status} {placed}/{n_target} loop windows: {cl_name}/{split}")

        n_bg_target = exp.get("background_windows_per_cell_line", 0)
        placed_bg = 0
        for _ in range(n_bg_target * 3):
            if placed_bg >= n_bg_target:
                break
            result = sample_background_window(train_chroms, window_size, sizes, blacklist, rng)
            if result is None:
                continue
            chrom, start, end = result
            rows.append({
                "split": "train", "cell_line": cl_name, "window_type": "background",
                "chrom": chrom, "start": start, "end": end,
                "bin_size": bin_size, "n_bins": n_bins,
                "loop_mid_x": "", "loop_mid_y": "", "loop_span_bp": "",
                "track_fingerprint": track_fp,
            })
            placed_bg += 1
        print(f"  {placed_bg}/{n_bg_target} background windows: {cl_name}/train")

    print(f"\n[prepare] Rejected: {n_span_rejected} loops (span too large), "
          f"{n_placement_failed} (blacklist/bounds after retries).")

    fieldnames = ["split", "cell_line", "window_type", "chrom", "start", "end",
                  "bin_size", "n_bins", "loop_mid_x", "loop_mid_y", "loop_span_bp",
                  "track_fingerprint"]
    with open(manifest_path, "w", newline="") as fh:
        w = csv.DictWriter(fh, fieldnames=fieldnames)
        w.writeheader(); w.writerows(rows)
    fingerprint_path.write_text(needed_fp)
    print(f"\n[prepare] Done: {len(rows)} windows -> {manifest_path}")


# ── Training ─────────────────────────────────────────────────────────────────

def train_model(model_name: str, cfg: dict, data_dir: Path, src_dir: Path,
                 exclude_track: str = "", shared_raw=None) -> dict:
    import torch
    from torch.utils.data import DataLoader, Dataset
    sys.path.insert(0, str(src_dir))
    from genomic_io import read_bigwig, _get_hicfile
    from oe_extraction import get_oe_matrix, MAX_OE

    tr_cfg = cfg["training"]
    raw = data_dir / "raw"

    def _resolve_existing(rel_path: str) -> str:
        """Same shared-then-local lookup as resolve_file(), minus the
        download fallback -- by training time prepare_data() has already
        ensured every file exists somewhere, so a miss here is a bug worth
        surfacing immediately rather than silently re-fetching 40GB mid-run."""
        if shared_raw is not None:
            cand = Path(shared_raw) / rel_path
            if cand.exists():
                return str(cand)
        return str(raw / rel_path)

    train_seed = int(tr_cfg.get("seed", cfg.get("experiment", {}).get("seed", 42)))
    torch.manual_seed(train_seed)
    torch.cuda.manual_seed_all(train_seed)
    print(f"[train] seed={train_seed}")

    track_names = sorted(next(iter(cfg["cell_lines"].values()))["tracks"].keys())
    if exclude_track:
        if exclude_track not in track_names:
            raise ValueError(f"exclude_track='{exclude_track}' not in {track_names}")
        track_names = [t for t in track_names if t != exclude_track]
        print(f"[train] LEAVE-ONE-TRACK-OUT: excluding '{exclude_track}'")
    n_tracks = len(track_names)
    track_fp = _track_fingerprint(track_names)
    max_oe = float(tr_cfg.get("max_oe", MAX_OE))
    print(f"[train] {n_tracks} tracks: {track_names}")

    class LoopDataset(Dataset):
        def __init__(self, manifest_path, split, normalizer=None, cache_dir=None, clip_pct=99.0):
            full = pd.read_csv(manifest_path)
            self.manifest = full[full["split"] == split].reset_index(drop=True)
            self.normalizer = normalizer; self.clip_pct = clip_pct
            self.cache_root = Path(cache_dir) if cache_dir else None
            if self.cache_root:
                self.cache_root.mkdir(parents=True, exist_ok=True)

        def __len__(self):
            return len(self.manifest)

        def _paths(self, cl):
            cl_tracks = cfg["cell_lines"][cl]["tracks"]
            return {name: _resolve_existing(cl_tracks[name]["path"]) for name in track_names} | \
                   {"hic": _resolve_existing(cfg["cell_lines"][cl].get("hic_path", f"{cl}.hic"))}

        def _cache(self, row):
            if self.cache_root is None:
                return None
            bs = int(row["bin_size"])
            return self.cache_root / (f"{row['cell_line']}_{row['chrom']}_{row['start']}_"
                                        f"{row['end']}_log2oe_v2_{bs}_{track_fp}.npz")

        def get_tracks(self, idx):
            row = self.manifest.iloc[idx]
            p = self._paths(row["cell_line"]); n = int(row["n_bins"])
            ch, s, e = row["chrom"], int(row["start"]), int(row["end"])
            arrays = [read_bigwig(p[name], ch, s, e, n, self.clip_pct) for name in track_names]
            return np.stack(arrays, axis=0)

        def __getitem__(self, idx):
            row = self.manifest.iloc[idx]; cl = row["cell_line"]
            ch, s, e = row["chrom"], int(row["start"]), int(row["end"])
            bs = int(row["bin_size"]); cf = self._cache(row)
            if cf is not None and cf.exists():
                data = np.load(cf)
                tracks, log2_oe, valid_mask = data["tracks"], data["log2_oe"], data["valid_mask"]
            else:
                tracks = self.get_tracks(idx)
                hf = _get_hicfile(self._paths(cl)["hic"])
                _, log2_oe, valid_mask = get_oe_matrix(hf, ch, s, e, bs, max_oe=max_oe)
                if cf is not None:
                    np.savez_compressed(cf, tracks=tracks.astype(np.float32),
                                         log2_oe=log2_oe, valid_mask=valid_mask)
            if self.normalizer is not None:
                s_n = self.normalizer[cl]
                tracks = (tracks - s_n["mean"][:, None]) / s_n["std"][:, None]
            if row["window_type"] == "loop" and pd.notna(row.get("loop_mid_x")):
                anchor_r = (int(row["loop_mid_x"]) - s) // bs
                anchor_c = (int(row["loop_mid_y"]) - s) // bs
            else:
                anchor_r = anchor_c = -1
            return {"tracks": torch.from_numpy(tracks).float(),
                    "target": torch.from_numpy(log2_oe).float(),
                    "valid_mask": torch.from_numpy(valid_mask.astype(np.float32)),
                    "anchor_r": anchor_r, "anchor_c": anchor_c, "cell_line": cl}

    manifest = data_dir / "loop_manifest_1kb.csv"
    run_dir = data_dir / "runs" / f"1kb_{model_name}_{track_fp}"
    run_dir.mkdir(parents=True, exist_ok=True)
    cache_dir = data_dir / "cache" / "windows_1kb"

    norm_path = data_dir / "runs" / f"normalizer_1kb_{track_fp}.json"
    config_cls = set(cfg["cell_lines"].keys())
    need_refit = True
    if norm_path.exists():
        with open(norm_path) as f:
            normalizer = json.load(f)
        if set(normalizer.keys()) == config_cls:
            need_refit = False
            print(f"[train] Loaded normalizer for {config_cls}.")

    if need_refit:
        train_df = pd.read_csv(manifest)
        train_df = train_df[train_df["split"] == "train"].reset_index(drop=True)
        raw_ds = LoopDataset(manifest, "train", normalizer=None, cache_dir=cache_dir)
        normalizer = {}
        for cl in train_df["cell_line"].unique():
            cl_pos = train_df[train_df["cell_line"] == cl].index.tolist()
            n_s = min(80, len(cl_pos))
            sample_idx = np.linspace(0, len(cl_pos) - 1, n_s).astype(int)
            sums = np.zeros(n_tracks); sq_sums = np.zeros(n_tracks); count = 0
            for si in sample_idx:
                t = raw_ds.get_tracks(cl_pos[si])
                sums += t.sum(axis=1); sq_sums += (t ** 2).sum(axis=1); count += t.shape[1]
            mean = sums / count
            std = np.sqrt(np.clip(sq_sums / count - mean ** 2, 1e-8, None))
            normalizer[cl] = {"mean": mean.tolist(), "std": std.tolist()}
            print(f" {cl}: " + " ".join(f"{n}={m:.3f}" for n, m in zip(track_names, mean)))
        run_dir.parent.mkdir(parents=True, exist_ok=True)
        with open(norm_path, "w") as f:
            json.dump(normalizer, f, indent=2)

    norm_np = {cl: {"mean": np.array(v["mean"], dtype=np.float32),
                     "std": np.array(v["std"], dtype=np.float32)} for cl, v in normalizer.items()}

    train_ds = LoopDataset(manifest, "train", normalizer=norm_np, cache_dir=cache_dir)
    test_ds = LoopDataset(manifest, "test", normalizer=norm_np, cache_dir=cache_dir)
    print(f"[train] {len(train_ds)} train, {len(test_ds)} test windows")

    bs = tr_cfg["batch_size"]
    train_loader = DataLoader(train_ds, batch_size=bs, shuffle=True, num_workers=0, pin_memory=True)
    test_loader = DataLoader(test_ds, batch_size=bs, shuffle=False, num_workers=0, pin_memory=True)
    n_bins = int(train_ds.manifest["n_bins"].iloc[0])

    device = "cuda"
    if model_name == "cnn":
        from model_cnn import CNNHiCModel
        model = CNNHiCModel(n_tracks=n_tracks, n_bins=n_bins).to(device)
        wd = float(tr_cfg["weight_decay"]); clip_norm = None
    elif model_name == "unet":
        from model_unet import UNetHiCModel
        model = UNetHiCModel(n_tracks=n_tracks, n_bins=n_bins).to(device)
        wd = float(tr_cfg["weight_decay"]); clip_norm = None
    elif model_name == "transformer":
        from model_transformer import TransformerHiCModel
        dropout = float(tr_cfg.get("transformer_dropout", 0.3))
        model = TransformerHiCModel(n_tracks=n_tracks, n_bins=n_bins, dropout=dropout).to(device)
        wd = float(tr_cfg.get("transformer_weight_decay", 1e-3))
        clip_norm = float(tr_cfg.get("transformer_grad_clip_norm", 0.5))
    elif model_name == "bilstm":
        from model_bilstm import BiLSTMHiCModel
        dropout = float(tr_cfg.get("bilstm_dropout", 0.2))
        model = BiLSTMHiCModel(n_tracks=n_tracks, n_bins=n_bins, dropout=dropout).to(device)
        wd = float(tr_cfg.get("bilstm_weight_decay", 1e-5))
        clip_norm = float(tr_cfg.get("bilstm_grad_clip_norm", 1.0))
    else:
        raise ValueError(f"Unknown model: {model_name}")

    n_params = sum(p.numel() for p in model.parameters())
    print(f"[train] {model_name}: {n_params:,} params")

    epochs = tr_cfg["epochs"][model_name]
    patience = tr_cfg.get("early_stopping_patience", 10)
    warmup = tr_cfg.get("transformer_warmup_steps", 0) if model_name == "transformer" else 0
    total_steps = len(train_loader) * epochs
    optimizer = torch.optim.Adam(model.parameters(), lr=float(tr_cfg["lr"]), weight_decay=wd)

    def lr_lambda(step):
        if warmup <= 0:
            return 1.0
        if step < warmup:
            return (step + 1) / warmup
        progress = min(1.0, (step - warmup) / max(1, total_steps - warmup))
        return max(0.05, 0.5 * (1 + math.cos(math.pi * progress)))

    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)

    log_path = run_dir / "training_log.csv"
    ckpt_path = run_dir / "best_model.pt"
    best_test_loss = float("inf"); epochs_no_improve = 0; stopped_at = None
    huber_delta = float(tr_cfg.get("huber_delta", 0.5))
    anchor_boost_radius = int(tr_cfg.get(f"{model_name}_anchor_boost_radius",
                                          tr_cfg.get("anchor_boost_radius", 5)))
    anchor_boost_factor = float(tr_cfg.get(f"{model_name}_anchor_boost_factor",
                                            tr_cfg.get("anchor_boost_factor", 10.0)))
    print(f"[train] anchor-weighted loss: radius={anchor_boost_radius} factor={anchor_boost_factor}")

    fields = ["epoch", "train_huber", "train_huber_weighted", "train_mse", "test_huber",
              "test_mse", "epoch_seconds", "lr"]
    with open(log_path, "w", newline="") as log_fh:
        writer = csv.DictWriter(log_fh, fieldnames=fields)
        writer.writeheader()
        for epoch in range(1, epochs + 1):
            t0 = time.time()
            model.train()
            tr_hub = tr_hub_weighted = tr_mse = n_tr = 0; last_lr = None
            for batch in train_loader:
                tracks = batch["tracks"].to(device, non_blocking=True)
                target = batch["target"].to(device, non_blocking=True)
                valid = batch["valid_mask"].to(device, non_blocking=True)
                anchors_r = batch["anchor_r"].numpy(); anchors_c = batch["anchor_c"].numpy()
                weight_maps = np.stack([
                    build_anchor_weight_map(n_bins, r if r >= 0 else None, c if c >= 0 else None,
                                             boost_radius=anchor_boost_radius,
                                             boost_factor=anchor_boost_factor)
                    for r, c in zip(anchors_r, anchors_c)])
                weight = torch.from_numpy(weight_maps).float().to(device, non_blocking=True)

                optimizer.zero_grad()
                pred = model(tracks)
                loss = masked_weighted_huber_loss(pred, target, valid, weight, delta=huber_delta)
                loss.backward()
                if clip_norm:
                    torch.nn.utils.clip_grad_norm_(model.parameters(), clip_norm)
                optimizer.step(); scheduler.step()
                last_lr = scheduler.get_last_lr()[0]
                tr_hub_weighted += loss.item()
                tr_hub += masked_huber_loss(pred, target, valid, delta=huber_delta).item()
                tr_mse += masked_mse_loss(pred, target, valid).item()
                n_tr += 1
            tr_hub /= max(1, n_tr); tr_hub_weighted /= max(1, n_tr); tr_mse /= max(1, n_tr)

            model.eval(); te_hub = te_mse = n_te = 0
            with torch.no_grad():
                for batch in test_loader:
                    tracks = batch["tracks"].to(device, non_blocking=True)
                    target = batch["target"].to(device, non_blocking=True)
                    valid = batch["valid_mask"].to(device, non_blocking=True)
                    pred = model(tracks)
                    te_hub += masked_huber_loss(pred, target, valid, delta=huber_delta).item()
                    te_mse += masked_mse_loss(pred, target, valid).item()
                    n_te += 1
            te_hub /= max(1, n_te); te_mse /= max(1, n_te)

            dt = time.time() - t0
            print(f"  epoch {epoch:3d}/{epochs}  tr_hub={tr_hub:.4f} tr_hub_w={tr_hub_weighted:.4f} "
                  f"tr_mse={tr_mse:.4f}  te_hub={te_hub:.4f} te_mse={te_mse:.4f}  "
                  f"lr={last_lr:.2e}  ({dt:.1f}s)")
            writer.writerow({"epoch": epoch, "train_huber": tr_hub,
                              "train_huber_weighted": tr_hub_weighted, "train_mse": tr_mse,
                              "test_huber": te_hub, "test_mse": te_mse,
                              "epoch_seconds": dt, "lr": last_lr})
            log_fh.flush()

            if te_hub < best_test_loss:
                best_test_loss = te_hub; epochs_no_improve = 0
                torch.save({"model_name": model_name, "n_bins": n_bins, "n_tracks": n_tracks,
                            "track_names": track_names, "epoch": epoch,
                            "model_state_dict": model.state_dict(),
                            "test_huber": te_hub, "test_mse": te_mse}, ckpt_path)
            else:
                epochs_no_improve += 1
                if epochs_no_improve >= patience:
                    stopped_at = epoch
                    print(f"[train] Early stop epoch {epoch}")
                    break

    return {"model": model_name, "n_tracks": n_tracks, "best_test_huber": best_test_loss,
            "best_epoch": int(torch.load(ckpt_path, map_location="cpu", weights_only=False)["epoch"]),
            "stopped_early_at": stopped_at}


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--model", default="bilstm", choices=["cnn", "unet", "transformer", "bilstm"])
    p.add_argument("--config-name", default="encode_sources_1kb.yaml")
    p.add_argument("--exclude-track", default="")
    p.add_argument("--prepare-only", action="store_true")
    p.add_argument("--code-dir", default="/code")
    p.add_argument("--data-dir", default="/data")
    p.add_argument("--shared-raw-dir", default="",
                    help="Read-only mount holding the already-uploaded raw data "
                         "(e.g. /shared for the 3d-genomics-prediction bucket). Files found "
                         "here are used as-is; anything missing falls back to downloading "
                         "into --data-dir/raw. Empty = download everything, the old behavior.")
    args = p.parse_args()

    code_dir = Path(args.code_dir)
    data_dir = Path(args.data_dir)
    shared_raw = Path(args.shared_raw_dir) if args.shared_raw_dir else None
    src_dir = code_dir / "src"
    sys.path.insert(0, str(src_dir))
    import genomic_io
    import genome_constants

    cfg = _cfg(code_dir / "configs", args.config_name)

    print("=" * 65)
    print(f"HepG2/HCT116 1kb loop prediction (HF Jobs) -- config: {args.config_name}")
    if shared_raw:
        print(f"shared read-only raw data: {shared_raw}")
    print("=" * 65)
    print("\n[Step 1] Preparing data...")
    prepare_data(cfg, data_dir, genomic_io, genome_constants, shared_raw=shared_raw)
    print("[Step 1] Done.")
    if args.prepare_only:
        return

    print(f"\n[Step 2] Training {args.model}"
          + (f" (excluding {args.exclude_track})" if args.exclude_track else "") + "...")
    r = train_model(args.model, cfg, data_dir, src_dir, args.exclude_track, shared_raw=shared_raw)
    s = f" (early@{r['stopped_early_at']})" if r["stopped_early_at"] else ""
    print(f"\n<- {args.model.upper()}{s}: best_test_huber={r['best_test_huber']:.4f}  "
          f"epoch={r['best_epoch']}")


if __name__ == "__main__":
    main()