reproducibility.py

"""Logging and reproducibility tracking for simulations.

Every simulation run captures:
  • Physics module versions
  • Solver algorithm + parameters
  • Configuration snapshot
  • Random seed (for stochastic processes)
  • Git commit hash and dirty status
  • System info (CPU, memory, Python version)
  • Numerical precision (float32 vs float64)
  • Timestamp and duration
  • User and hostname

This enables:
  ✓ Exact reproduction of any run
  ✓ Tracking physics model evolution
  ✓ Audit trail for research
  ✓ Comparison of runs with different configs
  ✓ Detection of bitwise identical simulations (numerical reproducibility)
"""

from dataclasses import dataclass, field, asdict
from typing import Dict, Any, Optional
import json
import time
import sys
import platform
import hashlib
import subprocess
from pathlib import Path
from datetime import datetime


# ============================================================================
# System Metadata
# ============================================================================

@dataclass(frozen=True)
class SystemMetadata:
    """Capture machine and environment information."""
    
    hostname: str
    platform_system: str
    platform_release: str
    python_version: str
    python_implementation: str
    
    # CPU info
    cpu_count: int
    
    # Precision
    float_dtype: str = "float64"  # default
    machine_epsilon: float = 2.220446049250313e-16  # default for float64
    
    # Numpy version (if available)
    numpy_version: Optional[str] = None
    scipy_version: Optional[str] = None
    
    @staticmethod
    def capture() -> "SystemMetadata":
        """Capture current system metadata."""
        try:
            import numpy as np
            numpy_version = str(np.__version__)
            float_dtype = str(np.float64)
            machine_epsilon = float(np.finfo(np.float64).eps)
        except ImportError:
            numpy_version = None
            float_dtype = "float64"
            machine_epsilon = 2.220446049250313e-16
        
        try:
            import scipy
            scipy_version = str(scipy.__version__)
        except ImportError:
            scipy_version = None
        
        import os
        
        return SystemMetadata(
            hostname=platform.node(),
            platform_system=platform.system(),
            platform_release=platform.release(),
            python_version=platform.python_version(),
            python_implementation=platform.python_implementation(),
            cpu_count=os.cpu_count() or 1,
            numpy_version=numpy_version,
            scipy_version=scipy_version,
            float_dtype=float_dtype,
            machine_epsilon=machine_epsilon,
        )


@dataclass(frozen=True)
class GitMetadata:
    """Capture Git repository state."""
    
    commit_hash: str  # Full SHA-1 (40 chars)
    commit_short: str  # Short SHA-1 (7 chars)
    branch: str
    is_dirty: bool  # Uncommitted changes
    uncommitted_count: int
    untracked_count: int
    
    @staticmethod
    def capture(repo_path: Optional[str] = None) -> "GitMetadata":
        """Capture current Git state."""
        if repo_path is None:
            repo_path = Path(__file__).parent.parent.parent  # Go up to workspace root
        
        try:
            # Get commit hash
            result = subprocess.run(
                ["git", "rev-parse", "HEAD"],
                cwd=repo_path,
                capture_output=True,
                text=True,
                timeout=5,
            )
            commit_hash = result.stdout.strip()
            commit_short = commit_hash[:7]
            
            # Get branch
            result = subprocess.run(
                ["git", "rev-parse", "--abbrev-ref", "HEAD"],
                cwd=repo_path,
                capture_output=True,
                text=True,
                timeout=5,
            )
            branch = result.stdout.strip()
            
            # Check if dirty
            result = subprocess.run(
                ["git", "status", "--porcelain"],
                cwd=repo_path,
                capture_output=True,
                text=True,
                timeout=5,
            )
            status_lines = result.stdout.strip().split("\n") if result.stdout.strip() else []
            uncommitted = [l for l in status_lines if l[0] in "M D"]
            untracked = [l for l in status_lines if l[0] == "?"]
            
            return GitMetadata(
                commit_hash=commit_hash,
                commit_short=commit_short,
                branch=branch,
                is_dirty=len(uncommitted) > 0,
                uncommitted_count=len(uncommitted),
                untracked_count=len(untracked),
            )
        except Exception as e:
            # Git not available or repo not found
            return GitMetadata(
                commit_hash="unknown",
                commit_short="unknown",
                branch="unknown",
                is_dirty=False,
                uncommitted_count=0,
                untracked_count=0,
            )


@dataclass(frozen=True)
class SolverMetadata:
    """Capture numerical integration method and parameters."""
    
    solver_type: str  # "euler", "rk4", "rk45", "bdf", etc.
    dt_nominal: float  # Nominal timestep [s]
    dt_min: Optional[float] = None
    dt_max: Optional[float] = None
    tolerance_relative: Optional[float] = None
    tolerance_absolute: Optional[float] = None
    max_steps: int = 1000000
    adaptive: bool = False
    
    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)


@dataclass(frozen=True)
class PhysicsModuleInfo:
    """Metadata for one physics module."""
    
    name: str
    version: str
    implementation: str  # "analytical", "numerical", "data-driven", "hybrid"
    enabled: bool
    parameters: Dict[str, Any] = field(default_factory=dict)


@dataclass(frozen=True)
class PhysicsMetadata:
    """Capture all physics modules and versions."""
    
    modules: Dict[str, PhysicsModuleInfo] = field(default_factory=dict)
    
    @staticmethod
    def from_engine(engine) -> "PhysicsMetadata":
        """Extract metadata from SimulationEngine."""
        modules = {}
        
        for module_name, module in engine.physics_modules.items():
            # Try to extract version info
            version = getattr(module, "version", "unknown")
            impl = getattr(module, "implementation", "unknown")
            
            modules[module_name] = PhysicsModuleInfo(
                name=module_name,
                version=version,
                implementation=impl,
                enabled=True,
                parameters=getattr(module, "config", {}),
            )
        
        return PhysicsMetadata(modules=modules)


@dataclass(frozen=True)
class RandomSeedMetadata:
    """Capture random number generator state."""
    
    numpy_seed: Optional[int] = None
    python_seed: Optional[int] = None
    random_state_hash: Optional[str] = None  # Hash of RNG state for validation
    
    @staticmethod
    def capture(seed: Optional[int] = None) -> "RandomSeedMetadata":
        """Capture RNG state."""
        import random
        
        if seed is not None:
            random.seed(seed)
            try:
                import numpy as np
                np.random.seed(seed)
                numpy_seed = seed
            except ImportError:
                numpy_seed = None
        else:
            numpy_seed = None
        
        return RandomSeedMetadata(
            numpy_seed=numpy_seed,
            python_seed=seed,
            random_state_hash=None,  # Could add state hash if needed
        )


# ============================================================================
# Complete Run Metadata
# ============================================================================

@dataclass
class SimulationMetadata:
    """Complete metadata for one simulation run."""
    
    # Identifiers
    run_id: str  # UUID for this run
    mission_name: str
    timestamp: str  # ISO 8601
    
    # Configuration
    config_hash: str  # SHA-256 of config JSON
    config_snapshot: Dict[str, Any]  # Full config for audit
    
    # System
    system: SystemMetadata
    git: GitMetadata
    hostname: str
    working_directory: str
    
    # Physics & Solver
    physics: PhysicsMetadata
    solver: SolverMetadata
    
    # Random seed
    random_seed: RandomSeedMetadata
    
    # Execution
    start_time: float  # Unix timestamp
    end_time: Optional[float] = None
    elapsed_seconds: Optional[float] = None
    
    # Results reference
    output_files: Dict[str, str] = field(default_factory=dict)  # { "hdf5": "file.h5", ... }
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to JSON-serializable dict."""
        return {
            "run_id": self.run_id,
            "mission_name": self.mission_name,
            "timestamp": self.timestamp,
            "config_hash": self.config_hash,
            "config_snapshot": self.config_snapshot,
            "system": asdict(self.system),
            "git": asdict(self.git),
            "hostname": self.hostname,
            "working_directory": self.working_directory,
            "physics": {
                k: asdict(v) for k, v in self.physics.modules.items()
            },
            "solver": self.solver.to_dict(),
            "random_seed": asdict(self.random_seed),
            "start_time": self.start_time,
            "end_time": self.end_time,
            "elapsed_seconds": self.elapsed_seconds,
            "output_files": self.output_files,
        }
    
    def to_json(self) -> str:
        """Serialize to JSON string."""
        return json.dumps(self.to_dict(), indent=2, default=str)
    
    def save_json(self, path: str) -> None:
        """Save metadata to JSON file."""
        Path(path).parent.mkdir(parents=True, exist_ok=True)
        with open(path, "w") as f:
            f.write(self.to_json())
    
    @staticmethod
    def create(
        mission_name: str,
        config: Dict[str, Any],
        engine,
        solver_config: Optional[Dict[str, Any]] = None,
        random_seed: Optional[int] = None,
        run_id: Optional[str] = None,
    ) -> "SimulationMetadata":
        """Create metadata for a new simulation run."""
        import uuid
        from pathlib import Path
        
        if run_id is None:
            run_id = str(uuid.uuid4())
        
        # Config hash
        config_json = json.dumps(config, sort_keys=True)
        config_hash = hashlib.sha256(config_json.encode()).hexdigest()
        
        # Solver config
        if solver_config is None:
            solver_config = {"solver_type": "euler", "dt_nominal": 10.0}
        
        solver_meta = SolverMetadata(**solver_config)
        
        return SimulationMetadata(
            run_id=run_id,
            mission_name=mission_name,
            timestamp=datetime.utcnow().isoformat() + "Z",
            config_hash=config_hash,
            config_snapshot=config,
            system=SystemMetadata.capture(),
            git=GitMetadata.capture(),
            hostname=platform.node(),
            working_directory=str(Path.cwd()),
            physics=PhysicsMetadata.from_engine(engine),
            solver=solver_meta,
            random_seed=RandomSeedMetadata.capture(random_seed),
            start_time=time.time(),
        )
    
    def mark_complete(self) -> None:
        """Mark run as complete."""
        self.end_time = time.time()
        self.elapsed_seconds = self.end_time - self.start_time


# ============================================================================
# Research Database Logging
# ============================================================================

class ResearchDatabase:
    """Centralized logging for research reproducibility."""
    
    def __init__(self, db_path: str = "research-database"):
        """Initialize research database."""
        self.db_path = Path(db_path)
        self.db_path.mkdir(parents=True, exist_ok=True)
        self.metadata_dir = self.db_path / "metadata"
        self.metadata_dir.mkdir(exist_ok=True)
        self.results_dir = self.db_path / "results"
        self.results_dir.mkdir(exist_ok=True)
        
        # Catalog file
        self.catalog_path = self.db_path / "catalog.json"
        self.catalog = self._load_catalog()
    
    def _load_catalog(self) -> Dict[str, Any]:
        """Load existing catalog or create new one."""
        if self.catalog_path.exists():
            with open(self.catalog_path, "r") as f:
                return json.load(f)
        
        return {
            "version": "1.0",
            "created_at": datetime.utcnow().isoformat() + "Z",
            "runs": {},
        }
    
    def _save_catalog(self) -> None:
        """Save catalog to JSON."""
        with open(self.catalog_path, "w") as f:
            json.dump(self.catalog, f, indent=2, default=str)
    
    def register_run(self, metadata: SimulationMetadata) -> None:
        """Register a simulation run in the catalog."""
        self.catalog["runs"][metadata.run_id] = {
            "mission": metadata.mission_name,
            "timestamp": metadata.timestamp,
            "config_hash": metadata.config_hash,
            "duration_sec": metadata.elapsed_seconds,
            "git_commit": metadata.git.commit_short,
            "metadata_file": f"metadata/{metadata.run_id}.json",
            "output_files": metadata.output_files,
        }
        self._save_catalog()
        
        # Save metadata file
        metadata_file = self.metadata_dir / f"{metadata.run_id}.json"
        metadata.save_json(str(metadata_file))
    
    def query_runs(self, config_hash: Optional[str] = None) -> list:
        """Query runs by config hash (for reproducibility)."""
        if config_hash is None:
            return list(self.catalog["runs"].values())
        
        return [
            v for v in self.catalog["runs"].values()
            if v["config_hash"] == config_hash
        ]
    
    def get_metadata(self, run_id: str) -> Optional[Dict[str, Any]]:
        """Load metadata for a specific run."""
        metadata_file = self.metadata_dir / f"{run_id}.json"
        if not metadata_file.exists():
            return None
        
        with open(metadata_file, "r") as f:
            return json.load(f)


# ============================================================================
# Reproducibility Report
# ============================================================================

def create_reproducibility_report(
    metadata: SimulationMetadata,
    result_summary: Dict[str, Any],
) -> str:
    """Create human-readable reproducibility report."""
    
    lines = [
        "",
        "="*70,
        "SIMULATION REPRODUCIBILITY REPORT",
        "="*70,
        "",
        f"Run ID:           {metadata.run_id}",
        f"Mission:          {metadata.mission_name}",
        f"Timestamp:        {metadata.timestamp}",
        "",
        "EXACT REPRODUCTION:",
        f"  Config SHA-256: {metadata.config_hash}",
        f"  Git Commit:     {metadata.git.commit_hash}",
        f"  Branch:         {metadata.git.branch}",
        f"  Dirty:          {metadata.git.is_dirty}",
        "",
        "SYSTEM ENVIRONMENT:",
        f"  Hostname:       {metadata.system.hostname}",
        f"  Platform:       {metadata.system.platform_system} {metadata.system.platform_release}",
        f"  Python:         {metadata.system.python_version} ({metadata.system.python_implementation})",
        f"  CPUs:           {metadata.system.cpu_count}",
        f"  Float Type:     {metadata.system.float_dtype}",
        f"  Machine Epsilon:{metadata.system.machine_epsilon:.2e}",
        "",
        "PHYSICS MODULES:",
    ]
    
    for name, module_info in metadata.physics.modules.items():
        lines.append(f"  {name}:")
        lines.append(f"    Version: {module_info.version}")
        lines.append(f"    Impl:    {module_info.implementation}")
    
    lines.extend([
        "",
        "SOLVER:",
        f"  Type:      {metadata.solver.solver_type}",
        f"  dt:        {metadata.solver.dt_nominal} s",
        f"  Adaptive:  {metadata.solver.adaptive}",
        "",
        "RANDOM SEED:",
        f"  NumPy:     {metadata.random_seed.numpy_seed}",
        f"  Python:    {metadata.random_seed.python_seed}",
        "",
        "EXECUTION:",
        f"  Duration:  {metadata.elapsed_seconds:.1f} seconds",
        "",
        "RESULTS:",
    ])
    
    for key, val in result_summary.items():
        if isinstance(val, float):
            lines.append(f"  {key}: {val:.6e}")
        else:
            lines.append(f"  {key}: {val}")
    
    lines.extend([
        "",
        "TO REPRODUCE THIS RUN EXACTLY:",
        f"  1. Checkout: git checkout {metadata.git.commit_hash}",
        f"  2. Use config SHA: {metadata.config_hash}",
        f"  3. Use solver:     {metadata.solver.solver_type} dt={metadata.solver.dt_nominal}",
        f"  4. Use seed:       {metadata.random_seed.python_seed}",
        "",
        "="*70,
        "",
    ])
    
    return "\n".join(lines)