paths.py

"""Repo-root path helpers. All scripts should import this instead of hardcoding hosts."""
from __future__ import annotations

import os
import shutil
import sys
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parent


def add_repo_to_sys_path() -> Path:
    root = str(REPO_ROOT)
    if root not in sys.path:
        sys.path.insert(0, root)
    return REPO_ROOT


def resolve_path(path: str | Path) -> Path:
    p = Path(path).expanduser()
    if p.is_absolute():
        return p
    return (REPO_ROOT / p).resolve()


def peptiverse_root() -> Path:
    env = os.environ.get("PEPTIVERSE_ROOT")
    if env:
        root = Path(env).expanduser().resolve()
    else:
        root = (REPO_ROOT.parent / "PeptiVerse").resolve()
    if not root.is_dir():
        raise FileNotFoundError(
            "PeptiVerse was not found. Clone it first and either place it next "
            f"to this repo ({REPO_ROOT.parent / 'PeptiVerse'}) or set PEPTIVERSE_ROOT.\n"
            "  git clone https://huggingface.co/ChatterjeeLab/PeptiVerse"
        )
    return root


def mafft_path() -> str:
    env = os.environ.get("MAFFT_PATH")
    if env:
        if Path(env).is_file() and os.access(env, os.X_OK):
            return env
        raise FileNotFoundError(f"MAFFT_PATH is set but not an executable: {env}")
    found = shutil.which("mafft")
    if found:
        return found
    raise FileNotFoundError(
        "MAFFT was not found on PATH. Install it and make sure `mafft` is "
        "executable, or set MAFFT_PATH to the binary. GFP FPredX oracles require MAFFT."
    )


def smiles_vocab_files() -> tuple[Path, Path]:
    return (
        REPO_ROOT / "smiles_tokenizer" / "new_vocab.txt",
        REPO_ROOT / "smiles_tokenizer" / "new_splits.txt",
    )


def selfies_vocab_path() -> Path:
    return REPO_ROOT / "smiles_tokenizer" / "selfies_vocab.json"


# --- Cas9 ---------------------------------------------------------------

def protein2pam_root() -> Path:
    """Root of the protein2pam package (provides the custom ESM loader).

    Weights come from the Hub ("Profluent-Bio/protein2pam-cas9_full"); this is
    only the Python package that defines the model class.
    """
    env = os.environ.get("PROTEIN2PAM_ROOT")
    root = Path(env).expanduser().resolve() if env else (REPO_ROOT / "cas9" / "protein2pam")
    if not root.is_dir():
        raise FileNotFoundError(
            "protein2pam was not found. Install it next to the Cas9 code at "
            f"{REPO_ROOT / 'cas9' / 'protein2pam'} or set PROTEIN2PAM_ROOT.\n"
            "  pip install protein2pam==0.2.0"
        )
    return root


def cas_predictor_root() -> Path:
    """Root of the cas_predictor package (model.py + lightning_module.py)."""
    env = os.environ.get("CAS_PREDICTOR_ROOT")
    root = Path(env).expanduser().resolve() if env else (REPO_ROOT / "cas9" / "cas_predictor")
    if not root.is_dir():
        raise FileNotFoundError(
            "cas_predictor was not found. Place it at "
            f"{REPO_ROOT / 'cas9' / 'cas_predictor'} or set CAS_PREDICTOR_ROOT."
        )
    return root


def cas9_classifier_ckpt() -> Path:
    """Cas9-likeness classifier checkpoint (hard constraint + objective)."""
    env = os.environ.get("CAS9_CLASSIFIER_CKPT")
    p = resolve_path(env) if env else (REPO_ROOT / "cas9" / "classifier_ckpt" / "last.ckpt")
    if not p.is_file():
        raise FileNotFoundError(
            f"Cas9 classifier checkpoint not found at {p}. Download it into "
            "cas9/classifier_ckpt/ or set CAS9_CLASSIFIER_CKPT."
        )
    return p


def cas9_classifier_config() -> Path | None:
    """Optional classifier config; None falls back to the ckpt hyper_parameters."""
    env = os.environ.get("CAS9_CLASSIFIER_CONFIG")
    if env:
        return resolve_path(env)
    default = REPO_ROOT / "cas9" / "classifier_ckpt" / "config.yaml"
    return default if default.is_file() else None


def cas9_hmm_db(name: str) -> Path:
    """HMM database under cas9/hmm/ (e.g. 'cas9_domain_detector_RuvCmix.hmm')."""
    env = os.environ.get("CAS9_HMM_DIR")
    root = Path(env).expanduser().resolve() if env else (REPO_ROOT / "cas9" / "hmm")
    p = root / name
    if not p.is_file():
        raise FileNotFoundError(
            f"HMM database {name} not found at {p}. Place it in cas9/hmm/ or set CAS9_HMM_DIR."
        )
    return p