evaluate_detector.py

#!/usr/bin/env python3
"""Evaluate a user-trained biological-marker detector on a labeled dataset."""

from __future__ import annotations

import argparse
import datetime as dt
import json
import os
import shutil
import subprocess
import sys
import time
from pathlib import Path


REPO_ROOT = Path(__file__).resolve().parent
PYTHONPATH_ROOT = REPO_ROOT / "src"
THIRD_PARTY_ROOT = REPO_ROOT / "third_party"
REGISTRY_PATH = REPO_ROOT / "pretrained_models" / "registry.json"
CATEGORIES = ("respiratory", "vocal", "temporal")


def _load_registry() -> dict[str, dict[str, str]]:
    """Load optional model aliases from ``pretrained_models/registry.json`` if present."""
    if not REGISTRY_PATH.is_file():
        return {}
    return json.loads(REGISTRY_PATH.read_text(encoding="utf-8"))


def _safe_name(value: str) -> str:
    """
    Convert a model alias/path into a short filesystem-safe run-directory label.
    """
    cleaned = "".join(ch if ch.isalnum() or ch in ("-", "_") else "_" for ch in str(value))
    return cleaned.strip("_") or "custom_model"


def _resolve_user_path(path_text: str) -> Path:
    """
    Resolve user-provided paths relative to the repository root first, then cwd.
    """
    raw = Path(path_text).expanduser()
    if raw.is_absolute():
        return raw.resolve()
    final_relative = (REPO_ROOT / raw).resolve()
    if final_relative.exists():
        return final_relative
    return raw.resolve()


def _infer_classifier_from_model_file(model_path: Path, classifier_hint: str | None) -> str:
    """
    Infer classifier family from the pickle filename, with an optional explicit override.
    """
    if classifier_hint:
        return classifier_hint
    filename = model_path.name.lower()
    if "logreg" in filename:
        return "logreg"
    if "svm" in filename:
        return "svm"
    raise ValueError(
        f"Cannot infer classifier type from {model_path}. "
        "Pass --classifier logreg or --classifier svm."
    )


def _candidate_custom_models(
    path: Path,
    classifier_hint: str | None = None,
) -> list[tuple[str, Path, Path | None, Path | None]]:
    """
    Return possible custom model candidates as "(classifier, model_pkl, lda_pkl, report_json)".

    Supported inputs:
    - a direct "logreg_model.pkl" / "svm_model.pkl" file,
    - a model directory printed by "train_detector.py" (".../models/logreg"),
    - a whole training run directory ("runs/train_logreg/<timestamp>").
    """
    if path.is_file():
        classifier = _infer_classifier_from_model_file(path, classifier_hint)
        model_dir = path.parent
        lda_path = model_dir / "lda_model.pkl"
        report_path = model_dir / "report.json"
        return [
            (
                classifier,
                path,
                lda_path if lda_path.is_file() else None,
                report_path if report_path.is_file() else None,
            )
        ]

    candidates: list[tuple[str, Path, Path | None, Path | None]] = []
    search_roots = [path, path / "models" / "logreg", path / "models" / "svm"]
    for root in search_roots:
        for classifier, filename in (("logreg", "logreg_model.pkl"), ("svm", "svm_model.pkl")):
            model_path = root / filename
            if not model_path.is_file():
                continue
            lda_path = root / "lda_model.pkl"
            report_path = root / "report.json"
            candidates.append(
                (
                    classifier,
                    model_path,
                    lda_path if lda_path.is_file() else None,
                    report_path if report_path.is_file() else None,
                )
            )
    return candidates


def _resolve_model_selection(
    registry: dict[str, dict[str, str]],
    model_text: str,
    classifier_hint: str | None,
) -> dict[str, object]:
    """
    Resolve ``--model`` to a concrete model pickle and optional LDA artifact.

    Accepts paths to model files or directories produced by ``train_detector.py``,
    or an optional alias from ``pretrained_models/registry.json`` when that file exists.
    """
    if model_text in registry:
        selected = dict(registry[model_text])
        return {
            "name": model_text,
            "source": "bundled",
            "classifier": str(selected["classifier"]),
            "description": str(selected.get("description", "")),
            "model_path": (REPO_ROOT / selected["model_path"]).resolve(),
            "lda_model_path": None,
            "report_path": (REPO_ROOT / selected["report_path"]).resolve()
            if selected.get("report_path")
            else None,
        }

    path = _resolve_user_path(model_text)
    candidates = _candidate_custom_models(path, classifier_hint)
    if classifier_hint:
        candidates = [candidate for candidate in candidates if candidate[0] == classifier_hint]
    if not candidates:
        raise FileNotFoundError(
            f"Could not find a model under {path}. Expected logreg_model.pkl or svm_model.pkl, "
            "either directly, inside a model directory, or under models/logreg|models/svm."
        )
    if len(candidates) > 1:
        formatted = ", ".join(f"{classifier}:{model_path}" for classifier, model_path, _lda, _report in candidates)
        raise ValueError(
            f"Multiple models found under {path}: {formatted}. "
            "Pass --classifier logreg or --classifier svm."
        )

    classifier, model_path, lda_path, report_path = candidates[0]
    return {
        "name": classifier,
        "source": "custom",
        "classifier": classifier,
        "description": f"User-provided model from {model_path}",
        "model_path": model_path.resolve(),
        "lda_model_path": lda_path.resolve() if lda_path is not None else None,
        "report_path": report_path.resolve() if report_path is not None else None,
    }


def _make_env() -> dict[str, str]:
    """
    Build a subprocess environment with ``third_party/`` and ``src/`` on PYTHONPATH.
    """
    env = os.environ.copy()
    old_pythonpath = env.get("PYTHONPATH", "")
    prefix = os.pathsep.join((str(THIRD_PARTY_ROOT), str(PYTHONPATH_ROOT)))
    env["PYTHONPATH"] = prefix if not old_pythonpath else prefix + os.pathsep + old_pythonpath
    env.setdefault("NUMBA_CACHE_DIR", "/tmp/dp_biomarkers_numba_cache")
    return env


def _run_step(name: str, command: list[str], env: dict[str, str]) -> None:
    """Run one command and stream the original project logs to the terminal."""
    started = time.time()
    print(f"\n[{dt.datetime.now().strftime('%H:%M:%S')}] START: {name}", flush=True)
    print("Command: " + " ".join(command), flush=True)
    completed = subprocess.run(command, env=env, check=False)
    elapsed = time.time() - started
    if completed.returncode != 0:
        raise SystemExit(
            f"\nStep failed: {name}\n"
            f"Return code: {completed.returncode}\n"
            f"Elapsed: {elapsed:.1f} s"
        )
    print(f"[{dt.datetime.now().strftime('%H:%M:%S')}] DONE: {name} ({elapsed:.1f} s)", flush=True)


def _has_preprocessed_marker_splits(path: Path) -> bool:
    """
    Return True when the user passed a preprocessed marker dataset.

    The eval wrapper expects labels and split names to be present in the NPZ
    files, exactly as produced by the original preprocessing CLI.
    """
    return all(
        (path / category / f"{category}_train.npz").is_file()
        for category in CATEGORIES
    )


def _raw_csv_root(path: Path) -> Path | None:
    """
    Detect already extracted raw marker CSV files.

    This is useful when the user has marker tables but has not yet fitted the
    preprocessing artifacts needed by the classifier.
    """
    candidates = [path, path / "raw_hybrid_synced"]
    for candidate in candidates:
        if all((candidate / f"{category}.csv").is_file() for category in CATEGORIES):
            return candidate
    return None


def _copy_available_preprocessed_splits(source: Path, target: Path) -> None:
    """
    Copy all available NPZ splits to a clean evaluation workspace.

    Unlike the training wrapper, evaluation keeps any additional split in the
    provided dataset. This lets a user evaluate `train`, `val`, `test`, or a
    custom split name if those NPZ files are present.
    """
    if target.exists():
        shutil.rmtree(target)
    for category in CATEGORIES:
        source_category = source / category
        target_category = target / category
        target_category.mkdir(parents=True, exist_ok=True)
        if not source_category.is_dir():
            raise FileNotFoundError(f"Missing category directory: {source_category}")
        copied = 0
        for npz_file in sorted(source_category.glob(f"{category}_*.npz")):
            shutil.copy2(npz_file, target_category / npz_file.name)
            copied += 1
        if copied == 0:
            raise FileNotFoundError(f"No NPZ split files found in {source_category}")
        for suffix in ("metadata.json", "preprocessor.pkl"):
            optional = source_category / f"{category}_{suffix}"
            if optional.is_file():
                shutil.copy2(optional, target_category / optional.name)


def _preprocess_raw_csv(raw_csv_root: Path, target: Path, env: dict[str, str]) -> None:
    """
    Fit preprocessing for each marker category from raw marker CSV files.

    The preprocessor is fitted only on the train split and then applied to all
    other available splits. This creates the same NPZ format expected by the
    original training/evaluation code.
    """
    for category in CATEGORIES:
        _run_step(
            f"preprocess {category}",
            [
                sys.executable,
                "-m",
                "dp_biomarkers.cli.preprocess_features_category",
                "--input",
                str(raw_csv_root / f"{category}.csv"),
                "--category",
                category,
                "--output-dir",
                str(target / category),
                "--fit-split",
                "train",
                "--normalization",
                "standard",
                "--handle-missing",
                "median",
            ],
            env,
        )


def _extract_and_preprocess_audio_dataset(dataset_root: Path, run_dir: Path, env: dict[str, str]) -> Path:
    """
    Build an evaluation marker dataset from a raw-audio directory.

    The audio directory must contain enough label information for the manifest
    builder: either top-level `real`/`authentic` and `deepfake` directories, or
    source subdirectories that contain those branches. The manifest builder then
    creates train/validation/test splits, marker extraction creates CSV files,
    and preprocessing converts them into NPZ marker matrices.
    """
    manifest = run_dir / "manifest.csv"
    raw_dir = run_dir / "raw_hybrid_synced"
    preprocessed_dir = run_dir / "preprocessed"

    _run_step(
        "create evaluation manifest",
        [
            sys.executable,
            "-m",
            "dp_biomarkers.cli.make_manifest",
            "--output",
            str(manifest),
            "--root",
            str(dataset_root),
            "--stratified-split",
            "--val-size",
            "0.15",
            "--test-size",
            "0.15",
            "--seed",
            "42",
        ],
        env,
    )
    _run_step(
        "extract biological markers",
        [
            sys.executable,
            "-m",
            "dp_biomarkers.cli.extract_hybrid",
            "--manifest",
            str(manifest),
            "--output-dir",
            str(raw_dir),
            "--min-duration",
            "0.1",
            "--progress-every",
            "25",
        ],
        env,
    )
    _preprocess_raw_csv(raw_dir, preprocessed_dir, env)
    return preprocessed_dir


def _prepare_dataset(dataset: Path, run_dir: Path, env: dict[str, str]) -> Path:
    """
    Accept a preprocessed marker dataset, raw marker CSV files, or raw audio.
    """
    prepared = run_dir / "preprocessed"
    if _has_preprocessed_marker_splits(dataset):
        _copy_available_preprocessed_splits(dataset, prepared)
        return prepared

    raw_csv = _raw_csv_root(dataset)
    if raw_csv is not None:
        _preprocess_raw_csv(raw_csv, prepared, env)
        return prepared

    if not dataset.is_dir():
        raise FileNotFoundError(f"Dataset path does not exist or is not a directory: {dataset}")
    return _extract_and_preprocess_audio_dataset(dataset, run_dir, env)


def _print_report_summary(report_path: Path) -> None:
    """Print a compact metric summary after the original CLI writes report.json."""
    if not report_path.is_file():
        print(f"Report not found: {report_path}", flush=True)
        return
    report = json.loads(report_path.read_text(encoding="utf-8"))
    print("\n=== EVALUATION SUMMARY ===", flush=True)
    for split_name, metrics in report.get("metrics", {}).items():
        print(
            f"{split_name}: "
            f"EER={float(metrics.get('eer', 0.0)):.4f}, "
            f"ROC-AUC={float(metrics.get('roc_auc', 0.0)):.4f}, "
            f"ACC={float(metrics.get('accuracy', 0.0)):.4f}, "
            f"F1={float(metrics.get('f1', 0.0)):.4f}",
            flush=True,
        )


def main() -> None:
    registry = _load_registry()
    parser = argparse.ArgumentParser(
        description="Evaluate a user-trained biological-marker detector."
    )
    parser.add_argument(
        "--model",
        required=True,
        help=(
            "Path to a trained model: logreg_model.pkl, svm_model.pkl, "
            "models/logreg|models/svm, or a full runs/train_* run directory."
        ),
    )
    parser.add_argument(
        "--classifier",
        choices=["logreg", "svm"],
        default=None,
        help="Required only when --model points to a custom path containing both logreg and svm models, or an ambiguous .pkl filename.",
    )
    parser.add_argument("--dataset", type=Path, required=True, help="Path to preprocessed markers, raw marker CSVs, or a labeled audio dataset.")
    args = parser.parse_args()

    selected = _resolve_model_selection(registry, str(args.model), args.classifier)
    classifier = str(selected["classifier"])
    dataset = args.dataset.expanduser().resolve()
    model_label = _safe_name(str(selected["name"]))
    run_dir = REPO_ROOT / "runs" / f"eval_{model_label}" / dt.datetime.now().strftime("%Y%m%d_%H%M")
    run_dir.mkdir(parents=True, exist_ok=True)
    env = _make_env()
    prepared = _prepare_dataset(dataset, run_dir, env)

    model_path = Path(str(selected["model_path"])).resolve()
    lda_model_path = selected.get("lda_model_path")
    module = "dp_biomarkers.cli.train_logreg" if classifier == "logreg" else "dp_biomarkers.cli.train_svm"
    model_dir = run_dir / "models" / classifier
    merged_dir = run_dir / "merged"
    command = [
        sys.executable,
        "-m",
        module,
        "--respiratory-dir",
        str(prepared / "respiratory"),
        "--vocal-dir",
        str(prepared / "vocal"),
        "--temporal-dir",
        str(prepared / "temporal"),
        "--output-dir",
        str(model_dir),
        "--merged-output-dir",
        str(merged_dir),
        "--train-split",
        "train",
        "--model-in",
        str(model_path),
        "--threshold-source",
        "fixed_0.5",
    ]
    if lda_model_path is not None:
        command.extend(["--lda-model-in", str(Path(str(lda_model_path)).resolve())])
    if classifier == "svm":
        command.extend(["--kernel", "rbf", "--max-iter", "20000"])

    print("=== Biological Marker Detector Evaluation ===", flush=True)
    print(f"Model: {args.model}", flush=True)
    print(f"Model source: {selected['source']}", flush=True)
    print(f"Description: {selected.get('description', '')}", flush=True)
    print(f"Model pickle: {model_path}", flush=True)
    if lda_model_path is not None:
        print(f"LDA projection: {lda_model_path}", flush=True)
    print(f"Dataset: {dataset}", flush=True)
    print(f"Output run directory: {run_dir}", flush=True)
    _run_step("evaluate model", command, env)

    report_path = model_dir / "report.json"
    _print_report_summary(report_path)
    print("\n=== DONE ===", flush=True)
    print(f"Run directory: {run_dir}", flush=True)
    print(f"Report: {report_path}", flush=True)


if __name__ == "__main__":
    main()