cpp_bfb_bridge.py

"""
cpp_bfb_bridge.py - wrapper around the Ambigram C++ BFB ILP solver
(Lu/Yu et al., Nat Commun 2023, s41467-023-41259-w).

Provides candidate BFB derivative chromosome paths for v4's posterior ranker.

Workflow:
1. Auto-detect Ambigram binary, htslib, and cbc from environment variables
   or command-line overrides.
2. Given a per-clone .txt file (Ambigram-canonical format), invoke
   `Ambigram --op bfb --in_lh CLONE.txt --lp_prefix PREFIX [--all]`
3. Parse the BFB path string from stdout (last non-Cbc line containing
   `\\d+[+-]\\|` patterns)
4. Convert path syntax `3-2-|1+2+3+2+|3-2-` to v4 path token list
   `['3-', '2-', '1+', '2+', '3+', '2+', '3-', '2-']`
   (the `|` separator marks foldback inversions but is already implicit in
   the alternating signs; `||` marks translocations between chromosomes)
5. Optionally synthesize a per-clone .txt file from a multi-clone .lh + clone CN column.
"""
from __future__ import annotations

import os
import re
import subprocess
import tempfile
from collections import defaultdict
from typing import Dict, List, Optional, Tuple


# --- portable defaults, overridable by environment or caller ------------------
DEFAULT_BINARY = os.environ.get("SVGRAM_AMBIGRAM_BINARY", "Ambigram")
DEFAULT_HTSLIB_DIR = os.environ.get("SVGRAM_HTSLIB_DIR", "")
DEFAULT_CBC_DIR = os.environ.get("SVGRAM_CBC_DIR", "")
DEFAULT_TIMEOUT = 60


# --- output parsing -----------------------------------------------------------

# A BFB path token: digits + sign, e.g. "3+", "12-".
TOKEN_RE = re.compile(r"\d+[+\-]")
# A line is a candidate BFB-path output line if it contains at least 2 tokens
# and may contain the FBI ('|') or translocation ('||') separators.
PATH_LINE_RE = re.compile(r"^[\d+\-|]+$")


def parse_bfb_path_line(line: str) -> List[str]:
    """Convert `3-2-|1+2+3+2+|3-2-` -> ['3-','2-','1+','2+','3+','2+','3-','2-'].

    `|` and `||` are dropped — they are descriptive separators (FBI / translocation)
    but path semantics are fully captured by the token sequence + signs.
    """
    line = line.strip()
    if not line:
        return []
    return TOKEN_RE.findall(line)


def extract_bfb_paths(stdout_text: str) -> List[List[str]]:
    """Scan binary stdout for BFB-path lines.

    The Ambigram binary prints BFB paths via `printBFB()` to stdout, mixed in
    with Cbc solver chatter. A BFB-path line consists exclusively of segment
    tokens (digit+sign) and `|` separators.
    """
    out: List[List[str]] = []
    for raw in stdout_text.splitlines():
        s = raw.strip()
        if not s:
            continue
        # quick reject: must contain at least one digit+sign token, and only
        # consist of tokens + `|`
        if not TOKEN_RE.search(s):
            continue
        # remove all tokens + pipes; remainder must be empty for a path line
        residue = TOKEN_RE.sub("", s).replace("|", "")
        if residue:
            continue
        tokens = parse_bfb_path_line(s)
        if len(tokens) >= 2:
            out.append(tokens)
    # de-dup while preserving order
    seen = set()
    uniq: List[List[str]] = []
    for p in out:
        key = tuple(p)
        if key in seen:
            continue
        seen.add(key)
        uniq.append(p)
    return uniq


# --- env + binary invocation --------------------------------------------------

def build_env(htslib_dir: str = DEFAULT_HTSLIB_DIR,
              cbc_dir: str = DEFAULT_CBC_DIR) -> Dict[str, str]:
    """Build the os.environ extension required for the binary to run."""
    env = os.environ.copy()
    ld_paths = [p for p in (htslib_dir, env.get("LD_LIBRARY_PATH", ""))
                if p]
    env["LD_LIBRARY_PATH"] = ":".join(ld_paths)
    if cbc_dir and cbc_dir not in env.get("PATH", "").split(":"):
        env["PATH"] = cbc_dir + ":" + env.get("PATH", "")
    return env


def run_bfb_binary(binary: str,
                   in_txt: str,
                   *,
                   lp_prefix: Optional[str] = None,
                   all_paths: bool = True,
                   reversed_strand: bool = False,
                   timeout: int = DEFAULT_TIMEOUT,
                   env: Optional[Dict[str, str]] = None) -> Tuple[int, str, str]:
    """Invoke the C++ BFB binary on a per-clone .txt file.

    Returns (returncode, stdout, stderr). Caller parses stdout via
    `extract_bfb_paths`.
    """
    if env is None:
        env = build_env()
    if lp_prefix is None:
        # binary writes <prefix>.lp / .sol / .mps.gz alongside; use a tempdir
        td = tempfile.mkdtemp(prefix="cppbfb_")
        lp_prefix = os.path.join(td, "run")
    # The Ambigram binary internally reads `./<lp_prefix>.sol` after cbc writes
    # `<lp_prefix>.sol`, so an absolute lp_prefix becomes `./tmp/...` which
    # resolves relative to cwd and breaks. Workaround: split into workdir
    # (chdir target) + basename (relative prefix passed to the binary).
    workdir = os.path.dirname(lp_prefix) or os.getcwd()
    lp_basename = os.path.basename(lp_prefix)
    cmd = [binary, "--op", "bfb",
           "--in_lh", in_txt,
           "--lp_prefix", lp_basename]
    if all_paths:
        cmd.append("--all")
    if reversed_strand:
        cmd.append("--reversed")

    try:
        proc = subprocess.run(
            cmd, cwd=workdir, env=env,
            capture_output=True, text=True,
            timeout=timeout,
        )
        return proc.returncode, proc.stdout, proc.stderr
    except subprocess.TimeoutExpired as e:
        return 124, e.stdout or "", e.stderr or "timeout"


# --- per-clone .txt synthesis from subclone .lh -------------------------------

def synthesize_per_clone_txt(
    subclone_lh: str,
    clone_name: str,
    clone_index: int,
    out_path: str,
) -> str:
    """Generate an Ambigram-canonical per-clone .txt from a subclone .lh.

    Used when the per-clone .txt does not already exist next to the .lh.

    Subclone .lh format observed:
        SAMPLE clone0,clone1,clone2
        SOURCE 1
        SINK 4
        SEG H:N:chr:start:end CN_clone0 CN_clone1 CN_clone2
        JUNC H:u:dir H:v:dir aggregate_count [other]

    Output Ambigram .txt format observed:
        SOURCE 1 12
        SINK 4 8
        SEG H:N:chr:start:end CN read_count
        JUNC H:u:dir H:v:dir count read_count U B

    For JUNC counts we lack per-clone splits in the .lh, so we apportion the
    aggregate count proportionally to clone CN coverage on the involved
    endpoints (best-effort fallback). Real per-clone .txt files are preferred.
    """
    sample = None
    sources: List[Tuple[int, int]] = []  # (segid, count) — count fabricated as 1
    sinks: List[Tuple[int, int]] = []
    seg_lines: List[Tuple[str, List[int]]] = []   # (raw_seg_token, all CN values)
    junc_lines: List[Tuple[str, str, int]] = []   # (u_token, v_token, agg_count)

    with open(subclone_lh) as fh:
        for raw in fh:
            parts = raw.strip().split()
            if not parts:
                continue
            tag = parts[0]
            if tag == "SAMPLE":
                sample = parts[1]
            elif tag == "SOURCE":
                segid = int(parts[1])
                sources.append((segid, 1))
            elif tag == "SINK":
                segid = int(parts[1])
                sinks.append((segid, 1))
            elif tag == "SEG":
                token = parts[1]
                cns = [int(round(float(x))) for x in parts[2:]]
                seg_lines.append((token, cns))
            elif tag in ("JUNC", "REF"):
                u, v = parts[1], parts[2]
                # take the first numeric field after the two endpoint tokens
                try:
                    agg = int(parts[3])
                except (IndexError, ValueError):
                    agg = 0
                junc_lines.append((u, v, agg))

    if sample is None:
        raise ValueError(f"{subclone_lh}: missing SAMPLE line")
    clone_names = sample.split(",")
    if clone_index >= len(clone_names):
        raise ValueError(
            f"{subclone_lh}: clone_index {clone_index} out of range "
            f"({len(clone_names)} clones)")
    if not seg_lines:
        raise ValueError(f"{subclone_lh}: no SEG lines")

    # Per-clone segment CN
    seg_cn: Dict[int, int] = {}
    for token, cns in seg_lines:
        segid = int(token.split(":")[1])
        # If this clone has no CN value for this seg, fall back to last column
        cn = cns[clone_index] if clone_index < len(cns) else cns[-1]
        seg_cn[segid] = cn

    with open(out_path, "w") as out:
        # SOURCE / SINK lines (fabricate count = clone CN at that segment)
        for segid, _ in sources:
            out.write(f"SOURCE {segid} {seg_cn.get(segid, 1)}\n")
        for segid, _ in sinks:
            out.write(f"SINK {segid} {seg_cn.get(segid, 1)}\n")
        # SEG lines: keep token; CN = clone-specific; read_count fabricated as
        # ratio of clone CN to the AVG_DP-equivalent (use cn directly)
        for token, cns in seg_lines:
            segid = int(token.split(":")[1])
            cn = seg_cn[segid]
            # fabricate read_count proportional to cn (Ambigram only checks > 0)
            read_count = max(1, cn)
            out.write(f"SEG {token} {cn} {read_count}\n")
        # JUNC lines: apportion aggregate count by clone CN at the endpoints
        # ratio = clone_cn_avg(u,v) / sum_over_clones(cn_avg(u,v))
        clone_count = len(clone_names)
        for u_tok, v_tok, agg in junc_lines:
            try:
                u_id = int(u_tok.split(":")[1])
                v_id = int(v_tok.split(":")[1])
            except (IndexError, ValueError):
                continue
            # sum-of-clone-CN-averages on this edge's endpoints
            u_cns = next(
                (cns for tok, cns in seg_lines
                 if int(tok.split(":")[1]) == u_id), [agg]
            )
            v_cns = next(
                (cns for tok, cns in seg_lines
                 if int(tok.split(":")[1]) == v_id), [agg]
            )
            this_clone = (
                (u_cns[clone_index] if clone_index < len(u_cns) else u_cns[-1]) +
                (v_cns[clone_index] if clone_index < len(v_cns) else v_cns[-1])
            )
            total = sum(
                ((u_cns[i] if i < len(u_cns) else u_cns[-1]) +
                 (v_cns[i] if i < len(v_cns) else v_cns[-1]))
                for i in range(clone_count)
            )
            count = int(round(agg * (this_clone / total))) if total > 0 else 0
            count = max(0, count)
            if count == 0:
                continue
            read_count = count  # placeholder
            out.write(
                f"JUNC {u_tok} {v_tok} {count} {read_count} U B\n"
            )

    return out_path


# --- per-clone .txt discovery -------------------------------------------------

def discover_per_clone_txt(subclone_lh: str,
                           clone_name: str) -> Optional[str]:
    """Look for an existing per-clone .txt file next to the subclone .lh.

    Convention observed in user's data:
        <prefix>.lh                    -> subclone graph
        <prefix>_<clone_name>.txt      -> per-clone Ambigram input

    Returns the path if found, else None.
    """
    base = subclone_lh
    if base.endswith(".lh"):
        base = base[:-3]
    candidates = [
        f"{base}_{clone_name}.txt",
        f"{base}.{clone_name}.txt",
        f"{base}_with_ref_{clone_name}.txt",  # observed in cell_line data
    ]
    for c in candidates:
        if os.path.isfile(c):
            return c
    return None


# --- top-level convenience ---------------------------------------------------

def get_bfb_candidates_for_clone(
    *,
    subclone_lh: str,
    clone_name: str,
    clone_index: int,
    binary: str = DEFAULT_BINARY,
    htslib_dir: str = DEFAULT_HTSLIB_DIR,
    cbc_dir: str = DEFAULT_CBC_DIR,
    timeout: int = DEFAULT_TIMEOUT,
    all_paths: bool = True,
    workdir: Optional[str] = None,
    verbose: bool = False,
) -> Tuple[List[List[str]], str]:
    """Return BFB path candidates for one clone.

    Returns (paths, info_str).
    paths: list of token lists (each is one BFB-path candidate).
    info_str: short status (which .txt was used, return code, # paths).

    On any failure (binary not found, segfault, parse failure) returns ([], info).
    """
    if not os.path.isfile(binary):
        return [], {'src': 'none', 'rc': -1, 'n_paths': 0,
                    'stderr_head': f"binary not found: {binary}"}

    # 1. find or synthesize the per-clone .txt
    in_txt = discover_per_clone_txt(subclone_lh, clone_name)
    synthesized = False
    if in_txt is None:
        if workdir is None:
            workdir = tempfile.mkdtemp(prefix="cppbfb_synth_")
        os.makedirs(workdir, exist_ok=True)
        in_txt = os.path.join(workdir, f"clone_{clone_name}.txt")
        try:
            synthesize_per_clone_txt(
                subclone_lh, clone_name, clone_index, in_txt)
            synthesized = True
        except Exception as e:
            return [], {'src': 'synth_failed', 'rc': -1, 'n_paths': 0,
                        'stderr_head': f"synth failed: {e}"}

    # 2. invoke binary
    env = build_env(htslib_dir=htslib_dir, cbc_dir=cbc_dir)
    if workdir is None:
        workdir = tempfile.mkdtemp(prefix="cppbfb_run_")
    os.makedirs(workdir, exist_ok=True)
    lp_prefix = os.path.join(workdir, f"clone_{clone_name}_run")
    rc, stdout, stderr = run_bfb_binary(
        binary, in_txt, lp_prefix=lp_prefix,
        all_paths=all_paths, timeout=timeout, env=env,
    )
    if verbose:
        print(f"[cpp_bfb] clone={clone_name} txt={in_txt} "
              f"synth={synthesized} rc={rc}")

    # 3. parse stdout (defensive: TimeoutExpired may yield bytes; decode safely)
    if isinstance(stdout, (bytes, bytearray)):
        stdout = stdout.decode('utf-8', errors='replace')
    if isinstance(stderr, (bytes, bytearray)):
        stderr = stderr.decode('utf-8', errors='replace')
    paths = extract_bfb_paths(stdout)
    info = {
        'src': 'synth' if synthesized else 'preexisting',
        'rc': rc,
        'n_paths': len(paths),
        'in_txt': in_txt,
        'stderr_head': stderr[:200] if (rc != 0 and not paths) else '',
    }
    return paths, info


def get_bfb_candidates_for_txt(
    *,
    in_txt: str,
    clone_name: str,
    binary: str = DEFAULT_BINARY,
    htslib_dir: str = DEFAULT_HTSLIB_DIR,
    cbc_dir: str = DEFAULT_CBC_DIR,
    timeout: int = DEFAULT_TIMEOUT,
    all_paths: bool = True,
    workdir: Optional[str] = None,
    verbose: bool = False,
) -> Tuple[List[List[str]], str]:
    """Return freshly generated BFB candidates from one clone-resolved graph.

    This direct mode is for v5.3 ``--clone-graph`` inputs.  It invokes the C++
    BFB ILP solver on ``in_txt`` and parses stdout.  It never discovers or
    reads existing ``.path`` outputs.
    """
    if not os.path.isfile(binary):
        return [], {'src': 'none', 'rc': -1, 'n_paths': 0,
                    'stderr_head': f"binary not found: {binary}"}
    if not os.path.isfile(in_txt):
        return [], {'src': 'missing_txt', 'rc': -1, 'n_paths': 0,
                    'stderr_head': f"input graph not found: {in_txt}"}

    env = build_env(htslib_dir=htslib_dir, cbc_dir=cbc_dir)
    if workdir is None:
        workdir = tempfile.mkdtemp(prefix="cppbfb_direct_")
    os.makedirs(workdir, exist_ok=True)
    lp_prefix = os.path.join(workdir, f"clone_{clone_name}_run")
    rc, stdout, stderr = run_bfb_binary(
        binary, in_txt, lp_prefix=lp_prefix,
        all_paths=all_paths, timeout=timeout, env=env,
    )
    if verbose:
        print(f"[cpp_bfb] clone={clone_name} txt={in_txt} direct rc={rc}")

    if isinstance(stdout, (bytes, bytearray)):
        stdout = stdout.decode('utf-8', errors='replace')
    if isinstance(stderr, (bytes, bytearray)):
        stderr = stderr.decode('utf-8', errors='replace')
    paths = extract_bfb_paths(stdout)
    info = {
        'src': 'direct_txt',
        'rc': rc,
        'n_paths': len(paths),
        'in_txt': in_txt,
        'stderr_head': stderr[:200] if (rc != 0 and not paths) else '',
    }
    return paths, info