# DQN Refactoring Implementation Guide

## Current Status: READY FOR EXECUTION

### Files Analysis

| File | Current Lines | Target | Status |
|------|---------------|--------|--------|
| `trainers/dqn.rs` | 4,975 | <1,000 each | **BACKUP CREATED** |
| `hyperopt/adapters/dqn.rs` | 3,162 | <1,000 each | Pending |
| `trainers/tft.rs` | 2,915 | <1,000 each | Pending |
| `trainers/mamba2.rs` | 544 | OK (under 1K) | ✅ No action needed |

### Extraction Map for `dqn.rs`

#### Module 1: `dqn/config.rs` (800 lines)
**Source Lines**: 48-747

**Contents**:
```rust
// Constants and type aliases (lines 48-56)
const EPISODE_LENGTH: usize = 200;
type FeatureVector = [f64; 54];
type FeatureVector51 = [f64; 51];

// FeatureStatistics (lines 77-157)
pub struct FeatureStatistics { /* Welford's algorithm */ }

// DQNHyperparameters (lines 425-610)
pub struct DQNHyperparameters { /* 60+ fields */ }

// DQNHyperparameters impl (lines 615-747)
impl DQNHyperparameters {
    pub fn conservative() -> Self { /* ... */ }
}
```

**Required Imports**:
```rust
// From trainers/mod.rs
use crate::trainers::TargetUpdateMode;
```

#### Module 2: `dqn/agent_wrapper.rs` (260 lines)
**Source Lines**: 164-424

**Contents**:
```rust
// DQNAgentType enum (lines 164-169)
pub enum DQNAgentType { /* Standard | RegimeConditional */ }

// QValueStats (lines 182-194)
pub struct QValueStats { /* C51 adaptive bounds */ }

// DQNAgentType impl (lines 196-421)
impl DQNAgentType {
    // 20+ unified API methods
}
```

**Required Imports**:
```rust
use candle_core::{Device, Tensor};
use crate::dqn::{
    dqn::{WorkingDQN},
    regime_conditional::{RegimeConditionalDQN, RegimeMetrics, RegimeType},
    action_space::FactoredAction,
    Experience,
};
use crate::MLError;
```

#### Module 3: `dqn/training_monitor.rs` (265 lines)
**Source Lines**: 728-1012

**Contents**:
```rust
// TrainingMonitor struct (lines 728-746)
struct TrainingMonitor { /* validation fields */ }

// TrainingMonitor impl (lines 748-1012)
impl TrainingMonitor {
    fn new(epoch: usize) -> Self { /* ... */ }
    fn track_reward(&mut self, reward: f32) { /* ... */ }
    fn track_action(&mut self, action: &FactoredAction) { /* ... */ }
    fn validate_rewards(&mut self) -> Result<()> { /* ... */ }
    fn validate_action_diversity(&self) -> Result<()> { /* ... */ }
    fn validate_q_value_balance(&self) -> Result<()> { /* ... */ }
    fn validate_all(&mut self) -> Result<()> { /* ... */ }
    // ... more validation methods
}
```

**Required Imports**:
```rust
use anyhow::Result;
use tracing::{info, warn};
use crate::dqn::action_space::FactoredAction;
```

#### Module 4: `dqn/trainer_core.rs` (600 lines estimated)
**Source Lines**: 1013-1680 (approx)

**Contents**:
```rust
// DQNTrainer struct (lines 1013-1116)
pub struct DQNTrainer {
    agent: Arc<RwLock<DQNAgentType>>,
    hyperparams: DQNHyperparameters,
    device: Device,
    // ... 20+ fields
}

// Debug impl (lines 1118-1124)
impl std::fmt::Debug for DQNTrainer { /* ... */ }

// Constructor impl (lines 1126-1450)
impl DQNTrainer {
    pub fn new(hyperparams: DQNHyperparameters) -> Result<Self> { /* ... */ }
    pub fn new_with_debug(hyperparams: DQNHyperparameters, debug_logging: bool) -> Result<Self> { /* ... */ }
    pub fn with_feature_cache(mut self, cache_dir: PathBuf) -> Self { /* ... */ }
}
```

**Required Imports**: (See full file - 40+ imports needed)

#### Module 5: `dqn/training_loop.rs` (1500 lines estimated)
**Source Lines**: 1464-2980 (main train methods)

**Contents**:
```rust
impl DQNTrainer {
    pub async fn train<F>(&mut self, ...) -> Result<TrainingMetrics> {
        // Main training loop: lines 1464-2800
    }

    pub async fn train_from_parquet<F>(&mut self, ...) -> Result<TrainingMetrics> {
        // Parquet training loop: lines 2981-3400
    }

    // Helper methods:
    fn calculate_epoch_metrics(...) -> Result<...> { /* ... */ }
    fn calculate_adaptive_bounds(...) -> (f64, f64) { /* ... */ }
    fn check_early_stopping(...) -> Option<String> { /* ... */ }
}
```

#### Module 6: `dqn/data_loading.rs` (600 lines estimated)
**Source Lines**: 3482-4100 (OHLCV extraction, feature creation)

**Contents**:
```rust
impl DQNTrainer {
    pub fn extract_ohlcv_bars_from_dbn(&self, file_path: &Path) -> Result<Vec<OHLCVBar>> {
        // DBN file parsing: lines 3482-3600
    }

    fn create_features(&self, ...) -> Result<Vec<f32>> {
        // Feature engineering: lines 3609-4100
    }
}
```

#### Module 7: `dqn/checkpointing.rs` (200 lines estimated)
**Source Lines**: Scattered throughout, needs extraction

**Contents**:
```rust
impl DQNTrainer {
    fn save_checkpoint(&self, epoch: usize) -> Result<()> { /* ... */ }
    fn load_checkpoint(&mut self, path: &str) -> Result<()> { /* ... */ }
}
```

#### Module 8: `dqn/mod.rs` (50 lines)
**New file**

**Contents**:
```rust
//! DQN Trainer Module
//!
//! Modularized DQN trainer for maintainability and testability.
//! This module preserves the original public API via re-exports.

// Internal modules
mod config;
mod agent_wrapper;
mod training_monitor;
mod trainer_core;
mod training_loop;
mod data_loading;
mod checkpointing;

// Public re-exports (maintain API compatibility)
pub use config::{
    DQNHyperparameters,
    FeatureStatistics,
    FeatureVector,
    FeatureVector51,
    EPISODE_LENGTH,
};

pub use agent_wrapper::{DQNAgentType, QValueStats};

pub use trainer_core::DQNTrainer;

// TrainingMonitor is internal, not part of public API
pub(crate) use training_monitor::TrainingMonitor;
```

## Implementation Scripts

### Script 1: Create `dqn/config.rs`

```bash
#!/bin/bash
# Extract config module from dqn.rs

SOURCE="ml/src/trainers/dqn.rs.backup"
TARGET="ml/src/trainers/dqn/config.rs"

cat > "$TARGET" <<'HEADER'
//! DQN Configuration and Hyperparameters
//!
//! Contains all configuration types for DQN training including:
//! - Constants and type aliases
//! - Feature normalization statistics (Welford's algorithm)
//! - DQN hyperparameters (60+ tunable parameters)

use crate::trainers::TargetUpdateMode;

HEADER

# Extract lines 48-747 (constants, FeatureStatistics, DQNHyperparameters)
sed -n '48,747p' "$SOURCE" >> "$TARGET"

echo "✅ Created $TARGET"
```

### Script 2: Create `dqn/agent_wrapper.rs`

```bash
#!/bin/bash
# Extract agent wrapper module from dqn.rs

SOURCE="ml/src/trainers/dqn.rs.backup"
TARGET="ml/src/trainers/dqn/agent_wrapper.rs"

cat > "$TARGET" <<'HEADER'
//! DQN Agent Type Wrapper
//!
//! Provides unified API for both standard and regime-conditional DQN agents.
//! Allows transparent switching between single-head and multi-head Q-networks.

use candle_core::{Device, Tensor};
use candle_nn::VarMap;

use crate::dqn::action_space::FactoredAction;
use crate::dqn::dqn::WorkingDQN;
use crate::dqn::regime_conditional::{RegimeConditionalDQN, RegimeMetrics, RegimeType};
use crate::dqn::Experience;
use crate::MLError;

HEADER

# Extract lines 164-424 (DQNAgentType, QValueStats)
sed -n '164,424p' "$SOURCE" >> "$TARGET"

echo "✅ Created $TARGET"
```

### Script 3: Create `dqn/training_monitor.rs`

```bash
#!/bin/bash
# Extract training monitor module from dqn.rs

SOURCE="ml/src/trainers/dqn.rs.backup"
TARGET="ml/src/trainers/dqn/training_monitor.rs"

cat > "$TARGET" <<'HEADER'
//! Training Monitoring and Validation
//!
//! Validates training progress to prevent common bugs:
//! - Constant rewards (reward shaping issues)
//! - Action collapse (policy degeneration)
//! - Q-value imbalance (exploration failure)

use anyhow::{Context, Result};
use tracing::{info, warn};

use crate::dqn::action_space::FactoredAction;

HEADER

# Extract lines 728-1012 (TrainingMonitor)
sed -n '728,1012p' "$SOURCE" >> "$TARGET"

echo "✅ Created $TARGET"
```

## Execution Checklist

- [x] Backup created: `dqn.rs.backup`
- [x] Directories created: `dqn/`, `tft/`, `mamba2/`, `shared/`
- [x] ADR written: `docs/ADR-001-dqn-refactoring.md`
- [ ] Extract `dqn/config.rs`
- [ ] Extract `dqn/agent_wrapper.rs`
- [ ] Extract `dqn/training_monitor.rs`
- [ ] Verify: `cargo check --package ml`
- [ ] Extract `dqn/trainer_core.rs`
- [ ] Extract `dqn/training_loop.rs`
- [ ] Extract `dqn/data_loading.rs`
- [ ] Extract `dqn/checkpointing.rs`
- [ ] Create `dqn/mod.rs` with re-exports
- [ ] Delete original `dqn.rs`
- [ ] Verify: `cargo check --package ml`
- [ ] Run tests: `cargo test --package ml --lib`
- [ ] Verify all 19 DQN tests pass
- [ ] Verify hyperopt adapter compiles

## Notes for Next Context

Due to file size (4,975 lines) and context limits, this refactoring requires:

1. **Automated extraction scripts** (provided above)
2. **Incremental verification** with `cargo check` after each module
3. **Careful import management** (40+ imports needed for trainer_core)
4. **Test validation** after completion

The architecture is sound. Implementation is mechanical but requires attention to:
- Import statements (many cross-dependencies)
- Method visibility (pub vs pub(crate))
- Re-exports maintaining public API

**Next agent should execute extraction scripts in order, verifying build after each step.**