#![allow(unsafe_code)] // Required for CUDA DtoD memcpy via cudarc.

//! Target Network Update Module
//!
//! Implements two strategies for updating target networks:
//! 1. **Polyak Averaging (Soft Updates)**: Gradual weight tracking via exponential moving average
//! 2. **Hard Updates**: Periodic full weight copy
//!
//! Rainbow DQN uses Polyak averaging with tau=0.001 for smoother Q-value stability.
//!
//! All weight updates run entirely on GPU -- zero host downloads.

use std::mem::ManuallyDrop;
use std::sync::Arc;

use cudarc::driver::{CudaStream, DevicePtr, DevicePtrMut};

use ml_core::cuda_autograd::GpuVarStore;
use ml_core::MLError;

/// Extract raw read-only device pointer from a `CudaSlice`.
fn raw_ptr(slice: &cudarc::driver::CudaSlice<f32>, stream: &CudaStream) -> u64 {
    let (ptr, guard) = slice.device_ptr(stream);
    let _no_drop = ManuallyDrop::new(guard);
    ptr
}

/// Extract raw mutable device pointer from a `CudaSlice`.
fn raw_ptr_mut(slice: &mut cudarc::driver::CudaSlice<f32>, stream: &CudaStream) -> u64 {
    let (ptr, guard) = slice.device_ptr_mut(stream);
    let _no_drop = ManuallyDrop::new(guard);
    ptr
}

/// Polyak averaging (soft target update) -- GPU-only.
///
/// Formula: `theta_target` = (1 - tau) * `theta_target` + tau * `theta_online`
///
/// Uses GPU elementwise kernels (affine + add) with zero host downloads.
///
/// # Arguments
/// * `online_vars` - `GpuVarStore` of the online Q-network
/// * `target_vars` - `GpuVarStore` of the target network
/// * `tau` - Interpolation coefficient (0.0 = no update, 1.0 = full copy)
/// * `stream` - CUDA stream for GPU operations
///
/// # Theory
/// Polyak averaging reduces Q-value oscillations by gradually tracking the online network.
/// Rainbow uses tau=0.001, giving a convergence half-life of ~693 steps.
///
/// **Benefits over Hard Updates**:
/// - 50-70% reduction in Q-value variance
/// - Smoother learning curves
/// - Better gradient stability
/// - No sudden target shifts
///
/// # Performance
/// Convergence half-life: `t_half` = ln(0.5) / ln(1 - tau)
/// - tau=0.001 -> 693 steps
/// - tau=0.01  -> 69 steps
/// - tau=0.1   -> 7 steps
pub fn polyak_update(
    online_vars: &GpuVarStore,
    target_vars: &mut GpuVarStore,
    tau: f64,
    stream: &Arc<CudaStream>,
) -> Result<(), MLError> {
    assert!(
        (0.0..=1.0).contains(&tau),
        "Tau must be in [0.0, 1.0], got {}",
        tau
    );

    let tau_f32 = tau as f32;
    let one_minus_tau = 1.0_f32 - tau_f32;

    // Pre-compile elementwise kernels (cached globally after first call)
    let ew = ml_core::cuda_autograd::elementwise::get_or_compile(stream)?;

    // Iterate over all online parameters and blend into target -- GPU-only
    for name in online_vars.param_names() {
        let online_param = online_vars.get(name).ok_or_else(|| {
            MLError::ModelError(format!("Online param '{name}' not found"))
        })?;
        let target_param = target_vars.get_mut(name).ok_or_else(|| {
            MLError::ModelError(format!("Target param '{name}' not found"))
        })?;

        let n = online_param.data.len();

        // GPU elementwise: online * tau
        let online_scaled = ew.affine(&online_param.data, tau_f32, 0.0, n)?;
        // GPU elementwise: target * (1 - tau)
        let target_scaled = ew.affine(&target_param.data, one_minus_tau, 0.0, n)?;
        // GPU elementwise: online_scaled + target_scaled
        let blended = ew.binary(&online_scaled, &target_scaled, n, 0)?;

        // DtoD copy result back into target param buffer
        let num_bytes = n * std::mem::size_of::<f32>();
        let src_ptr = raw_ptr(&blended, stream);
        let dst_ptr = raw_ptr_mut(&mut target_param.data, stream);
        // SAFETY: src and dst are valid device allocations of at least `num_bytes`.
        // Both belong to the same CUDA context. The operation is ordered on `stream`.
        unsafe {
            cudarc::driver::result::memcpy_dtod_async(
                dst_ptr, src_ptr, num_bytes, stream.cu_stream(),
            ).map_err(|e| MLError::ModelError(format!("polyak dtod '{name}': {e}")))?;
        }
    }

    Ok(())
}

/// Hard update (copy all weights) -- GPU-only via `DtoD` memcpy.
///
/// Used for:
/// 1. Initial target network setup
/// 2. Legacy hard update strategy (every N steps)
pub fn hard_update(
    online_vars: &GpuVarStore,
    target_vars: &mut GpuVarStore,
    stream: &Arc<CudaStream>,
) -> Result<(), MLError> {
    for name in online_vars.param_names() {
        let online_param = online_vars.get(name).ok_or_else(|| {
            MLError::ModelError(format!("Online param '{name}' not found"))
        })?;
        let target_param = target_vars.get_mut(name).ok_or_else(|| {
            MLError::ModelError(format!("Target param '{name}' not found"))
        })?;

        let n = online_param.data.len();
        let num_bytes = n * std::mem::size_of::<f32>();

        let src_ptr = raw_ptr(&online_param.data, stream);
        let dst_ptr = raw_ptr_mut(&mut target_param.data, stream);
        // SAFETY: src and dst are valid device allocations of at least `num_bytes`.
        // Both belong to the same CUDA context. The operation is ordered on `stream`.
        unsafe {
            cudarc::driver::result::memcpy_dtod_async(
                dst_ptr, src_ptr, num_bytes, stream.cu_stream(),
            ).map_err(|e| MLError::ModelError(format!("hard_update dtod '{name}': {e}")))?;
        }
    }

    Ok(())
}

/// Calculate convergence half-life for a given tau
///
/// Formula: `t_half` = ln(0.5) / ln(1 - tau)
///
/// Returns the number of steps for the target network to reach
/// 50% of the distance to the online network.
pub fn convergence_half_life(tau: f64) -> f64 {
    assert!(
        tau > 0.0 && tau < 1.0,
        "Tau must be in (0.0, 1.0), got {}",
        tau
    );
    (-0.5_f64.ln()) / (-(1.0 - tau).ln())
}

/// Compute network divergence (L2 norm of parameter differences) -- GPU-only.
///
/// Measures how far the target network has drifted from the online network.
/// Uses GPU subtract -> square -> sum reduction per param, with only single
/// scalar readbacks (one f32 per param) for the final summation.
pub fn compute_network_divergence(
    online_vars: &GpuVarStore,
    target_vars: &GpuVarStore,
    stream: &Arc<CudaStream>,
) -> Result<f64, MLError> {
    let ew = ml_core::cuda_autograd::elementwise::get_or_compile(stream)?;
    let reductions = ml_core::cuda_autograd::ReductionKernels::new(stream)?;

    let mut total_divergence = 0.0;
    let mut param_count = 0;

    for name in online_vars.param_names() {
        let Some(online_param) = online_vars.get(name) else { continue; };
        let Some(target_param) = target_vars.get(name) else { continue; };

        let n = online_param.data.len();
        if n == 0 { continue; }

        // GPU: diff = online - target
        let diff = ew.binary(&online_param.data, &target_param.data, n, 1)?; // op=1 is sub
        // GPU: diff_sq = diff * diff
        let diff_sq = ew.binary(&diff, &diff, n, 2)?; // op=2 is mul
        // GPU: sum_sq = sum(diff_sq) -- single scalar readback
        // ALLOWED: single scalar readback per param (exits system for divergence metric)
        let sum_sq = reductions.sum(&diff_sq, n)?;
        total_divergence += (sum_sq as f64).sqrt();
        param_count += 1;
    }

    if param_count > 0 {
        Ok(total_divergence / param_count as f64)
    } else {
        Ok(0.0)
    }
}

#[cfg(test)]
#[allow(clippy::let_underscore_must_use)]
mod tests {
    use super::*;
    use tracing::info;

    fn make_stream() -> Arc<CudaStream> {
        let device = ml_core::device::MlDevice::cuda(0).expect("CUDA required");
        device.cuda_stream().expect("stream").clone()
    }

    fn create_test_varstore(value: f32, stream: &Arc<CudaStream>) -> GpuVarStore {
        let mut store = GpuVarStore::new(stream.clone());

        // Create test tensors with uniform value
        let weight_host = vec![value; 10 * 10];
        let bias_host = vec![value; 10];

        let mut w_data = stream.alloc_zeros::<f32>(100).unwrap();
        stream.memcpy_htod(&weight_host, &mut w_data).unwrap();
        store.register("layer1.weight", w_data, vec![10, 10]).unwrap();

        let mut b_data = stream.alloc_zeros::<f32>(10).unwrap();
        stream.memcpy_htod(&bias_host, &mut b_data).unwrap();
        store.register("layer1.bias", b_data, vec![10]).unwrap();

        store
    }

    fn get_average_value(store: &GpuVarStore, stream: &Arc<CudaStream>) -> f32 {
        let mut sum = 0.0_f32;
        let mut count = 0;

        for name in store.param_names() {
            let param = store.get(name).expect("param must exist");
            let n = param.data.len();
            let mut buf = vec![0.0_f32; n];
            stream.memcpy_dtoh(&param.data, &mut buf).unwrap(); // test-only readback
            let param_sum: f32 = buf.iter().sum();
            sum += param_sum / n as f32;
            count += 1;
        }

        sum / count as f32
    }

    #[test]
    fn test_polyak_single_update() {
        let stream = make_stream();
        let online_vars = create_test_varstore(1.0, &stream);
        let mut target_vars = create_test_varstore(0.0, &stream);

        polyak_update(&online_vars, &mut target_vars, 0.1, &stream).unwrap();

        let avg = get_average_value(&target_vars, &stream);
        assert!(
            (avg - 0.1).abs() < 0.01,
            "Expected target ~=0.1, got {}",
            avg
        );
        info!(target = %format!("{:.3}", avg), "Single Polyak update (expected 0.1)");
    }

    #[test]
    fn test_hard_update_correctness() {
        let stream = make_stream();
        let online_vars = create_test_varstore(1.0, &stream);
        let mut target_vars = create_test_varstore(0.0, &stream);

        hard_update(&online_vars, &mut target_vars, &stream).unwrap();

        let avg = get_average_value(&target_vars, &stream);
        assert!((avg - 1.0).abs() < 1e-6, "Expected 1.0, got {}", avg);
        info!(target = %format!("{:.3}", avg), "Hard update (expected 1.0)");
    }

    #[test]
    fn test_convergence_half_life_calculation() {
        let half_life = convergence_half_life(0.001);
        assert!(
            (half_life - 693.0).abs() < 1.0,
            "Expected ~=693, got {}",
            half_life
        );
        info!(half_life = %format!("{:.0}", half_life), "Rainbow tau=0.001 half-life (steps)");

        let half_life_fast = convergence_half_life(0.01);
        assert!(
            (half_life_fast - 69.0).abs() < 1.0,
            "Expected ~=69, got {}",
            half_life_fast
        );
        info!(half_life = %format!("{:.0}", half_life_fast), "Fast tau=0.01 half-life (steps)");
    }

    #[test]
    fn test_gradual_convergence() {
        let stream = make_stream();
        let online_vars = create_test_varstore(1.0, &stream);
        let mut target_vars = create_test_varstore(0.0, &stream);

        let mut weights = vec![];
        for _ in 0..100 {
            polyak_update(&online_vars, &mut target_vars, 0.01, &stream).unwrap();
            weights.push(get_average_value(&target_vars, &stream));
        }

        for i in 1..weights.len() {
            assert!(
                weights[i] >= weights[i - 1] - 1e-6,
                "Non-monotonic at step {}: {:.4} -> {:.4}",
                i,
                weights[i - 1],
                weights[i]
            );
        }

        let final_weight = weights[99];
        assert!(
            final_weight > 0.6 && final_weight < 1.0,
            "Final weight should be 0.6-1.0, got {}",
            final_weight
        );
        info!(
            weight_first = %format!("{:.4}", weights[0]),
            weight_last = %format!("{:.4}", final_weight),
            "Gradual convergence"
        );
    }

    #[test]
    #[should_panic(expected = "Tau must be in [0.0, 1.0]")]
    fn test_invalid_tau_negative() {
        let stream = make_stream();
        let online_vars = create_test_varstore(1.0, &stream);
        let mut target_vars = create_test_varstore(0.0, &stream);
        let _ = polyak_update(&online_vars, &mut target_vars, -0.1, &stream);
    }

    #[test]
    #[should_panic(expected = "Tau must be in [0.0, 1.0]")]
    fn test_invalid_tau_too_large() {
        let stream = make_stream();
        let online_vars = create_test_varstore(1.0, &stream);
        let mut target_vars = create_test_varstore(0.0, &stream);
        let _ = polyak_update(&online_vars, &mut target_vars, 1.5, &stream);
    }
}
