"""Utility functions for docking pose filtering."""

from __future__ import annotations

import json
import os
import shlex
import subprocess
import tempfile
import warnings
from contextlib import contextmanager
from pathlib import Path
from typing import Any, Literal

import numpy as np
import pandas as pd
from rdkit import Chem
from rdkit.Chem import AllChem

from hedgehog.configs.logger import logger
from hedgehog.utils.nvmolkit import maybe_enable_nvmolkit
from hedgehog.utils.parallel import parallel_map, resolve_n_jobs

# ---------------------------------------------------------------------------
# Shared helpers
# ---------------------------------------------------------------------------


def _project_root() -> Path:
    # src/hedgehog/docking_filters/utils.py -> project root
    return Path(__file__).resolve().parent.parent.parent.parent


@contextmanager
def _suppress_posecheck_fast_warnings():
    """Suppress known third-party warning from posecheck-fast/torch interop."""
    with warnings.catch_warnings():
        warnings.filterwarnings(
            "ignore",
            message="The given NumPy array is not writable",
            category=UserWarning,
        )
        yield


@contextmanager
def _suppress_interaction_filter_warnings():
    """Suppress known noisy MDAnalysis warnings during interaction filtering."""
    with warnings.catch_warnings():
        warnings.filterwarnings(
            "ignore",
            message="Deprecated in version 2.8.0",
            category=DeprecationWarning,
            module=r"MDAnalysis\.topology\.tables",
        )
        warnings.filterwarnings(
            "ignore",
            message="No `bonds` attribute in this AtomGroup. Guessing bonds based on atoms coordinates",
            category=UserWarning,
            module=r"MDAnalysis\.converters\.RDKit",
        )
        yield


def _resolve_existing_path(base_folder: Path, path: str | Path) -> Path:
    """Resolve a path (possibly relative) to an existing absolute path."""
    p = Path(path)
    if p.is_absolute():
        return p

    candidates = [
        (base_folder / p).resolve(),
        (_project_root() / p).resolve(),
        (Path.cwd() / p).resolve(),
    ]
    for c in candidates:
        if c.exists():
            return c
    return (base_folder / p).resolve()


def _make_passthrough_df(n_mols: int, columns: dict[str, Any]) -> pd.DataFrame:
    """Build a 'pass everyone' DataFrame with mol_idx and the given constant columns.

    Args:
        n_mols: Number of molecules (rows).
        columns: Mapping of column name -> constant value for every row.

    Returns:
        DataFrame with ``mol_idx`` in ``range(n_mols)`` plus the requested columns.
    """
    data: dict[str, Any] = {"mol_idx": range(n_mols)}
    for col, val in columns.items():
        data[col] = [val] * n_mols if isinstance(val, float) and np.isnan(val) else val
    return pd.DataFrame(data)


# Sentinel for conformer-RMSD error results.
_CONFORMER_ERROR: dict[str, Any] = {
    "min_rmsd": float("inf"),
    "n_conformers_generated": 0,
    "passed": False,
}


def _error_result(**defaults: Any) -> dict[str, Any]:
    """Build a default error-result dict.

    Returns a dict with ``error: None`` merged with *defaults*.
    """
    out = dict(defaults)
    out.setdefault("error", None)
    return out


def _fail_result(error: str, **defaults: Any) -> dict[str, Any]:
    """Build a failure-result dict with a specific error message."""
    out = dict(defaults)
    out["error"] = error
    return out


def _aggregate_filter_results(
    raw_results: list[dict[str, Any]],
    col_mapping: dict[str, str],
    pass_col: str,
    filter_name: str,
) -> pd.DataFrame:
    """Aggregate per-molecule filter results into a DataFrame.

    Args:
        raw_results: List of dicts returned by per-molecule worker functions.
        col_mapping: Mapping of ``{result_key: dataframe_column}``.
        pass_col: Name of the boolean pass column in the DataFrame.
        filter_name: Human-readable filter name for log messages.

    Returns:
        DataFrame with ``mol_idx``, mapped columns, and ``pass_col``.
    """
    results: list[dict[str, Any]] = []
    for i, res in enumerate(raw_results):
        if res.get("error"):
            logger.warning("%s failed for mol %d: %s", filter_name, i, res["error"])
        row: dict[str, Any] = {"mol_idx": i}
        for src_key, dst_col in col_mapping.items():
            row[dst_col] = res[src_key]
        row[pass_col] = res["passed"]
        results.append(row)

    df = pd.DataFrame(results)
    logger.info(
        "%s: %d/%d passed",
        filter_name,
        int(df[pass_col].sum()),
        len(df),
    )
    return df


def _resolve_worker_command(
    env_var: str,
    venv_name: str,
    entry_name: str,
    module_path: str,
    install_hint: str,
) -> list[str]:
    """Resolve a worker command from environment, venv entry-point, or venv python.

    Resolution order:
    1. Environment variable (shell-split).
    2. Platform-specific entry-point script inside the venv.
    3. ``python -m <module_path>`` using the venv python.

    Args:
        env_var: Environment variable name (e.g. ``HEDGEHOG_SHEPHERD_WORKER_CMD``).
        venv_name: Venv directory name under project root (e.g. ``.venv-shepherd-worker``).
        entry_name: Entry-point script base name (e.g. ``hedgehog-shepherd-worker``).
        module_path: Dotted module path for ``python -m`` fallback.
        install_hint: Human-readable install instruction for the error message.

    Returns:
        Command tokens ready for ``subprocess.run()``.

    Raises:
        RuntimeError: If no command could be resolved.
    """
    raw_cmd = os.environ.get(env_var, "").strip()
    if raw_cmd:
        parts = shlex.split(raw_cmd)
        if not parts:
            raise RuntimeError(f"{env_var} is empty")
        return parts

    project_root = _project_root()
    if os.name == "nt":
        worker_entry = project_root / venv_name / "Scripts" / f"{entry_name}.exe"
        venv_python = project_root / venv_name / "Scripts" / "python.exe"
    else:
        worker_entry = project_root / venv_name / "bin" / entry_name
        venv_python = project_root / venv_name / "bin" / "python"

    if worker_entry.exists():
        return [str(worker_entry)]
    if venv_python.exists():
        return [str(venv_python), "-m", module_path]

    raise RuntimeError(
        f"No {entry_name} command found. Set {env_var} or run: {install_hint}"
    )


# ---------------------------------------------------------------------------
# Conformer RMSD helpers (unified for symmetry-aware and naive modes)
# ---------------------------------------------------------------------------


def _select_embedding_params(
    method_name: str, random_seed: int
) -> AllChem.EmbedParameters:
    """Return RDKit embedding parameters for the given method name."""
    if method_name == "ETKDGv3":
        params = AllChem.ETKDGv3()
    elif method_name == "ETKDGv2":
        params = AllChem.ETKDGv2()
    else:
        params = AllChem.ETKDG()
    params.randomSeed = random_seed
    return params


def _prepare_reference_mol_for_conformer(
    mol: Chem.Mol, include_hydrogens: bool
) -> Chem.Mol:
    """Prepare the reference molecule for conformer RMSD computation."""
    if include_hydrogens:
        return Chem.AddHs(mol, addCoords=True)
    return Chem.RemoveHs(Chem.Mol(mol), sanitize=False)


def _check_conformer_rmsd_single(args: tuple) -> dict[str, Any]:
    """Unified conformer RMSD checker -- supports both symmetry-aware and naive modes.

    When ``use_symmetry`` is True, uses ``GetBestRMS`` which considers molecular
    symmetry (via substructure match) when computing RMSD.  When False, uses the
    naive ``AlignMol`` (Kabsch, identity permutation).

    Args:
        args: Tuple of
            (mol, num_conformers, max_rmsd, method_name, random_seed,
             include_hydrogens, max_matches, early_stop_on_pass, use_symmetry,
             use_nvmolkit).

    Returns:
        Dict with keys: min_rmsd, n_conformers_generated, passed, error.
    """
    (
        mol,
        num_conformers,
        max_rmsd,
        method_name,
        random_seed,
        include_hydrogens,
        max_matches,
        early_stop_on_pass,
        use_symmetry,
        use_nvmolkit,
    ) = args
    try:
        if use_nvmolkit:
            maybe_enable_nvmolkit(
                project_root=_project_root(),
                context="docking_filters.conformer_deviation",
            )

        reference_mol = _prepare_reference_mol_for_conformer(mol, include_hydrogens)
        params = _select_embedding_params(method_name, random_seed)

        mol_confs = Chem.Mol(reference_mol)
        mol_confs.RemoveAllConformers()
        AllChem.EmbedMultipleConfs(mol_confs, numConfs=num_conformers, params=params)
        n_generated = mol_confs.GetNumConformers()

        if n_generated == 0:
            return _fail_result("no conformers generated", **_CONFORMER_ERROR)

        if use_symmetry:
            from rdkit.Chem import rdMolAlign

            rmsd_values: list[float] = []
            for conf_id in range(n_generated):
                rmsd = rdMolAlign.GetBestRMS(
                    mol_confs,
                    reference_mol,
                    prbId=conf_id,
                    refId=0,
                    maxMatches=max_matches,
                )
                rmsd_values.append(rmsd)
                if early_stop_on_pass and rmsd <= max_rmsd:
                    break
        else:
            rmsd_values = []
            for conf_id in range(n_generated):
                conf_mol = Chem.Mol(mol_confs)
                conf_mol.RemoveAllConformers()
                conf_mol.AddConformer(mol_confs.GetConformer(conf_id), assignId=True)
                rmsd = AllChem.AlignMol(conf_mol, reference_mol)
                rmsd_values.append(rmsd)

        min_rmsd = min(rmsd_values)
        return _error_result(
            min_rmsd=min_rmsd,
            n_conformers_generated=n_generated,
            passed=min_rmsd <= max_rmsd,
        )
    except Exception as e:
        return _fail_result(str(e), **_CONFORMER_ERROR)


# ---------------------------------------------------------------------------
# Search-box filter
# ---------------------------------------------------------------------------


def _autobox_from_reference_ligand(
    base_folder: Path, docking_cfg: dict[str, Any]
) -> tuple[tuple[float, float, float], tuple[float, float, float]] | None:
    """Build a Vina/GNINA-style search box from autobox_ligand + autobox_add.

    Returns:
        (center_xyz, size_xyz) or None if insufficient config.
    """
    tool_cfg = docking_cfg.get("gnina_config") or docking_cfg.get("smina_config") or {}
    autobox_ligand = tool_cfg.get("autobox_ligand") or docking_cfg.get("autobox_ligand")
    if not autobox_ligand:
        return None

    ref_path = _resolve_existing_path(base_folder, autobox_ligand)
    if not ref_path.exists():
        return None

    suppl = Chem.SDMolSupplier(str(ref_path))
    ref_mol = next((m for m in suppl if m is not None), None)
    if ref_mol is None or ref_mol.GetNumConformers() == 0:
        return None

    conf = ref_mol.GetConformer()
    xs: list[float] = []
    ys: list[float] = []
    zs: list[float] = []
    for i in range(ref_mol.GetNumAtoms()):
        p = conf.GetAtomPosition(i)
        xs.append(float(p.x))
        ys.append(float(p.y))
        zs.append(float(p.z))

    min_x, max_x = min(xs), max(xs)
    min_y, max_y = min(ys), max(ys)
    min_z, max_z = min(zs), max(zs)

    add = float(tool_cfg.get("autobox_add", 0.0) or 0.0)
    min_x -= add
    max_x += add
    min_y -= add
    max_y += add
    min_z -= add
    max_z += add

    center = ((min_x + max_x) / 2, (min_y + max_y) / 2, (min_z + max_z) / 2)
    size = (max_x - min_x, max_y - min_y, max_z - min_z)
    return center, size


def apply_search_box_filter(
    mols: list[Chem.Mol],
    base_folder: Path,
    docking_cfg: dict[str, Any],
    config: dict[str, Any],
) -> pd.DataFrame:
    """Fast filter: ensure docked poses are inside the configured search box.

    Supports explicit center/size, or autobox_ligand-based boxes.

    Returns:
        DataFrame with columns: mol_idx, frac_atoms_outside_box, pass_search_box
    """
    enabled = bool(config.get("enabled", True))
    if not enabled:
        return _make_passthrough_df(
            len(mols), {"frac_atoms_outside_box": 0.0, "pass_search_box": True}
        )

    # Prefer explicit center/size if present, otherwise fall back to autobox reference.
    center = docking_cfg.get("center")
    size = docking_cfg.get("size")
    if not (
        isinstance(center, (list, tuple))
        and len(center) >= 3
        and isinstance(size, (list, tuple))
        and len(size) >= 3
    ):
        box = _autobox_from_reference_ligand(base_folder, docking_cfg)
        if box is None:
            logger.warning(
                "Search-box filter enabled but no box could be resolved; passing all poses."
            )
            return _make_passthrough_df(
                len(mols), {"frac_atoms_outside_box": 0.0, "pass_search_box": True}
            )
        center, size = box

    cx, cy, cz = float(center[0]), float(center[1]), float(center[2])
    sx, sy, sz = float(size[0]), float(size[1]), float(size[2])
    hx, hy, hz = sx / 2.0, sy / 2.0, sz / 2.0

    max_outside = float(config.get("max_outside_fraction", 0.0) or 0.0)

    frac_outside: list[float] = []
    passed: list[bool] = []

    for mol in mols:
        if mol is None or mol.GetNumConformers() == 0:
            frac_outside.append(1.0)
            passed.append(False)
            continue
        conf = mol.GetConformer()
        n_atoms = mol.GetNumAtoms()
        outside = 0
        for i in range(n_atoms):
            p = conf.GetAtomPosition(i)
            x, y, z = float(p.x), float(p.y), float(p.z)
            if (
                x < cx - hx
                or x > cx + hx
                or y < cy - hy
                or y > cy + hy
                or z < cz - hz
                or z > cz + hz
            ):
                outside += 1
        frac = outside / n_atoms if n_atoms else 1.0
        frac_outside.append(frac)
        passed.append(frac <= max_outside)

    df = pd.DataFrame(
        {
            "mol_idx": range(len(mols)),
            "frac_atoms_outside_box": frac_outside,
            "pass_search_box": passed,
        }
    )
    logger.info(
        "Search-box filter: %d/%d passed (max_outside_fraction=%.3f)",
        int(df["pass_search_box"].sum()),
        len(df),
        max_outside,
    )
    return df


# ---------------------------------------------------------------------------
# Molecule loading
# ---------------------------------------------------------------------------


def load_molecules_from_sdf(sdf_path: str | Path) -> list[Chem.Mol]:
    """Load molecules from SDF file."""
    sdf_path = Path(sdf_path)
    if not sdf_path.exists():
        raise FileNotFoundError(f"SDF file not found: {sdf_path}")

    suppl = Chem.SDMolSupplier(str(sdf_path))
    mols = [mol for mol in suppl if mol is not None]
    logger.info("Loaded %d molecules from %s", len(mols), sdf_path)
    return mols


# ---------------------------------------------------------------------------
# Pose quality filter (PoseCheck legacy)
# ---------------------------------------------------------------------------


def apply_pose_quality_filter(
    mols: list[Chem.Mol],
    protein_pdb: str | Path,
    config: dict[str, Any],
) -> pd.DataFrame:
    """
    Apply pose quality filter using PoseCheck.

    Checks for steric clashes and ligand strain energy.

    Args:
        mols: List of RDKit molecules with 3D coordinates
        protein_pdb: Path to protein PDB file
        config: Filter configuration dict

    Returns:
        DataFrame with columns: mol_idx, clashes, strain_energy, pass
    """
    from posecheck import PoseCheck

    max_clashes = config.get("max_clashes", 2)
    max_strain = config.get("max_strain_energy", 10.0)

    logger.info(
        "Running pose quality filter (max_clashes=%d, max_strain=%.1f)",
        max_clashes,
        max_strain,
    )

    pc = PoseCheck()
    pc.load_protein_from_pdb(str(protein_pdb))
    pc.ligands = mols

    # Calculate metrics
    clashes = pc.calculate_clashes()
    strain_energies = pc.calculate_strain_energy()

    # Build results DataFrame
    df = pd.DataFrame(
        {
            "mol_idx": range(len(mols)),
            "clashes": clashes,
            "strain_energy": strain_energies,
        }
    )
    df["pass_pose_quality"] = (df["clashes"] <= max_clashes) & (
        df["strain_energy"] <= max_strain
    )
    logger.info(
        "Pose quality filter: %d/%d passed",
        int(df["pass_pose_quality"].sum()),
        len(df),
    )
    return df


# ---------------------------------------------------------------------------
# Interaction filter (ProLIF)
# ---------------------------------------------------------------------------


def apply_interaction_filter(
    mols: list[Chem.Mol],
    protein_pdb: str | Path,
    config: dict[str, Any],
) -> pd.DataFrame:
    """
    Apply interaction filter using ProLIF.

    Checks for required protein-ligand interactions.

    Args:
        mols: List of RDKit molecules with 3D coordinates
        protein_pdb: Path to protein PDB file
        config: Filter configuration dict

    Returns:
        DataFrame with columns: mol_idx, n_hbonds, interactions, pass
    """
    with _suppress_interaction_filter_warnings():
        import MDAnalysis as mda
        import prolif as plf

    min_hbonds = config.get("min_hbonds", 0)
    required_residues = config.get("required_residues", [])
    forbidden_residues = config.get("forbidden_residues", [])
    interaction_types = config.get(
        "interaction_types", ["HBDonor", "HBAcceptor", "Hydrophobic", "VdWContact"]
    )

    logger.info("Running interaction filter (min_hbonds=%d)", min_hbonds)

    with _suppress_interaction_filter_warnings():
        # Load protein
        u = mda.Universe(str(protein_pdb))
        protein = plf.Molecule.from_mda(u)

        # Convert RDKit mols to ProLIF molecules
        plf_mols = []
        for mol in mols:
            mol_h = Chem.AddHs(mol, addCoords=True)
            plf_mols.append(plf.Molecule.from_rdkit(mol_h))

        # Calculate interactions (exclude FaceToFace due to edge case errors)
        safe_interactions = [it for it in interaction_types if it != "FaceToFace"]
        fp = plf.Fingerprint(interactions=safe_interactions)
        n_jobs = resolve_n_jobs(config)
        fp.run_from_iterable(plf_mols, protein, n_jobs=n_jobs)

    df_interactions = fp.to_dataframe(drop_empty=False)

    # Build results DataFrame
    results = []
    for i in range(len(mols)):
        if i < len(df_interactions):
            row = df_interactions.iloc[i]
            # Count H-bonds
            hbond_cols = [
                c for c in row.index if "HBDonor" in str(c) or "HBAcceptor" in str(c)
            ]
            n_hbonds = sum(row[hbond_cols].values) if hbond_cols else 0

            # Check required residues
            has_required = True
            for res in required_residues:
                res_cols = [c for c in row.index if res in str(c)]
                if res_cols and not any(row[res_cols].values):
                    has_required = False
                    break

            # Check forbidden residues
            has_forbidden = False
            for res in forbidden_residues:
                res_cols = [c for c in row.index if res in str(c)]
                if res_cols and any(row[res_cols].values):
                    has_forbidden = True
                    break

            passed = n_hbonds >= min_hbonds and has_required and not has_forbidden
            interactions_str = ",".join([str(c) for c in row.index if row[c]])
        else:
            n_hbonds = 0
            passed = min_hbonds == 0 and not required_residues
            interactions_str = ""

        results.append(
            {
                "mol_idx": i,
                "n_hbonds": n_hbonds,
                "interactions": interactions_str,
                "pass_interactions": passed,
            }
        )

    df = pd.DataFrame(results)
    logger.info(
        "Interaction filter: %d/%d passed",
        int(df["pass_interactions"].sum()),
        len(df),
    )
    return df


# ---------------------------------------------------------------------------
# Shepherd-Score filter
# ---------------------------------------------------------------------------


def _check_shepherd_score_single(mol: Chem.Mol) -> dict[str, Any]:
    """Compute shape Tanimoto score for a single molecule.

    Args:
        mol: RDKit molecule.

    Returns:
        Dict with keys: score, passed, error.
    """
    import torch
    from shepherd_score.score.gaussian_overlap import shape_tanimoto

    ref_positions_list = _SHEPHERD_REF_POSITIONS_LIST
    alpha = _SHEPHERD_ALPHA
    min_score = _SHEPHERD_MIN_SCORE
    try:
        if ref_positions_list is None:
            raise RuntimeError(
                "Shepherd-Score worker configuration was not initialized"
            )
        ref_positions = torch.tensor(ref_positions_list, dtype=torch.float32)

        mol_h = Chem.AddHs(mol, addCoords=True)
        mol_conf = mol_h.GetConformer()
        mol_positions = torch.tensor(mol_conf.GetPositions(), dtype=torch.float32)

        score = shape_tanimoto(ref_positions, mol_positions, alpha=alpha).item()
        passed = score >= min_score
        return {"score": score, "passed": passed, "error": None}
    except Exception as e:
        return {"score": 0.0, "passed": False, "error": str(e)}


_SHEPHERD_REF_POSITIONS_LIST: list[list[float]] | None = None
_SHEPHERD_ALPHA: float = 0.81
_SHEPHERD_MIN_SCORE: float = 0.5


def _init_shepherd_score_worker(
    ref_positions_list: list[list[float]],
    alpha: float,
    min_score: float,
) -> None:
    global _SHEPHERD_REF_POSITIONS_LIST, _SHEPHERD_ALPHA, _SHEPHERD_MIN_SCORE
    _SHEPHERD_REF_POSITIONS_LIST = ref_positions_list
    _SHEPHERD_ALPHA = float(alpha)
    _SHEPHERD_MIN_SCORE = float(min_score)


def _soft_skip_shepherd_filter(mols: list[Chem.Mol], reason: str) -> pd.DataFrame:
    logger.warning(
        "Shepherd-Score backend unavailable (%s). Skipping filter. "
        "To enable it, run: uv run hedgehog setup shepherd-worker --yes",
        reason,
    )
    return _make_passthrough_df(
        len(mols),
        {"shape_score": float("nan"), "pass_shepherd_score": True},
    )


def _resolve_shepherd_worker_command() -> list[str]:
    return _resolve_worker_command(
        env_var="HEDGEHOG_SHEPHERD_WORKER_CMD",
        venv_name=".venv-shepherd-worker",
        entry_name="hedgehog-shepherd-worker",
        module_path="hedgehog.workers.shepherd_worker",
        install_hint="uv run hedgehog setup shepherd-worker --yes",
    )


def _write_sdf(path: Path, mols: list[Chem.Mol]) -> None:
    writer = Chem.SDWriter(str(path))
    writer.SetKekulize(False)
    for i, mol in enumerate(mols):
        current = Chem.Mol(mol)
        current.SetProp("_Name", f"mol_{i}")
        writer.write(current)
    writer.close()


def _apply_shepherd_score_filter_worker(
    mols: list[Chem.Mol],
    reference_mol: Chem.Mol,
    config: dict[str, Any],
    progress_cb=None,
) -> pd.DataFrame:
    def _auto_install_worker() -> None:
        from hedgehog.setup import ensure_shepherd_worker

        project_root = _project_root()
        old_auto = os.environ.get("HEDGEHOG_AUTO_INSTALL")
        try:
            os.environ["HEDGEHOG_AUTO_INSTALL"] = "1"
            ensure_shepherd_worker(project_root, python_bin=worker_python)
        finally:
            if old_auto is None:
                os.environ.pop("HEDGEHOG_AUTO_INSTALL", None)
            else:
                os.environ["HEDGEHOG_AUTO_INSTALL"] = old_auto

    auto_install_worker = bool(config.get("auto_install_worker", True))
    worker_python = config.get("worker_python")
    try:
        worker_cmd = _resolve_shepherd_worker_command()
    except Exception as exc:
        if not auto_install_worker:
            raise
        logger.info(
            "Shepherd worker command was not found, attempting auto-install "
            "(python=%s): %s",
            worker_python or "auto",
            exc,
        )
        _auto_install_worker()
        worker_cmd = _resolve_shepherd_worker_command()

    min_shape_score = float(config.get("min_shape_score", 0.5))
    alpha = float(config.get("alpha", 0.81))
    timeout_sec = int(config.get("worker_timeout_sec", 900))

    def _run_worker(worker_cmd: list[str]) -> pd.DataFrame:
        with tempfile.TemporaryDirectory(prefix="hedgehog_shepherd_worker_") as tmp_dir:
            tmp_path = Path(tmp_dir)
            input_sdf = tmp_path / "input.sdf"
            reference_sdf = tmp_path / "reference.sdf"
            params_json = tmp_path / "params.json"
            output_json = tmp_path / "output.json"
            meta_json = tmp_path / "meta.json"

            _write_sdf(input_sdf, mols)
            _write_sdf(reference_sdf, [reference_mol])
            params_json.write_text(
                json.dumps(
                    {
                        "alpha": alpha,
                        "min_shape_score": min_shape_score,
                    }
                ),
                encoding="utf-8",
            )

            cmd = [
                *worker_cmd,
                "--input-sdf",
                str(input_sdf),
                "--reference-sdf",
                str(reference_sdf),
                "--params-json",
                str(params_json),
                "--output-json",
                str(output_json),
                "--meta-json",
                str(meta_json),
            ]

            proc = subprocess.run(
                cmd,
                check=False,
                capture_output=True,
                text=True,
                timeout=timeout_sec,
            )
            if proc.returncode != 0:
                stderr = proc.stderr.strip()
                stdout = proc.stdout.strip()
                detail = stderr or stdout or "worker exited with non-zero status"
                raise RuntimeError(
                    f"Shepherd worker failed (code={proc.returncode}): {detail}"
                )

            if not output_json.exists():
                raise RuntimeError("Shepherd worker did not produce output.json")

            payload = json.loads(output_json.read_text(encoding="utf-8"))
            if not isinstance(payload, list):
                raise RuntimeError("Shepherd worker output.json has invalid format")

            results: list[dict[str, Any]] = []
            for i in range(len(mols)):
                item = (
                    payload[i]
                    if i < len(payload) and isinstance(payload[i], dict)
                    else {}
                )
                score = float(item.get("shape_score", 0.0))
                passed = bool(item.get("pass_shepherd_score", score >= min_shape_score))
                err = item.get("error")
                if err:
                    logger.warning("Shepherd-Score failed for mol %d: %s", i, err)
                results.append(
                    {
                        "mol_idx": i,
                        "shape_score": score,
                        "pass_shepherd_score": passed,
                    }
                )

        if progress_cb is not None:
            progress_cb(len(mols), len(mols))

        df = pd.DataFrame(results)
        logger.info(
            "Shepherd-Score filter: %d/%d passed",
            int(df["pass_shepherd_score"].sum()),
            len(df),
        )
        return df

    try:
        return _run_worker(worker_cmd)
    except Exception as exc:
        if not auto_install_worker:
            raise
        logger.info(
            "Shepherd worker execution failed, attempting auto-repair (python=%s): %s",
            worker_python or "auto",
            exc,
        )
        _auto_install_worker()
        worker_cmd = _resolve_shepherd_worker_command()
        return _run_worker(worker_cmd)


def _apply_shepherd_score_filter_inprocess(
    mols: list[Chem.Mol],
    reference_mol: Chem.Mol,
    config: dict[str, Any],
    progress_cb=None,
) -> pd.DataFrame:
    """
    Apply Shepherd-Score filter for 3D shape similarity.

    Args:
        mols: List of RDKit molecules with 3D coordinates
        reference_mol: Reference molecule for comparison
        config: Filter configuration dict

    Returns:
        DataFrame with columns: mol_idx, shape_score, pass
    """
    min_shape_score = config.get("min_shape_score", 0.5)
    alpha = config.get("alpha", 0.81)

    logger.info("Running Shepherd-Score filter (min_shape_score=%.2f)", min_shape_score)

    # Get reference coordinates as plain Python list for picklability
    ref_h = Chem.AddHs(reference_mol, addCoords=True)
    ref_conf = ref_h.GetConformer()
    ref_positions_list = ref_conf.GetPositions().tolist()

    n_jobs = resolve_n_jobs(config)
    raw_results = parallel_map(
        _check_shepherd_score_single,
        mols,
        n_jobs,
        progress=progress_cb,
        initializer=_init_shepherd_score_worker,
        initargs=(ref_positions_list, alpha, min_shape_score),
    )

    return _aggregate_filter_results(
        raw_results,
        col_mapping={"score": "shape_score"},
        pass_col="pass_shepherd_score",
        filter_name="Shepherd-Score filter",
    )


def apply_shepherd_score_filter(
    mols: list[Chem.Mol],
    reference_mol: Chem.Mol,
    config: dict[str, Any],
    progress_cb=None,
) -> pd.DataFrame:
    """Apply Shepherd-Score filter with backend fallback model."""
    backend = str(config.get("backend", "auto")).lower()
    valid_backends = {"auto", "worker", "inprocess"}
    if backend not in valid_backends:
        logger.warning(
            "Unknown shepherd_score.backend=%r, using 'auto'. Valid values: auto|worker|inprocess",
            backend,
        )
        backend = "auto"

    min_shape_score = float(config.get("min_shape_score", 0.5))
    logger.info(
        "Running Shepherd-Score filter (backend=%s, min_shape_score=%.2f)",
        backend,
        min_shape_score,
    )

    if backend in {"auto", "worker"}:
        try:
            return _apply_shepherd_score_filter_worker(
                mols, reference_mol, config, progress_cb=progress_cb
            )
        except Exception as exc:  # noqa: BLE001
            if backend == "worker":
                return _soft_skip_shepherd_filter(mols, str(exc))
            logger.warning(
                "Shepherd worker backend failed, trying in-process backend: %s",
                exc,
            )

    if backend in {"auto", "inprocess"}:
        try:
            return _apply_shepherd_score_filter_inprocess(
                mols, reference_mol, config, progress_cb=progress_cb
            )
        except Exception as exc:  # noqa: BLE001
            return _soft_skip_shepherd_filter(mols, str(exc))

    return _soft_skip_shepherd_filter(mols, "no backend available")


# ---------------------------------------------------------------------------
# PoseBusters / posecheck-fast filter
# ---------------------------------------------------------------------------

_POSEBUSTERS_ERROR: dict[str, Any] = {
    "no_clashes": False,
    "no_volume_clash": False,
    "not_too_far_away": False,
    "no_internal_clash": False,
    "passed": False,
}


def _check_posebusters_fast_single(args: tuple) -> dict[str, Any]:
    """Run posecheck-fast checks for a single molecule.

    Args:
        args: Tuple of (mol, protein_coords, protein_atom_names, config).

    Returns:
        Dict with keys: no_clashes, no_volume_clash, not_too_far_away,
        no_internal_clash, passed, error.
    """
    from posecheck_fast import check_intermolecular_distance

    mol, protein_coords, protein_atom_names, config = args
    try:
        if mol is None or mol.GetNumConformers() == 0:
            return _fail_result("no conformer", **_POSEBUSTERS_ERROR)

        conf = mol.GetConformer()
        pos_pred = np.array([conf.GetPositions()])  # (1, n_atoms, 3)
        atom_names_pred = np.array([atom.GetSymbol() for atom in mol.GetAtoms()])

        # check_intermolecular_distance runs all checks in one call:
        # clashes, volume overlap, distance to protein, and internal geometry.
        # Returns {"results": {"no_clashes": [...], ...}} with one bool per pose.
        with _suppress_posecheck_fast_warnings():
            raw = check_intermolecular_distance(
                mol_orig=mol,
                pos_pred=pos_pred,
                pos_cond=protein_coords,
                atom_names_pred=atom_names_pred,
                atom_names_cond=protein_atom_names,
                clash_cutoff=config.get("clash_cutoff", 0.75),
                clash_cutoff_volume=config.get("volume_clash_cutoff", 0.075),
                max_distance=config.get("max_distance", 5.0),
            )
        r = raw["results"]

        no_clashes = r["no_clashes"][0]
        no_volume_clash = r["no_volume_clash"][0]
        not_too_far = r["not_too_far_away"][0]
        no_internal_clash = r["no_internal_clash"][0]

        passed = no_clashes and no_volume_clash and not_too_far and no_internal_clash
        return _error_result(
            no_clashes=no_clashes,
            no_volume_clash=no_volume_clash,
            not_too_far_away=not_too_far,
            no_internal_clash=no_internal_clash,
            passed=passed,
        )
    except Exception as e:
        return _fail_result(str(e), **_POSEBUSTERS_ERROR)


def apply_posebusters_fast_filter(
    mols: list[Chem.Mol],
    protein_pdb: str | Path,
    config: dict[str, Any],
    progress_cb=None,
) -> pd.DataFrame:
    """Apply pose quality filter using posecheck-fast (~100x faster than PoseCheck).

    Checks for steric clashes (VDW), volume overlap (ShapeTverskyIndex),
    distance to protein, and internal geometry (bond lengths/angles).

    Args:
        mols: List of RDKit molecules with 3D coordinates.
        protein_pdb: Path to protein PDB file.
        config: Filter configuration dict with keys:
            clash_cutoff (float): Relative VDW distance threshold (default 0.75).
            volume_clash_cutoff (float): ShapeTverskyIndex threshold (default 0.075).
            max_distance (float): Max min-distance to protein in Angstrom (default 5.0).

    Returns:
        DataFrame with columns: mol_idx, no_clashes, no_volume_clash,
        not_too_far_away, no_internal_clash, pass_pose_quality.
    """
    logger.info(
        "Running posecheck-fast pose quality filter "
        "(clash_cutoff=%.2f, volume_cutoff=%.3f, max_dist=%.1f)",
        config.get("clash_cutoff", 0.75),
        config.get("volume_clash_cutoff", 0.075),
        config.get("max_distance", 5.0),
    )

    # Load protein once and extract coordinates + atom names.
    # Explicit RemoveHs is needed because MolFromPDBFile(removeHs=True) often
    # retains explicit H from prepared PDB files, causing dimension mismatches
    # inside posecheck-fast's ShapeTverskyIndex call.
    protein_mol = Chem.MolFromPDBFile(str(protein_pdb), removeHs=True, sanitize=False)
    if protein_mol is not None:
        protein_mol = Chem.RemoveHs(protein_mol, sanitize=False)
    if protein_mol is None:
        logger.error("Failed to load protein PDB: %s", protein_pdb)
        return pd.DataFrame(
            {
                "mol_idx": range(len(mols)),
                "pass_pose_quality": False,
            }
        )

    protein_conf = protein_mol.GetConformer()
    protein_coords = np.array(protein_conf.GetPositions())  # (n_atoms, 3)
    protein_atom_names = np.array([atom.GetSymbol() for atom in protein_mol.GetAtoms()])

    items = [(mol, protein_coords, protein_atom_names, config) for mol in mols]
    # Force sequential for posecheck-fast (uses torch internally, forking unsafe)
    raw_results = parallel_map(
        _check_posebusters_fast_single, items, n_jobs=1, progress=progress_cb
    )

    return _aggregate_filter_results(
        raw_results,
        col_mapping={
            "no_clashes": "no_clashes",
            "no_volume_clash": "no_volume_clash",
            "not_too_far_away": "not_too_far_away",
            "no_internal_clash": "no_internal_clash",
        },
        pass_col="pass_pose_quality",
        filter_name="posecheck-fast filter",
    )


def apply_posecheck_fast_filter(
    mols: list[Chem.Mol],
    protein_pdb: str | Path,
    config: dict[str, Any],
) -> pd.DataFrame:
    """Backward-compatible alias for posecheck-fast backend."""
    return apply_posebusters_fast_filter(mols, protein_pdb, config)


# ---------------------------------------------------------------------------
# Conformer deviation filters (symmetry-aware & naive)
# ---------------------------------------------------------------------------


def _resolve_conformer_backend(
    config: dict[str, Any],
) -> Literal["symmetry_rmsd", "naive"]:
    """Determine which conformer RMSD backend to use from config.

    Returns:
        ``"symmetry_rmsd"`` or ``"naive"``.
    """
    backend = str(config.get("backend", "symmetry_rmsd")).lower()
    if backend not in {"symmetry_rmsd", "naive"}:
        logger.warning(
            "Unknown conformer_deviation.backend=%r, using 'symmetry_rmsd'. "
            "Valid values: symmetry_rmsd|naive",
            backend,
        )
        return "symmetry_rmsd"
    return backend  # type: ignore[return-value]


def apply_symmetry_rmsd_filter(
    mols: list[Chem.Mol],
    config: dict[str, Any],
    progress_cb=None,
) -> pd.DataFrame:
    """Apply conformer deviation filter using symmetry-corrected RMSD.

    Uses symmetry-aware RDKit GetBestRMS instead of
    naive AllChem.AlignMol. This gives chemically correct RMSD for
    symmetric molecules (e.g., benzene, biphenyl).

    Args:
        mols: List of RDKit molecules with 3D coordinates.
        config: Filter configuration dict (same keys as conformer_deviation).

    Returns:
        DataFrame with columns: mol_idx, min_conformer_rmsd,
        n_conformers_generated, pass_conformer_deviation.
    """
    num_conformers = config.get("num_conformers", 50)
    max_rmsd = config.get("max_rmsd_to_conformer", 3.0)
    random_seed = config.get("random_seed", 42)
    method = config.get("conformer_method", "ETKDGv3")
    include_hydrogens = bool(config.get("include_hydrogens", False))
    max_matches = int(config.get("max_matches", 10000))
    early_stop_on_pass = bool(config.get("early_stop_on_pass", True))
    use_nvmolkit = bool(config.get("use_nvmolkit", True))

    logger.info(
        "Running symmetry-RMSD conformer filter (max_rmsd=%.1f, n_confs=%d)",
        max_rmsd,
        num_conformers,
    )

    n_jobs = resolve_n_jobs(config)
    items = [
        (
            mol,
            num_conformers,
            max_rmsd,
            method,
            random_seed,
            include_hydrogens,
            max_matches,
            early_stop_on_pass,
            True,  # use_symmetry
            use_nvmolkit,
        )
        for mol in mols
    ]
    raw_results = parallel_map(
        _check_conformer_rmsd_single, items, n_jobs, progress=progress_cb
    )

    return _aggregate_filter_results(
        raw_results,
        col_mapping={
            "min_rmsd": "min_conformer_rmsd",
            "n_conformers_generated": "n_conformers_generated",
        },
        pass_col="pass_conformer_deviation",
        filter_name="Symmetry-RMSD filter",
    )


def apply_conformer_deviation_filter(
    mols: list[Chem.Mol],
    config: dict[str, Any],
    progress_cb=None,
) -> pd.DataFrame:
    """
    Apply conformer deviation filter (naive, identity-permutation RMSD).

    Checks if docked pose is geometrically plausible by comparing
    to generated conformers.

    Args:
        mols: List of RDKit molecules with 3D coordinates
        config: Filter configuration dict

    Returns:
        DataFrame with columns: mol_idx, min_conformer_rmsd,
        n_conformers_generated, pass_conformer_deviation
    """
    num_conformers = config.get("num_conformers", 50)
    max_rmsd = config.get("max_rmsd_to_conformer", 3.0)
    random_seed = config.get("random_seed", 42)
    method = config.get("conformer_method", "ETKDGv3")
    use_nvmolkit = bool(config.get("use_nvmolkit", True))

    logger.info(
        "Running conformer deviation filter (max_rmsd=%.1f, n_confs=%d)",
        max_rmsd,
        num_conformers,
    )

    n_jobs = resolve_n_jobs(config)
    # Naive mode: always include hydrogens, no max_matches, no early stop.
    items = [
        (
            mol,
            num_conformers,
            max_rmsd,
            method,
            random_seed,
            True,  # include_hydrogens (original naive behavior)
            10000,  # max_matches (unused in naive path)
            False,  # early_stop_on_pass (not supported in naive)
            False,  # use_symmetry
            use_nvmolkit,
        )
        for mol in mols
    ]
    raw_results = parallel_map(
        _check_conformer_rmsd_single, items, n_jobs, progress=progress_cb
    )

    return _aggregate_filter_results(
        raw_results,
        col_mapping={
            "min_rmsd": "min_conformer_rmsd",
            "n_conformers_generated": "n_conformers_generated",
        },
        pass_col="pass_conformer_deviation",
        filter_name="Conformer deviation filter",
    )
