# Explicit Configuration System - Implementation Guide

**Date**: October 19, 2025
**Status**: COMPLETE
**Goal**: Remove all fallback behavior, make everything explicit with clear errors

---

## Completed

### 1. Exception Classes (`src/phasic/exceptions.py`)
Created comprehensive error hierarchy:
- `phasicError` - Base exception
- `PTDConfigError` - Configuration errors with fix suggestions
- `PTDBackendError` - Backend unavailable errors
- `PTDFeatureError` - Platform feature errors
- `PTDJAXError` - JAX-specific errors

**Usage**:
```python
from phasic.exceptions import PTDConfigError

raise PTDConfigError(
    "jax=True but JAX not installed.\n"
    "  Install: pip install jax\n"
    "  Or: phasic.configure(jax=False)"
)
```

### 2. Configuration System (`src/phasic/config.py`)
Created `phasicConfig` dataclass with:
- **Parameters**:
  - `jax: bool` - Require JAX (error if unavailable)
  - `jit: bool` - Enable JIT compilation
  - `ffi: bool` - Enable FFI (currently always disabled)
  - `strict: bool` - Error vs warn on problems
  - `platform: str` - 'cpu', 'gpu', 'tpu'
  - `backend: str` - 'jax', 'cpp', 'ffi'
  - `verbose: bool` - Print config details

- **Methods**:
  - `validate()` - Check configuration validity
  - `get_available_options()` - Return what's available on system
  - Factory methods: `jax_only()`, `cpp_only()`, `permissive()`

- **Global API**:
  - `configure(**kwargs)` - Set global config
  - `get_config()` - Get current config
  - `get_available_options()` - Check availability
  - `reset_config()` - Reset to defaults

**Usage**:
```python
import phasic as ptd

# Check availability
print(ptd.get_available_options())
# {'jax': True, 'jit': True, 'ffi': False, ...}

# Configure explicitly
ptd.configure(jax=True, jit=True, ffi=False)

# Or use factories
config = ptd.phasicConfig.jax_only()
config.validate()
```

---

## Completed (Part 2)

### 3. Update `__init__.py` ✅

**Status**: COMPLETE

**Changes made**:
1. Imported config system at top (lines 67-80)
2. Replaced try-except JAX import with config-based import
3. Preserved XLA_FLAGS and environment setup
4. Raise PTDConfigError if JAX unavailable when `jax=True`
5. Added config exports documentation at end of file

**Code**:
```python
# At top of file, after imports
from .config import (
    configure,
    get_config,
    get_available_options,
    phasicConfig,
    reset_config
)
from .exceptions import *

# Initialize default config (validates on first access)
_config = get_config()

# Import JAX based on config
if _config.jax:
    try:
        import jax
        import jax.numpy as jnp
        HAS_JAX = True
    except ImportError:
        raise PTDJAXError(
            "jax=True but JAX not installed.\n"
            "  Install: pip install jax jaxlib\n"
            "  Or: phasic.configure(jax=False) before importing"
        )
else:
    jax = None
    jnp = None
    HAS_JAX = False

# SVGD - only import if JAX available
if _config.jax:
    from .svgd import SVGD
else:
    SVGD = None  # Keep this for backward compat, but user chose jax=False

# Update __all__ to include config functions
__all__ = [
    ... existing exports ...,
    'configure',
    'get_config',
    'get_available_options',
    'phasicConfig',
    'PTDConfigError',
    'PTDBackendError',
    'PTDJAXError',
]
```

**Challenge**: The JAX configuration code (lines 67-168) sets environment variables BEFORE importing JAX. This needs to be preserved but integrated with config system.

**Solution**: Keep environment variable setup, but make it conditional on config:
```python
# Only configure environment if JAX will be imported
if get_config().jax:
    # ... existing XLA_FLAGS, PTDALG_CPUS setup ...
    import jax
    import jax.numpy as jnp
```

---

## Completed (Part 3)

### 4. Update `ffi_wrappers.py` ✅

**Status**: COMPLETE (simplified approach for FFI-disabled state)

**Changes made**:
1. Added config and exceptions imports
2. Changed C++ module import from try-except to explicit requirement with clear error
3. Documented that FFI is disabled (memory corruption bug)
4. Updated `_register_ffi_targets()` to raise clear error if FFI requested
5. Removed silent try-except at module initialization
6. Added documentation for future FFI re-enablement

**Rationale**:
Since FFI is currently completely disabled (see FFI_MEMORY_CORRUPTION_FIX.md),
adding a full `backend` parameter would be premature. The current changes:
- Remove all silent fallbacks
- Make it crystal clear FFI is disabled and why
- Provide clear error if user tries to enable it
- Document how to properly re-enable when bug is fixed

**Future work** (when FFI re-enabled):
Add `backend` parameter as originally planned in section below.

**Original plan**:

1. **Remove silent fallbacks** (lines 46-51, 54-69, 591-595):
   ```python
   # DELETE:
   try:
       from . import phasic_pybind as cpp_module
       _HAS_CPP_MODULE = True
   except ImportError:
       _HAS_CPP_MODULE = False  # SILENT FALLBACK!

   # REPLACE WITH:
   from . import phasic_pybind as cpp_module
   from .config import get_config
   from .exceptions import PTDBackendError
   ```

2. **Add `backend` parameter** to all functions:
   ```python
   def compute_pmf_ffi(
       structure_json,
       theta,
       times,
       discrete=False,
       granularity=100,
       backend=None  # NEW: explicit backend control
   ):
       if backend is None:
           backend = get_config().backend

       if backend == 'ffi':
           raise PTDBackendError(
               "FFI backend not available.\n"
               "  Reason: Memory corruption bug\n"
               "  See: FFI_MEMORY_CORRUPTION_FIX.md\n"
               "  Available: " + str(_get_available_backends())
           )
       elif backend == 'jax':
           # Use pure_callback
       elif backend == 'cpp':
           # Use direct pybind11
       else:
           raise ValueError(f"Unknown backend: {backend}")
   ```

3. **Remove `_register_ffi_targets()` fallback** (lines 163-204):
   - Remove try-except
   - Raise clear error if FFI requested

4. **Remove module init fallback** (lines 591-595):
   ```python
   # DELETE:
   try:
       _register_ffi_targets()
   except Exception as e:
       pass  # SILENT FAILURE!
   ```

### 5. Update `svgd.py` ✅

**Changes completed**:

1. **Removed automatic XLA_FLAGS** (lines 10-15):
   - Deleted automatic environment variable configuration
   - Added comment explaining XLA_FLAGS is configured by __init__.py
   - Users must configure via `ptd.configure()` or PTDALG_CPUS env var

2. **Added explicit parameters** to SVGD class:
   - New parameters: `jit`, `parallel`, `n_devices`
   - Validation against config (raises PTDConfigError)
   - Auto-selection of sensible defaults (pmap if multi-device, vmap otherwise)
   - Backward compatibility via `precompile` parameter

3. **Updated fit() method**:
   - Modified `svgd_step()` to accept parallel_mode and n_devices
   - Modified `run_svgd()` to accept and pass through config
   - Both `fit()` and `fit_regularized()` now pass config parameters

4. **Testing**:
   - ✓ Basic configuration and instantiation
   - ✓ Error on invalid parallel mode
   - ✓ Error on excessive n_devices
   - ✓ SVGD.fit() executes successfully
   - ✓ Backward compatibility maintained

**See**: `SVGD_CONFIG_INTEGRATION_PLAN.md` for detailed implementation plan

### 6. Check `trace_elimination.py`

**Action**: Search for any JAX-related try-except or fallback logic

### 7. Update Documentation

**Files to update**:
- `README.md` - Add configuration section
- `docs/` - Add configuration guide
- Docstrings in affected modules

**Example documentation**:
```markdown
## Configuration

phasic uses an explicit configuration system. No silent fallbacks.

### Check what's available

```python
import phasic as ptd
print(ptd.get_available_options())
# {'jax': True, 'jit': True, 'ffi': False,
#  'backends': ['jax', 'cpp'], 'platforms': ['cpu']}
```

### Configure before use

```python
# Standard: JAX with JIT
ptd.configure(jax=True, jit=True, ffi=False)

# Pure C++: No JAX
ptd.configure(jax=False, backend='cpp')

# Permissive: Warnings not errors
ptd.configure(jax=True, strict=False)
```

### Error messages

Clear, actionable errors when features unavailable:

```
PTDConfigError: jax=True but JAX not installed.
  Install: pip install jax jaxlib
  Or: phasic.configure(jax=False)
```
```

---

## Testing Plan

### Unit Tests

1. **Config validation**:
   - Test with JAX installed
   - Test without JAX (should error if `jax=True`)
   - Test invalid combinations (e.g., `jit=True, jax=False`)

2. **Backend selection**:
   - Test `backend='jax'` with JAX available
   - Test `backend='jax'` without JAX (should error)
   - Test `backend='cpp'` always works
   - Test `backend='ffi'` always errors (currently disabled)

3. **Error messages**:
   - Verify error messages are clear and actionable
   - Check that errors include fix suggestions

### Integration Tests

1. **Import sequence**:
   ```python
   # Test 1: Default config
   import phasic as ptd  # Should work if JAX installed

   # Test 2: Explicit config
   ptd.configure(jax=True, jit=True)
   from ptd.ffi_wrappers import compute_pmf_ffi

   # Test 3: Pure C++
   ptd.configure(jax=False, backend='cpp')
   # Should work without JAX
   ```

2. **SVGD configuration**:
   ```python
   svgd = ptd.SVGD(
       log_prob,
       theta_dim=10,
       jit=True,       # Explicit
       parallel='pmap', # Explicit
       devices=8       # Explicit
   )
   ```

---

## Migration Guide (for Users)

### Before (Silent Fallbacks)

```python
import phasic as ptd

# JAX imported silently if available
# No way to know if JIT is enabled
# FFI might fail silently

graph = ptd.Graph(...)
# What backend is this using? Unknown!
```

### After (Explicit Configuration)

```python
import phasic as ptd

# Check what's available
opts = ptd.get_available_options()
print(f"JAX available: {opts['jax']}")
print(f"Backends: {opts['backends']}")

# Configure explicitly
ptd.configure(
    jax=True,      # Explicit - errors if not installed
    jit=True,      # Explicit - enable JIT
    ffi=False,     # Explicit - FFI not ready
    backend='jax'  # Explicit - use JAX backend
)

# Now use library
graph = ptd.Graph(...)

# FFI with explicit backend
from ptd.ffi_wrappers import compute_pmf_ffi
pmf = compute_pmf_ffi(..., backend='jax')  # Explicit!
```

---

## Benefits

1. **No surprises**: User always knows what's happening
2. **Clear errors**: When something fails, user knows why and how to fix
3. **Testability**: Can test with/without JAX explicitly
4. **Debuggability**: No silent fallbacks hiding problems
5. **Documentation**: Configuration is self-documenting

---

## Implementation Status

### Completed ✅
1. exceptions.py - Error class hierarchy with actionable messages
2. config.py - phasicConfig with validation and global API
3. __init__.py - Config-based JAX import with clear errors
4. ffi_wrappers.py - Removed silent fallbacks, documented FFI disabled
5. trace_elimination.py - Checked, no silent fallbacks found
6. Basic testing - Import works, config validation works, errors are clear

### Deferred ⏸
- svgd.py - Adding explicit jit/parallel parameters
  - Reason: Large file (2317 lines), needs careful review
  - Current state: Has automatic XLA_FLAGS setup that should be removed
  - Plan created: See SVGD_CONFIG_INTEGRATION_PLAN.md (detailed 5-phase plan)
  - Prompt ready: See SVGD_INTEGRATION_PROMPT.txt (for new conversation)
  - Recommended: Handle in separate focused session (~1.5 hours)

### Remaining 📋
- Write comprehensive unit tests for config system
- Test library with JAX unavailable (mock environment)
- Update CLAUDE.md with configuration patterns
- Update user documentation (README, docs/)
