test.py

"""Benchmark script to measure query speeds for OMolIndex."""

import time
import numpy as np
import argparse
from pathlib import Path
import json
from typing import Dict, List, Callable

from src.index import OMolIndex
from src.query import QueryFilter


class QueryBenchmark:
    """Benchmark different query types and patterns"""

    def __init__(self, index: OMolIndex):
        self.index = index
        self.results = []

    def time_query(self, name: str, query_filter: QueryFilter, n_runs: int = 5) -> Dict:
        """Time a query multiple times and return statistics"""
        times = []
        result_counts = []

        # Warm-up run
        self.index.query(query_filter)

        # Timed runs
        for _ in range(n_runs):
            start = time.perf_counter()
            indices = self.index.query(query_filter)
            elapsed = time.perf_counter() - start

            times.append(elapsed)
            result_counts.append(len(indices))

        result = {
            "name": name,
            "n_results": result_counts[0],
            "time_mean": np.mean(times),
            "time_std": np.std(times),
            "time_min": np.min(times),
            "time_max": np.max(times),
            "times_all": times,
        }

        self.results.append(result)
        return result

    def print_result(self, result: Dict):
        """Pretty print a single result"""
        print(f"\n{'=' * 70}")
        print(f"Query: {result['name']}")
        print(f"Results found: {result['n_results']:,}")
        print(
            f"Time (mean ± std): {result['time_mean'] * 1000:.2f} ± {result['time_std'] * 1000:.2f} ms"
        )
        print(
            f"Time range: [{result['time_min'] * 1000:.2f}, {result['time_max'] * 1000:.2f}] ms"
        )
        print(
            f"Throughput: {result['n_results'] / result['time_mean']:,.0f} results/sec"
        )

    def run_all_benchmarks(self, n_runs: int = 5):
        """Run a comprehensive suite of benchmarks"""

        print(f"\n{'=' * 70}")
        print("STARTING QUERY BENCHMARKS")
        print(f"{'=' * 70}")
        print(f"Number of runs per query: {n_runs}")
        print(f"Total structures in index: {self.index.h5file.attrs['n_structures']:,}")

        # 1. Simple single-field queries
        print(f"\n{'#' * 70}")
        print("1. SIMPLE SINGLE-FIELD QUERIES")
        print(f"{'#' * 70}")

        result = self.time_query("Charge = 0", QueryFilter(charge_range=(0, 0)), n_runs)
        self.print_result(result)

        result = self.time_query("Spin = 1", QueryFilter(spin_range=(1, 1)), n_runs)
        self.print_result(result)

        result = self.time_query(
            "Num atoms 10-20", QueryFilter(num_atoms_range=(10, 20)), n_runs
        )
        self.print_result(result)

        result = self.time_query(
            "HOMO-LUMO gap > 5 eV",
            QueryFilter(homo_lumo_gap_range=(5.0, float("inf"))),
            n_runs,
        )
        self.print_result(result)

        result = self.time_query(
            "Unrestricted = True", QueryFilter(unrestricted=True), n_runs
        )
        self.print_result(result)

        result = self.time_query(
            "Has warnings = True", QueryFilter(has_warnings=True), n_runs
        )
        self.print_result(result)

        # 2. Element queries (using binary mask - should be fast!)
        print(f"\n{'#' * 70}")
        print("2. ELEMENT QUERIES (Binary Mask)")
        print(f"{'#' * 70}")

        result = self.time_query("Contains Carbon", QueryFilter(elements=["C"]), n_runs)
        self.print_result(result)

        result = self.time_query(
            "Contains Carbon AND Nitrogen", QueryFilter(elements=["C", "N"]), n_runs
        )
        self.print_result(result)

        result = self.time_query(
            "Contains C, N, O, H", QueryFilter(elements=["C", "N", "O", "H"]), n_runs
        )
        self.print_result(result)

        result = self.time_query(
            "Exclude Transition Metals (Fe, Co, Ni)",
            QueryFilter(exclude_elements=["Fe", "Co", "Ni"]),
            n_runs,
        )
        self.print_result(result)

        result = self.time_query(
            "Contains Rare Element (Ir)", QueryFilter(elements=["Ir"]), n_runs
        )
        self.print_result(result)

        # 3. Multi-field queries (realistic use cases)
        print(f"\n{'#' * 70}")
        print("3. MULTI-FIELD QUERIES")
        print(f"{'#' * 70}")

        result = self.time_query(
            "Charge=0, Spin=1, 10-20 atoms",
            QueryFilter(
                charge_range=(0, 0), spin_range=(1, 1), num_atoms_range=(10, 20)
            ),
            n_runs,
        )
        self.print_result(result)

        result = self.time_query(
            "Organic (C,H,N,O), charge=0, <30 atoms",
            QueryFilter(
                elements=["C", "H", "N", "O"],
                charge_range=(0, 0),
                num_atoms_range=(1, 30),
            ),
            n_runs,
        )
        self.print_result(result)

        result = self.time_query(
            "Small organic, large gap",
            QueryFilter(
                elements=["C", "H", "N", "O"],
                num_atoms_range=(5, 15),
                homo_lumo_gap_range=(6.0, float("inf")),
                charge_range=(0, 0),
            ),
            n_runs,
        )
        self.print_result(result)

        result = self.time_query(
            "Unrestricted, high spin, no warnings",
            QueryFilter(unrestricted=True, spin_range=(2, 10), has_warnings=False),
            n_runs,
        )
        self.print_result(result)

        # 4. Complex queries
        print(f"\n{'#' * 70}")
        print("4. COMPLEX QUERIES")
        print(f"{'#' * 70}")

        result = self.time_query(
            "Complex: 6 conditions",
            QueryFilter(
                charge_range=(0, 0),
                spin_range=(1, 1),
                num_atoms_range=(10, 25),
                homo_lumo_gap_range=(3.0, 8.0),
                unrestricted=False,
                has_warnings=False,
            ),
            n_runs,
        )
        self.print_result(result)

        result = self.time_query(
            "Transition metal complex",
            QueryFilter(
                elements=["Fe", "C", "N"],
                charge_range=(-1, 1),
                spin_range=(1, 5),
                num_atoms_range=(10, 50),
            ),
            n_runs,
        )
        self.print_result(result)

        # 5. String matching queries (slower - requires decoding)
        print(f"\n{'#' * 70}")
        print("5. STRING MATCHING QUERIES")
        print(f"{'#' * 70}")

        result = self.time_query(
            "Composition pattern 'C6'", QueryFilter(composition_pattern="C6"), n_runs
        )
        self.print_result(result)

        # Check available domains first
        domains = self.index.get_unique_domains()
        if domains:
            result = self.time_query(
                f"Data domain filter: {domains[0]}",
                QueryFilter(data_domains=[domains[0]]),
                n_runs,
            )
            self.print_result(result)

        # 6. Limited result queries
        print(f"\n{'#' * 70}")
        print("6. LIMITED RESULT QUERIES")
        print(f"{'#' * 70}")

        result = self.time_query(
            "First 100 results only", QueryFilter(max_results=100), n_runs
        )
        self.print_result(result)

        result = self.time_query(
            "First 1000 with element filter",
            QueryFilter(elements=["C"], max_results=1000),
            n_runs,
        )
        self.print_result(result)

        # 7. Full scan (worst case)
        print(f"\n{'#' * 70}")
        print("7. FULL SCAN (Baseline)")
        print(f"{'#' * 70}")

        result = self.time_query("No filters (full scan)", QueryFilter(), n_runs)
        self.print_result(result)

        # Summary
        self.print_summary()

    def print_summary(self):
        """Print summary statistics of all benchmarks"""
        print(f"\n{'=' * 70}")
        print("BENCHMARK SUMMARY")
        print(f"{'=' * 70}")

        times = [r["time_mean"] for r in self.results]
        result_counts = [r["n_results"] for r in self.results]

        print(f"\nTotal queries run: {len(self.results)}")
        print(f"Mean query time: {np.mean(times) * 1000:.2f} ms")
        print(f"Median query time: {np.median(times) * 1000:.2f} ms")
        print(
            f"Fastest query: {np.min(times) * 1000:.2f} ms ({self.results[np.argmin(times)]['name']})"
        )
        print(
            f"Slowest query: {np.max(times) * 1000:.2f} ms ({self.results[np.argmax(times)]['name']})"
        )

        print(f"\nMean results per query: {np.mean(result_counts):,.0f}")
        print(f"Median results per query: {np.median(result_counts):,.0f}")

        # Find queries by speed categories
        fast_queries = [r for r in self.results if r["time_mean"] < 0.01]  # <10ms
        medium_queries = [
            r for r in self.results if 0.01 <= r["time_mean"] < 0.1
        ]  # 10-100ms
        slow_queries = [r for r in self.results if r["time_mean"] >= 0.1]  # >100ms

        print(f"\nQuery speed distribution:")
        print(f"  Fast (<10ms): {len(fast_queries)}")
        print(f"  Medium (10-100ms): {len(medium_queries)}")
        print(f"  Slow (>100ms): {len(slow_queries)}")

    def save_results(self, output_path: str):
        """Save benchmark results to JSON"""
        with open(output_path, "w") as f:
            json.dump(self.results, f, indent=2)
        print(f"\nResults saved to {output_path}")


def main():
    parser = argparse.ArgumentParser(
        description="Benchmark query performance for OMolIndex"
    )
    parser.add_argument(
        "--index-path", required=True, help="Path to the HDF5 index file"
    )
    parser.add_argument(
        "--n-runs", type=int, default=5, help="Number of runs per query (default: 5)"
    )
    parser.add_argument(
        "--output",
        default="benchmark_results.json",
        help="Output file for results (default: benchmark_results.json)",
    )

    args = parser.parse_args()

    # Load index
    print(f"Loading index from {args.index_path}...")
    index = OMolIndex()
    index.load(args.index_path)

    try:
        # Run benchmarks
        benchmark = QueryBenchmark(index)
        benchmark.run_all_benchmarks(n_runs=args.n_runs)

        # Save results
        benchmark.save_results(args.output)

    finally:
        index.close()


if __name__ == "__main__":
    main()