//! Global-norm gradient clipping. The norm is computed across *all* parameter
//! gradients jointly (the same as `torch.nn.utils.clip_grad_norm_`): if the total
//! L2 norm exceeds `max_norm`, every gradient is scaled by `max_norm / total`.
//!
//! The norm math is host-only and testable; [`clip_grad_norm`] wraps it over the
//! parameter `Var`s (GPU round-trip), gated behind `not(no_cuda)`.

/// Compute the global L2 norm over a set of flat gradient buffers.
pub fn global_l2_norm(grads: &[Vec<f32>]) -> f32 {
    let mut sumsq = 0.0f64;
    for g in grads {
        for &x in g {
            sumsq += (x as f64) * (x as f64);
        }
    }
    sumsq.sqrt() as f32
}

/// The scale factor to apply for clipping to `max_norm` (1.0 if already under).
pub fn clip_scale(total_norm: f32, max_norm: f32) -> f32 {
    if total_norm > max_norm && total_norm > 0.0 {
        max_norm / total_norm
    } else {
        1.0
    }
}

#[cfg(not(no_cuda))]
mod gpu {
    use super::{clip_scale, global_l2_norm};
    use xtrain_autodiff::tape::Var;
    use xtrain_tensor::{Device, Tensor};

    /// First multiply every parameter's `.grad()` by `pre_scale` (use `1/batch`
    /// to turn accumulated summed grads into a batch mean; `1.0` for no-op), then
    /// clip the result to a joint global L2 norm of `max_norm`, writing the final
    /// grads back via the tape's grad slot. Returns the post-pre_scale total norm
    /// (the value the clip threshold is compared against, handy for logging).
    /// Parameters without a grad contribute 0.
    pub fn clip_grad_norm(params: &[Var], max_norm: f32, pre_scale: f32) -> f32 {
        let device = params[0].value().device();
        let grads: Vec<Option<Vec<f32>>> = params
            .iter()
            .map(|p| {
                p.grad()
                    .map(|g| g.to_device(Device::Cpu).as_slice::<f32>().to_vec())
            })
            .collect();

        // Norm is measured on the (pre_scale-applied) grads — what the optimizer
        // will actually see if no clipping is needed.
        let scaled_present: Vec<Vec<f32>> = grads
            .iter()
            .flatten()
            .map(|g| g.iter().map(|x| x * pre_scale).collect())
            .collect();
        let total = global_l2_norm(&scaled_present);
        let factor = pre_scale * clip_scale(total, max_norm);
        if (factor - 1.0).abs() < f32::EPSILON {
            return total; // pre_scale==1 and under threshold → grads untouched
        }

        for (p, g) in params.iter().zip(&grads) {
            if let Some(g) = g {
                let shape = p.grad().unwrap().shape().to_vec();
                let scaled: Vec<f32> = g.iter().map(|x| x * factor).collect();
                let t = Tensor::from_slice(&scaled, &shape).to_device(device);
                p.zero_grad();
                Var::push_grad(p, t);
            }
        }
        total
    }
}

#[cfg(not(no_cuda))]
mod gpu_norm {
    use super::clip_scale;
    use xtrain_autodiff::tape::Var;
    use xtrain_tensor::{DType, Device, Tensor};

    /// GPU-side global-norm grad clip (Phase T7): compute the joint L2 norm of all
    /// `pre_scale`-applied grads with a device reduction, then rescale every grad
    /// in place by `pre_scale·clip_factor` — no per-step grad roundtrip to host
    /// (only the single scalar norm comes back). Returns the post-pre_scale total
    /// norm. Params without a grad contribute 0 and are skipped on rescale.
    pub fn clip_grad_norm_gpu(params: &[Var], max_norm: f32, pre_scale: f32) -> f32 {
        let device = params[0].value().device();
        // sum-of-squares of the RAW grads accumulated on device.
        let acc = Tensor::zeros(&[1], DType::F32, device);
        for p in params {
            if let Some(g) = p.grad() {
                unsafe {
                    xtrain_cuda::ffi::launch_sumsq_accum_f32(
                        g.data_ptr() as *const f32,
                        acc.data_ptr() as *mut f32,
                        g.numel() as i32,
                        std::ptr::null_mut(),
                    );
                }
            }
        }
        xtrain_cuda::device::synchronize().expect("grad-norm reduce sync failed");
        let raw_sumsq = acc.to_device(Device::Cpu).as_slice::<f32>()[0];
        // Norm of the pre_scale-applied grads = pre_scale · sqrt(raw_sumsq).
        let total = pre_scale * raw_sumsq.max(0.0).sqrt();
        let factor = pre_scale * clip_scale(total, max_norm);
        if (factor - 1.0).abs() >= f32::EPSILON {
            for p in params {
                if let Some(g) = p.grad() {
                    unsafe {
                        xtrain_cuda::ffi::launch_scale_inplace_f32(
                            g.data_ptr() as *mut f32,
                            factor,
                            g.numel() as i32,
                            std::ptr::null_mut(),
                        );
                    }
                }
            }
            xtrain_cuda::device::synchronize().expect("grad rescale sync failed");
        }
        total
    }
}

#[cfg(not(no_cuda))]
pub use gpu::clip_grad_norm;
#[cfg(not(no_cuda))]
pub use gpu_norm::clip_grad_norm_gpu;

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn norm_and_scale() {
        // grads = [3,4] → norm 5.
        let g = vec![vec![3.0f32], vec![4.0]];
        assert!((global_l2_norm(&g) - 5.0).abs() < 1e-6);
        // Clip to 2.5 → scale 0.5.
        assert!((clip_scale(5.0, 2.5) - 0.5).abs() < 1e-6);
        // Already under → no scaling.
        assert!((clip_scale(5.0, 10.0) - 1.0).abs() < 1e-6);
    }
}
