#!/usr/bin/env python3
"""Add priority staleness tracking to PER buffer"""

import re

def add_staleness_tracking(filepath):
    with open(filepath, 'r') as f:
        content = f.read()

    # 1. Add field to struct
    struct_pattern = r'(pub struct PrioritizedReplayBuffer \{[^}]+priorities: Arc<Mutex<SegmentTree>>,)'
    struct_replacement = r'\1\n    priority_update_steps: Arc<RwLock<Vec<u64>>>,'
    content = re.sub(struct_pattern, struct_replacement, content)

    # 2. Initialize in new()
    new_pattern = r'(priorities: Arc::new\(Mutex::new\(SegmentTree::new\(config\.capacity\)\)\),)'
    new_replacement = r'\1\n            priority_update_steps: Arc::new(RwLock::new(vec![0; config.capacity])),'
    content = re.sub(new_pattern, new_replacement, content)

    # 3. Update push() - add current_step variable
    push_pattern1 = r'(let index = self\.position\.fetch_add\(1, Ordering::AcqRel\) % self\.config\.capacity;)'
    push_replacement1 = r'\1\n        let current_step = self.training_step.load(Ordering::Acquire) as u64;'
    content = re.sub(push_pattern1, push_replacement1, content)

    # 3b. Update push() - initialize priority update step
    push_pattern2 = r'(\{[\s\n]+let mut tree = self\.priorities\.lock\(\);[\s\n]+tree\.update\(index, priority\)\?;[\s\n]+\})'
    push_replacement2 = r'''\1

        // Initialize priority update step
        {
            let mut update_steps = self.priority_update_steps.write();
            if index < update_steps.len() {
                update_steps[index] = current_step;
            }
        }'''
    content = re.sub(push_pattern2, push_replacement2, content)

    # 4. Add staleness decay to sample() at the beginning
    sample_pattern = r'(pub fn sample\(\s+&self,\s+batch_size: usize,\s+\) -> Result<\(Vec<Experience>, Vec<f32>, Vec<usize>\), MLError> \{)'
    sample_replacement = r'''\1
        // Apply staleness decay before sampling
        // Default: max_age = 10000 steps, decay_factor = 0.9
        self.apply_staleness_decay(
            self.training_step.load(Ordering::Acquire) as u64,
            10000,
            0.9,
        )?;
'''
    content = re.sub(sample_pattern, sample_replacement, content)

    # 5. Update update_priorities() - add current_step
    update_pattern1 = r'(let mut max_priority = f32::from_bits\(self\.max_priority\.load\(Ordering::Acquire\) as u32\);)'
    update_replacement1 = r'\1\n        let current_step = self.training_step.load(Ordering::Acquire) as u64;'
    content = re.sub(update_pattern1, update_replacement1, content)

    # 5b. Update update_priorities() - track update steps
    update_pattern2 = r'(self\.max_priority[\s\n]+\.store\(max_priority\.to_bits\(\) as u64, Ordering::Release\);)'
    update_replacement2 = r'''\1

        // Update priority update steps
        {
            let mut update_steps = self.priority_update_steps.write();
            for &idx in indices {
                if idx < update_steps.len() {
                    update_steps[idx] = current_step;
                }
            }
        }'''
    content = re.sub(update_pattern2, update_replacement2, content)

    # 6. Add apply_staleness_decay method before clear()
    clear_pattern = r'(    /// Reset buffer \(clear all experiences\)[\s\n]+    pub fn clear\(&self\) \{)'
    staleness_method = r'''    /// Apply staleness decay to old priorities
    pub fn apply_staleness_decay(
        &self,
        current_step: u64,
        max_age: u64,
        decay_factor: f64,
    ) -> Result<(), MLError> {
        let size = self.size.load(Ordering::Acquire);
        if size == 0 {
            return Ok(());
        }

        let update_steps = self.priority_update_steps.read();
        let mut tree = self.priorities.lock();

        for idx in 0..size {
            if idx >= update_steps.len() {
                break;
            }

            let last_update = update_steps[idx];
            let age = current_step.saturating_sub(last_update);

            if age > max_age {
                let current_priority = tree.get_priority(idx);
                if current_priority > 0.0 {
                    let decay = decay_factor.powi((age / max_age) as i32) as f32;
                    let new_priority = (current_priority * decay).max(self.config.min_priority);
                    tree.update(idx, new_priority)?;
                }
            }
        }

        Ok(())
    }

    \1'''
    content = re.sub(clear_pattern, staleness_method, content)

    # 7. Update clear() to reset update steps
    clear_update_pattern = r'(\{[\s\n]+let mut tree = self\.priorities\.lock\(\);[\s\n]+for i in 0\.\.self\.config\.capacity \{[\s\n]+let _ = tree\.update\(i, 0\.0\);[\s\n]+\}[\s\n]+\})'
    clear_update_replacement = r'''\1

        {
            let mut update_steps = self.priority_update_steps.write();
            for step in update_steps.iter_mut() {
                *step = 0;
            }
        }'''
    content = re.sub(clear_update_pattern, clear_update_replacement, content)

    # 8. Add tests before the closing brace of mod tests
    test_pattern = r'(    \}\n\}\n)$'
    tests_addition = r'''
    #[test]
    fn test_priority_staleness_decay() {
        let config = PrioritizedReplayConfig {
            capacity: 100,
            initial_priority: 1.0,
            min_priority: 1e-6,
            ..Default::default()
        };
        let buffer = PrioritizedReplayBuffer::new(config)
            .expect("Failed to create prioritized replay buffer in test");

        // Push 10 experiences at step 0
        buffer.set_training_step(0);
        for _ in 0..10 {
            buffer
                .push(create_test_experience())
                .expect("Failed to push experience in test");
        }

        // Get initial priorities
        let tree = buffer.priorities.lock();
        let initial_priorities: Vec<f32> = (0..10).map(|i| tree.get_priority(i)).collect();
        drop(tree);

        // Simulate 15000 steps passing without priority updates
        buffer.set_training_step(15000);

        // Apply staleness decay (max_age=10000, decay_factor=0.9)
        buffer
            .apply_staleness_decay(15000, 10000, 0.9)
            .expect("Failed to apply staleness decay in test");

        // Check that priorities have decayed
        let tree = buffer.priorities.lock();
        for i in 0..10 {
            let current_priority = tree.get_priority(i);
            let initial_priority = initial_priorities[i];

            // Age = 15000 steps, max_age = 10000
            // Decay should be: decay_factor^(age/max_age) = 0.9^1
            let expected_decay = 0.9_f64.powi((15000 / 10000) as i32) as f32;
            let expected_priority = (initial_priority * expected_decay).max(1e-6);

            // Allow small floating point tolerance
            assert!(
                (current_priority - expected_priority).abs() < 0.01,
                "Priority at index {} should have decayed from {} to ~{}, got {}",
                i,
                initial_priority,
                expected_priority,
                current_priority
            );

            // Verify priority was reduced
            assert!(
                current_priority < initial_priority,
                "Priority at index {} should be reduced after staleness decay",
                i
            );
        }
        drop(tree);

        // Test that recently updated priorities don't decay
        buffer.set_training_step(16000);

        // Update some priorities
        let indices = vec![0, 1, 2];
        let new_priorities = vec![5.0, 6.0, 7.0];
        buffer
            .update_priorities(&indices, &new_priorities)
            .expect("Failed to update priorities in test");

        // Move forward only 5000 steps (below max_age threshold)
        buffer.set_training_step(21000);
        buffer
            .apply_staleness_decay(21000, 10000, 0.9)
            .expect("Failed to apply staleness decay in test");

        // Recently updated priorities should not decay significantly
        let tree = buffer.priorities.lock();
        for &idx in &indices {
            let priority = tree.get_priority(idx);
            // Age = 5000, which is < max_age (10000), so no decay should occur
            assert!(
                priority > 4.5,
                "Recently updated priority at index {} should not decay significantly, got {}",
                idx,
                priority
            );
        }
    }

    #[test]
    fn test_staleness_tracking_on_push() {
        let config = PrioritizedReplayConfig {
            capacity: 100,
            ..Default::default()
        };
        let buffer = PrioritizedReplayBuffer::new(config)
            .expect("Failed to create prioritized replay buffer in test");

        // Push experiences at different training steps
        buffer.set_training_step(100);
        buffer
            .push(create_test_experience())
            .expect("Failed to push experience in test");

        buffer.set_training_step(200);
        buffer
            .push(create_test_experience())
            .expect("Failed to push experience in test");

        // Verify update steps were recorded
        let update_steps = buffer.priority_update_steps.read();
        assert_eq!(update_steps[0], 100, "First experience should have update step 100");
        assert_eq!(update_steps[1], 200, "Second experience should have update step 200");
    }

    #[test]
    fn test_staleness_tracking_on_update() {
        let config = PrioritizedReplayConfig {
            capacity: 100,
            ..Default::default()
        };
        let buffer = PrioritizedReplayBuffer::new(config)
            .expect("Failed to create prioritized replay buffer in test");

        // Push experiences
        buffer.set_training_step(100);
        for _ in 0..10 {
            buffer
                .push(create_test_experience())
                .expect("Failed to push experience in test");
        }

        // Update priorities at a later step
        buffer.set_training_step(500);
        let indices = vec![0, 1, 2];
        let priorities = vec![2.0, 3.0, 4.0];
        buffer
            .update_priorities(&indices, &priorities)
            .expect("Failed to update priorities in test");

        // Verify update steps were updated
        let update_steps = buffer.priority_update_steps.read();
        assert_eq!(update_steps[0], 500, "Updated experience should have new update step");
        assert_eq!(update_steps[1], 500, "Updated experience should have new update step");
        assert_eq!(update_steps[2], 500, "Updated experience should have new update step");
        assert_eq!(update_steps[3], 100, "Non-updated experience should retain old update step");
    }
\1'''
    content = re.sub(test_pattern, tests_addition, content)

    with open(filepath, 'w') as f:
        f.write(content)

    print(f"Successfully updated {filepath}")

if __name__ == "__main__":
    add_staleness_tracking("/home/jgrusewski/Work/foxhunt/ml/src/dqn/prioritized_replay.rs")