//! GPU variable store — replaces Candle's `VarMap`/`VarBuilder`.
//!
//! Stores named weight tensors as `CudaSlice<f32>` on GPU. Provides:
//! - Named parameter registration and lookup
//! - Linear layer builder (returns [`GpuLinear`])
//! - Flat parameter view for single-kernel optimizer updates
//! - Save/load via safetensors-compatible byte buffers

use std::collections::BTreeMap;
use std::sync::Arc;

use cudarc::driver::{CudaSlice, CudaStream};

use crate::MLError;
use super::gpu_tensor::GpuTensor;
use super::init;
use super::linear::GpuLinear;

/// A named GPU parameter (weight or bias).
pub struct GpuParam {
    /// The GPU buffer holding this parameter's data.
    pub data: CudaSlice<f32>,
    /// Shape of the parameter tensor.
    pub shape: Vec<usize>,
}

impl std::fmt::Debug for GpuParam {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "GpuParam {{ shape: {:?}, numel: {} }}", self.shape, self.data.len())
    }
}

/// GPU variable store replacing Candle's `VarMap`.
///
/// Stores all trainable parameters as named `CudaSlice<f32>` buffers.
/// Parameters are ordered by insertion order (via `BTreeMap` on name).
///
/// ## Flat parameter view
///
/// For optimizers that operate on a single contiguous buffer (like the fused
/// DQN Adam kernel), call [`flatten()`](GpuVarStore::flatten) to get all
/// parameters concatenated into one `CudaSlice<f32>`. After optimizer update,
/// call [`unflatten()`](GpuVarStore::unflatten) to scatter back.
pub struct GpuVarStore {
    /// Named parameters, ordered by name for deterministic iteration.
    params: BTreeMap<String, GpuParam>,
    /// The CUDA stream used for all allocations and transfers.
    stream: Arc<CudaStream>,
}

impl std::fmt::Debug for GpuVarStore {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let names: Vec<&str> = self.params.keys().map(|s| s.as_str()).collect();
        let total: usize = self.params.values().map(|p| p.data.len()).sum();
        write!(
            f,
            "GpuVarStore {{ params: {}, total_elements: {}, names: {:?} }}",
            self.params.len(),
            total,
            names
        )
    }
}

impl GpuVarStore {
    /// Create a new empty variable store on the given stream.
    pub fn new(stream: Arc<CudaStream>) -> Self {
        Self {
            params: BTreeMap::new(),
            stream,
        }
    }

    /// Reference to the CUDA stream.
    pub fn stream(&self) -> &Arc<CudaStream> {
        &self.stream
    }

    /// Alias for `stream()` — some callers use `cuda_stream()`.
    pub fn cuda_stream(&self) -> &Arc<CudaStream> {
        &self.stream
    }

    /// Read-only access to the params map.
    ///
    /// Callers that need to iterate over parameters directly can use this.
    pub fn data(&self) -> &BTreeMap<String, GpuParam> {
        &self.params
    }

    /// Create a var store from an `MlDevice`.
    ///
    /// Extracts the CUDA stream from the device. Returns an error if the
    /// device is CPU-only.
    #[cfg(feature = "cuda")]
    pub fn with_device(device: &crate::device::MlDevice) -> Result<Self, MLError> {
        let stream = device.cuda_stream()?;
        Ok(Self::new(Arc::clone(stream)))
    }

    /// Register a named parameter from an existing `CudaSlice`.
    ///
    /// # Errors
    ///
    /// Returns an error if a parameter with the same name already exists.
    pub fn register<S: Into<String>>(
        &mut self,
        name: S,
        data: CudaSlice<f32>,
        shape: Vec<usize>,
    ) -> Result<(), MLError> {
        let name = name.into();
        if self.params.contains_key(&name) {
            return Err(MLError::ConfigError(format!(
                "GpuVarStore: duplicate parameter name '{name}'"
            )));
        }
        self.params.insert(name, GpuParam { data, shape });
        Ok(())
    }

    /// Look up a parameter by name.
    pub fn get(&self, name: &str) -> Option<&GpuParam> {
        self.params.get(name)
    }

    /// Mutably look up a parameter by name.
    pub fn get_mut(&mut self, name: &str) -> Option<&mut GpuParam> {
        self.params.get_mut(name)
    }

    /// Iterator over (name, param) pairs in sorted order.
    pub fn iter(&self) -> impl Iterator<Item = (&str, &GpuParam)> {
        self.params.iter().map(|(k, v)| (k.as_str(), v))
    }

    /// Number of registered parameters.
    pub fn len(&self) -> usize {
        self.params.len()
    }

    /// Whether the store is empty.
    pub fn is_empty(&self) -> bool {
        self.params.is_empty()
    }

    /// Total number of scalar elements across all parameters.
    pub fn total_elements(&self) -> usize {
        self.params.values().map(|p| p.data.len()).sum()
    }

    /// Names of all parameters in sorted order.
    pub fn param_names(&self) -> Vec<&str> {
        self.params.keys().map(|s| s.as_str()).collect()
    }

    /// List all parameter names (alias for `param_names`).
    pub fn vars(&self) -> Vec<&str> {
        self.param_names()
    }

    /// Return all parameters as `GpuTensor` values (cloned to new allocations).
    #[cold]
    pub fn all_vars(&self) -> Result<Vec<GpuTensor>, MLError> {
        let mut out = Vec::with_capacity(self.params.len());
        for (name, param) in &self.params {
            let mut host = vec![0.0_f32; param.data.len()];
            self.stream.memcpy_dtoh(&param.data, &mut host).map_err(|e| { // gpu-exit: checkpoint export
                MLError::ModelError(format!("all_vars DtoH '{name}': {e}"))
            })?;
            let t = GpuTensor::from_host(&host, param.shape.clone(), &self.stream)?;
            out.push(t);
        }
        Ok(out)
    }

    // ── Layer builders ────────────────────────────────────────────────

    /// Create a linear layer with Xavier initialization and register its
    /// weights and bias in this store under `{prefix}.weight` / `{prefix}.bias`.
    ///
    /// Returns a [`GpuLinear`] that holds references to the registered buffers.
    pub fn linear(
        &mut self,
        prefix: &str,
        in_dim: usize,
        out_dim: usize,
    ) -> Result<GpuLinear, MLError> {
        let w_name = format!("{prefix}.weight");
        let b_name = format!("{prefix}.bias");

        let w_data = init::xavier_uniform(in_dim, out_dim, &self.stream)?;
        let b_data = init::bias_uniform(in_dim, out_dim, &self.stream)?;

        self.register(w_name.clone(), w_data, vec![out_dim, in_dim])?;
        self.register(b_name.clone(), b_data, vec![out_dim])?;

        Ok(GpuLinear {
            weight_name: w_name,
            bias_name: b_name,
            in_dim,
            out_dim,
        })
    }

    /// Alias for `linear()` — Xavier-initialized linear layer.
    ///
    /// Provided for call sites that want the initialization scheme explicit in the name.
    pub fn linear_xavier(
        &mut self,
        prefix: &str,
        in_dim: usize,
        out_dim: usize,
    ) -> Result<GpuLinear, MLError> {
        self.linear(prefix, in_dim, out_dim)
    }

    /// Create a linear layer with near-zero Xavier initialization.
    ///
    /// For distributional output heads where uniform softmax at init is desired.
    pub fn linear_near_zero(
        &mut self,
        prefix: &str,
        in_dim: usize,
        out_dim: usize,
    ) -> Result<GpuLinear, MLError> {
        let w_name = format!("{prefix}.weight");
        let b_name = format!("{prefix}.bias");

        let w_data = init::near_zero_xavier(in_dim, out_dim, &self.stream)?;
        let b_data = init::zeros(out_dim, &self.stream)?;

        self.register(w_name.clone(), w_data, vec![out_dim, in_dim])?;
        self.register(b_name.clone(), b_data, vec![out_dim])?;

        Ok(GpuLinear {
            weight_name: w_name,
            bias_name: b_name,
            in_dim,
            out_dim,
        })
    }

    /// Create a linear layer with Kaiming initialization (for ReLU networks).
    pub fn linear_kaiming(
        &mut self,
        prefix: &str,
        in_dim: usize,
        out_dim: usize,
    ) -> Result<GpuLinear, MLError> {
        let w_name = format!("{prefix}.weight");
        let b_name = format!("{prefix}.bias");

        let w_data = init::kaiming_uniform(in_dim, out_dim, &self.stream)?;
        let b_data = init::bias_uniform(in_dim, out_dim, &self.stream)?;

        self.register(w_name.clone(), w_data, vec![out_dim, in_dim])?;
        self.register(b_name.clone(), b_data, vec![out_dim])?;

        Ok(GpuLinear {
            weight_name: w_name,
            bias_name: b_name,
            in_dim,
            out_dim,
        })
    }

    // ── Flat parameter view ──────────────────────────────────────────

    /// Flatten all parameters into a single contiguous GPU buffer.
    ///
    /// Returns (flat_buffer, offsets) where offsets maps param_name → (offset, count).
    /// The flat buffer is a fresh allocation — no aliasing with individual params.
    #[cold]
    pub fn flatten(&self) -> Result<(GpuTensor, BTreeMap<String, (usize, usize)>), MLError> {
        let total = self.total_elements();
        if total == 0 {
            let buf = self.stream.alloc_zeros::<f32>(1).map_err(|e| {
                MLError::ModelError(format!("flatten alloc: {e}"))
            })?;
            return Ok((GpuTensor::from_flat(buf), BTreeMap::new()));
        }

        // Gather all params to host, concatenate, re-upload
        let mut host_flat = Vec::with_capacity(total);
        let mut offsets = BTreeMap::new();
        for (name, param) in &self.params {
            let offset = host_flat.len();
            let count = param.data.len();
            let mut host_param = vec![0.0_f32; count];
            self.stream.memcpy_dtoh(&param.data, &mut host_param).map_err(|e| { // gpu-exit: checkpoint export
                MLError::ModelError(format!("flatten DtoH '{name}': {e}"))
            })?;
            host_flat.extend_from_slice(&host_param);
            offsets.insert(name.clone(), (offset, count));
        }

        let tensor = GpuTensor::from_host(&host_flat, vec![total], &self.stream)?;
        Ok((tensor, offsets))
    }

    /// Scatter a flat parameter buffer back into individual parameter tensors.
    ///
    /// `offsets` must match what was returned by [`flatten()`](Self::flatten).
    pub fn unflatten(
        &mut self,
        flat: &GpuTensor,
        offsets: &BTreeMap<String, (usize, usize)>,
    ) -> Result<(), MLError> {
        let flat_host = flat.to_host(&self.stream)?;
        for (name, &(offset, count)) in offsets {
            if let Some(param) = self.params.get_mut(name) {
                let slice = &flat_host[offset..offset + count];
                self.stream.memcpy_htod(slice, &mut param.data).map_err(|e| {
                    MLError::ModelError(format!("unflatten HtoD '{name}': {e}"))
                })?;
            }
        }
        Ok(())
    }

    // ── Save / Load ──────────────────────────────────────────────────

    /// Export all parameters as a map of `(name -> (shape, f32_bytes))`.
    ///
    /// This produces the data needed for safetensors serialization.
    /// The caller can use the `safetensors` crate to write the file.
    #[cold]
    pub fn export_to_host(
        &self,
    ) -> Result<BTreeMap<String, (Vec<usize>, Vec<f32>)>, MLError> {
        let mut out = BTreeMap::new();
        for (name, param) in &self.params {
            let mut host = vec![0.0_f32; param.data.len()];
            self.stream.memcpy_dtoh(&param.data, &mut host).map_err(|e| { // gpu-exit: checkpoint export
                MLError::ModelError(format!("export DtoH '{name}': {e}"))
            })?;
            out.insert(name.clone(), (param.shape.clone(), host));
        }
        Ok(out)
    }

    /// Import parameters from host data (e.g., loaded from safetensors).
    ///
    /// Only imports parameters whose names match existing registered params.
    /// Shape must match exactly.
    pub fn import_from_host(
        &mut self,
        data: &BTreeMap<String, (Vec<usize>, Vec<f32>)>,
    ) -> Result<(), MLError> {
        for (name, (shape, values)) in data {
            if let Some(param) = self.params.get_mut(name) {
                let expected_numel: usize = param.shape.iter().product();
                let actual_numel = values.len();
                if expected_numel != actual_numel {
                    return Err(MLError::DimensionMismatch {
                        expected: expected_numel,
                        actual: actual_numel,
                    });
                }
                if &param.shape != shape {
                    return Err(MLError::ConfigError(format!(
                        "Shape mismatch for '{name}': expected {:?}, got {:?}",
                        param.shape, shape
                    )));
                }
                self.stream.memcpy_htod(values, &mut param.data).map_err(|e| {
                    MLError::ModelError(format!("import HtoD '{name}': {e}"))
                })?;
            }
        }
        Ok(())
    }
}

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

    fn make_stream() -> Arc<CudaStream> {
        cudarc::driver::CudaContext::new(0)
            .expect("CUDA required")
            .new_stream()
            .expect("fork stream")
    }

    #[test]
    fn test_register_and_get() {
        let stream = make_stream();
        let mut store = GpuVarStore::new(stream.clone());
        let data = init::zeros(12, &stream).unwrap();
        store.register("layer.weight", data, vec![3, 4]).unwrap();
        assert_eq!(store.len(), 1);
        assert!(store.get("layer.weight").is_some());
        assert!(store.get("nonexistent").is_none());
    }

    #[test]
    fn test_duplicate_name_rejected() {
        let stream = make_stream();
        let mut store = GpuVarStore::new(stream.clone());
        let d1 = init::zeros(4, &stream).unwrap();
        let d2 = init::zeros(4, &stream).unwrap();
        store.register("x", d1, vec![4]).unwrap();
        assert!(store.register("x", d2, vec![4]).is_err());
    }

    #[test]
    fn test_linear_builder() {
        let stream = make_stream();
        let mut store = GpuVarStore::new(stream);
        let _linear = store.linear("fc1", 64, 128).unwrap();
        assert_eq!(store.len(), 2);
        assert!(store.get("fc1.weight").is_some());
        assert!(store.get("fc1.bias").is_some());
        let w = store.get("fc1.weight").unwrap();
        assert_eq!(w.shape, vec![128, 64]);
        assert_eq!(w.data.len(), 128 * 64);
    }

    #[test]
    fn test_total_elements() {
        let stream = make_stream();
        let mut store = GpuVarStore::new(stream);
        store.linear("fc1", 32, 64).unwrap();
        store.linear("fc2", 64, 16).unwrap();
        // fc1: 32*64 + 64 = 2112, fc2: 64*16 + 16 = 1040
        assert_eq!(store.total_elements(), 2112 + 1040);
    }

    #[test]
    fn test_flatten_unflatten_roundtrip() {
        let stream = make_stream();
        let mut store = GpuVarStore::new(stream.clone());
        store.linear("fc1", 8, 16).unwrap();
        store.linear("fc2", 16, 4).unwrap();

        // Export original values
        let original = store.export_to_host().unwrap(); // test-only readback

        // Flatten
        let (flat, offsets) = store.flatten().unwrap();
        assert_eq!(flat.numel(), store.total_elements());

        // Zero out the store params
        for (_, param) in store.params.iter_mut() {
            stream.memset_zeros(&mut param.data).unwrap();
        }

        // Unflatten should restore
        store.unflatten(&flat, &offsets).unwrap();
        let restored = store.export_to_host().unwrap(); // test-only readback

        for (name, (shape, values)) in &original {
            let (r_shape, r_values) = restored.get(name).unwrap();
            assert_eq!(shape, r_shape);
            for (a, b) in values.iter().zip(r_values.iter()) {
                assert!((a - b).abs() < 1e-6, "mismatch in {name}");
            }
        }
    }

    #[test]
    fn test_export_import() {
        let stream = make_stream();
        let mut store = GpuVarStore::new(stream.clone());
        store.linear("fc1", 8, 4).unwrap();

        let exported = store.export_to_host().unwrap(); // test-only readback

        // Create a fresh store with same structure
        let mut store2 = GpuVarStore::new(stream);
        store2.linear("fc1", 8, 4).unwrap();

        store2.import_from_host(&exported).unwrap();
        let re_exported = store2.export_to_host().unwrap(); // test-only readback

        for (name, (_, values)) in &exported {
            let (_, values2) = re_exported.get(name).unwrap();
            for (a, b) in values.iter().zip(values2.iter()) {
                assert!((a - b).abs() < 1e-6, "import/export mismatch in {name}");
            }
        }
    }
}
