//! Top-level inference engine — `OccWorldCandle`.
//!
//! Provides the public-facing API:
//! - `OccWorldCandle::load` — load from a SafeTensors checkpoint
//! - `OccWorldCandle::dummy` — random weights for testing / benchmarking
//! - `OccWorldCandle::predict` — infer 15 future occupancy frames
//!
//! The `dummy` constructor allows end-to-end benchmarking (wall-clock timing,
//! shape verification, memory footprint) before the Phase-5 checkpoint exists.

use std::path::Path;
use std::time::Instant;

use candle_core::{DType, Device, Tensor};
use candle_nn::VarBuilder;

use crate::config::OccWorldConfig;
use crate::error::OccWorldError;
use crate::transformer::OccWorldTransformer;
use crate::vqvae::{decode_to_logits, encode_occupancy, VQVAEComponents};

// ── Output types ─────────────────────────────────────────────────────────────

/// A predicted future trajectory waypoint in 3-D grid coordinates.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct TrajectoryWaypoint {
    /// Frame index within the prediction horizon (0 = first predicted frame).
    pub frame: usize,
    /// Grid X position of the predicted agent centroid.
    pub grid_x: f32,
    /// Grid Y position of the predicted agent centroid.
    pub grid_y: f32,
    /// Grid Z position of the predicted agent centroid.
    pub grid_z: f32,
    /// Confidence score in `[0, 1]`.
    pub confidence: f32,
}

/// Outputs produced by one call to `OccWorldCandle::predict`.
pub struct InferenceOutput {
    /// Predicted semantic class for each voxel.
    ///
    /// Shape: `(1, 15, 200, 200, 16)`, dtype `u8`.
    /// Values are class indices in `[0, num_classes)`.
    pub sem_pred: Tensor,

    /// Trajectory priors extracted from the predicted occupancy.
    ///
    /// One waypoint per predicted frame, centred on the non-free voxel
    /// with the highest occupancy probability.  Empty when the model
    /// predicts all frames as free space.
    pub trajectory_priors: Vec<TrajectoryWaypoint>,

    /// Wall-clock time for the full `predict` call in milliseconds.
    pub inference_ms: f64,
}

// ── Main engine ───────────────────────────────────────────────────────────────

/// Native Rust OccWorld inference engine backed by Candle.
///
/// # Loading
///
/// ```no_run
/// # use wifi_densepose_occworld_candle::inference::OccWorldCandle;
/// # use wifi_densepose_occworld_candle::config::OccWorldConfig;
/// # use candle_core::Device;
/// # use std::path::Path;
/// let cfg = OccWorldConfig::default();
/// match OccWorldCandle::load(Path::new("/path/to/occworld.safetensors"), cfg) {
///     Ok(engine) => { /* use engine */ }
///     Err(_) => { /* fall back to Python bridge */ }
/// }
/// ```
pub struct OccWorldCandle {
    // Note: Device does not implement Debug; derive manually below.
    config: OccWorldConfig,
    vqvae: VQVAEComponents,
    transformer: OccWorldTransformer,
    device: Device,
}

impl std::fmt::Debug for OccWorldCandle {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("OccWorldCandle")
            .field("config", &self.config)
            .finish_non_exhaustive()
    }
}

impl OccWorldCandle {
    /// Load model weights from a SafeTensors checkpoint.
    ///
    /// Returns `Err` if the checkpoint does not exist, so callers can
    /// gracefully fall back to the Python bridge (`wifi-densepose-worldmodel`).
    pub fn load(
        checkpoint_path: &Path,
        config: OccWorldConfig,
    ) -> Result<Self, OccWorldError> {
        if !checkpoint_path.exists() {
            return Err(OccWorldError::CheckpointNotFound(
                checkpoint_path.display().to_string(),
            ));
        }

        let device = pick_device();

        // Load weights through the safe file-read path in `model::load_safetensors`.
        // This avoids the `unsafe` mmap block forbidden by our lint config, at the
        // cost of reading the full file into memory rather than memory-mapping it.
        // Switch to `VarBuilder::from_mmaped_safetensors` (in a crate that allows
        // unsafe) once the checkpoint is large enough that mmap matters.
        let tensors = crate::model::load_safetensors(checkpoint_path, &device)?;
        let vb = VarBuilder::from_tensors(tensors, DType::F32, &device);

        let vqvae = VQVAEComponents::new(&config, vb.clone()).map_err(OccWorldError::Candle)?;
        let transformer =
            OccWorldTransformer::new(config.clone(), vb).map_err(OccWorldError::Candle)?;

        Ok(Self {
            config,
            vqvae,
            transformer,
            device,
        })
    }

    /// Construct with random weights for testing and benchmarking.
    ///
    /// All shapes are correct; no checkpoint is required.
    pub fn dummy(config: OccWorldConfig, device: Device) -> Result<Self, OccWorldError> {
        let vqvae =
            VQVAEComponents::dummy(&config, &device).map_err(OccWorldError::Candle)?;
        let transformer =
            OccWorldTransformer::dummy(config.clone(), &device).map_err(OccWorldError::Candle)?;
        Ok(Self {
            config,
            vqvae,
            transformer,
            device,
        })
    }

    /// Infer 15 future occupancy frames from 16 past frames.
    ///
    /// # Arguments
    /// * `past_occupancy` — `(1, 16, 200, 200, 16)` tensor of `u8` class indices.
    ///
    /// # Returns
    /// [`InferenceOutput`] containing:
    /// - `sem_pred`: `(1, 15, 200, 200, 16)` u8 predicted class indices
    /// - `trajectory_priors`: one waypoint per predicted frame
    /// - `inference_ms`: wall-clock latency
    pub fn predict(&self, past_occupancy: &Tensor) -> Result<InferenceOutput, OccWorldError> {
        let t0 = Instant::now();

        let cfg = &self.config;
        let (b, f_in, h, w, d) = past_occupancy.dims5().map_err(OccWorldError::Candle)?;

        if h != cfg.grid_h || w != cfg.grid_w || d != cfg.grid_d {
            return Err(OccWorldError::ShapeMismatch(format!(
                "expected past_occupancy (_, _, {}, {}, {}), got (_, _, {h}, {w}, {d})",
                cfg.grid_h, cfg.grid_w, cfg.grid_d
            )));
        }

        // ── Step 1: VQVAE encode each past frame ──────────────────────────
        // Flatten batch*frames: (B, F, H, W, D) → (B*F, H, W, D)
        let occ_flat = past_occupancy
            .reshape((b * f_in, h, w, d))
            .map_err(OccWorldError::Candle)?;

        // Cast to u32 for class embedding (input is u8)
        let occ_u32 = occ_flat
            .to_dtype(DType::U32)
            .map_err(OccWorldError::Candle)?;

        // Class embedding → (B*F, base_channels, H, W*D)
        let embedded = self
            .vqvae
            .class_embed
            .forward(&occ_u32, cfg.grid_d)
            .map_err(OccWorldError::Candle)?;

        // Encode (stub) → (B*F, z_channels, token_h, token_w)
        let z = encode_occupancy(&embedded, cfg, &self.device)?;

        // quant_conv → (B*F, embed_dim, token_h, token_w)
        let z_e = self
            .vqvae
            .quant_conv
            .forward(&z)
            .map_err(OccWorldError::Candle)?;

        // Vector quantisation → z_q (B*F, embed_dim, token_h, token_w), indices
        // Reshape to (B*F, H*W, embed_dim) for VQCodebook.encode
        let (bf, e_dim, th, tw) = z_e.dims4().map_err(OccWorldError::Candle)?;
        let z_e_flat = z_e
            .permute((0, 2, 3, 1)) // (B*F, th, tw, embed_dim)
            .map_err(OccWorldError::Candle)?
            .reshape((bf, th * tw, e_dim))
            .map_err(OccWorldError::Candle)?;

        let (z_q_flat, _indices) = self
            .vqvae
            .codebook
            .encode(&z_e_flat)
            .map_err(OccWorldError::Candle)?;

        // Back to (B*F, embed_dim, th, tw) → (B, F, embed_dim, th, tw)
        let z_q = z_q_flat
            .reshape((bf, th, tw, e_dim))
            .map_err(OccWorldError::Candle)?
            .permute((0, 3, 1, 2)) // (B*F, embed_dim, th, tw)
            .map_err(OccWorldError::Candle)?
            .reshape((b, f_in, e_dim, th, tw))
            .map_err(OccWorldError::Candle)?;

        // ── Step 2: Transformer predicts future token logits ──────────────
        // Output: (B, F_out, vocab, th, tw)
        let pred_logits = self.transformer.forward(&z_q)?;

        let f_out = pred_logits.dim(1).map_err(OccWorldError::Candle)?;

        // ── Step 3: Argmax over vocab dim → predicted token indices ───────
        let pred_indices = pred_logits
            .argmax(2) // (B, F_out, th, tw)  — over vocab dim
            .map_err(OccWorldError::Candle)?;

        // ── Step 4: Decode token indices → z_q values ────────────────────
        // Flatten to (B*F_out * th * tw,) for codebook lookup
        let idx_flat = pred_indices
            .flatten_all()
            .map_err(OccWorldError::Candle)?;
        let z_decoded = self
            .vqvae
            .codebook
            .decode(&idx_flat)
            .map_err(OccWorldError::Candle)?; // (B*F_out*th*tw, embed_dim)

        // Reshape to (B*F_out, embed_dim, th, tw) for post_quant_conv
        let z_dec_4d = z_decoded
            .reshape((b * f_out, e_dim, th, tw))
            .map_err(OccWorldError::Candle)?;

        let z_post = self
            .vqvae
            .post_quant_conv
            .forward(&z_dec_4d)
            .map_err(OccWorldError::Candle)?;

        // ── Step 5: Decode to class logits (stub) → class predictions ─────
        let class_logits = decode_to_logits(&z_post, cfg, &self.device)?;
        // class_logits: (B*F_out, num_classes, H, W, D)
        // Argmax over class dim → (B*F_out, H, W, D)
        let sem_flat = class_logits
            .argmax(1)
            .map_err(OccWorldError::Candle)?
            .to_dtype(DType::U8)
            .map_err(OccWorldError::Candle)?;

        let sem_pred = sem_flat
            .reshape((b, f_out, cfg.grid_h, cfg.grid_w, cfg.grid_d))
            .map_err(OccWorldError::Candle)?;

        // ── Step 6: Extract trajectory priors ─────────────────────────────
        let trajectory_priors = extract_trajectory_priors(&sem_pred, cfg, f_out)?;

        let inference_ms = t0.elapsed().as_secs_f64() * 1000.0;

        Ok(InferenceOutput {
            sem_pred,
            trajectory_priors,
            inference_ms,
        })
    }
}

// ── Trajectory prior extraction ───────────────────────────────────────────────

/// Extract one trajectory waypoint per predicted frame.
///
/// For each frame, finds the non-free voxel with the highest probability
/// (approximated by the centroid of all non-free voxels, weighted equally).
/// Returns an empty `Vec` when all frames are predicted as free space.
fn extract_trajectory_priors(
    sem_pred: &Tensor,
    cfg: &OccWorldConfig,
    f_out: usize,
) -> Result<Vec<TrajectoryWaypoint>, OccWorldError> {
    // sem_pred: (1, F_out, H, W, D) u8
    // Pull to CPU Vec for coordinate extraction — lightweight post-processing
    let data: Vec<u8> = sem_pred
        .flatten_all()
        .map_err(OccWorldError::Candle)?
        .to_vec1()
        .map_err(OccWorldError::Candle)?;

    let h = cfg.grid_h;
    let w = cfg.grid_w;
    let d = cfg.grid_d;
    let frame_stride = h * w * d;

    let mut waypoints = Vec::with_capacity(f_out);
    for fi in 0..f_out {
        let frame_slice = &data[fi * frame_stride..(fi + 1) * frame_stride];
        let mut sum_x = 0.0f64;
        let mut sum_y = 0.0f64;
        let mut sum_z = 0.0f64;
        let mut count = 0usize;

        for (idx, &cls) in frame_slice.iter().enumerate() {
            if cls != cfg.free_class {
                let xi = idx / (w * d);
                let yi = (idx % (w * d)) / d;
                let zi = idx % d;
                sum_x += xi as f64;
                sum_y += yi as f64;
                sum_z += zi as f64;
                count += 1;
            }
        }

        if count > 0 {
            let n = count as f64;
            waypoints.push(TrajectoryWaypoint {
                frame: fi,
                grid_x: (sum_x / n) as f32,
                grid_y: (sum_y / n) as f32,
                grid_z: (sum_z / n) as f32,
                confidence: (count as f32) / (frame_stride as f32),
            });
        }
    }
    Ok(waypoints)
}

// ── Device selection ──────────────────────────────────────────────────────────

fn pick_device() -> Device {
    #[cfg(feature = "cuda")]
    if let Ok(d) = Device::cuda_if_available(0) {
        return d;
    }
    Device::Cpu
}

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

    fn small_cfg() -> OccWorldConfig {
        OccWorldConfig {
            grid_h: 8,
            grid_w: 8,
            grid_d: 4,
            num_classes: 4,
            free_class: 3,
            base_channels: 8,
            z_channels: 8,
            codebook_size: 4,
            embed_dim: 8,
            num_frames: 2,
            token_h: 4,
            token_w: 4,
            num_heads: 2,
            num_layers: 1,
            ffn_hidden: 16,
        }
    }

    #[test]
    fn test_dummy_predict_shape() -> Result<(), OccWorldError> {
        let device = Device::Cpu;
        let cfg = small_cfg();
        let engine = OccWorldCandle::dummy(cfg.clone(), device.clone())?;

        // (1, 2, 8, 8, 4) — batch=1, 2 past frames (matches num_frames)
        let past = Tensor::zeros(
            (1, cfg.num_frames, cfg.grid_h, cfg.grid_w, cfg.grid_d),
            DType::U8,
            &device,
        )
        .map_err(OccWorldError::Candle)?;

        let out = engine.predict(&past)?;
        let dims = out.sem_pred.dims();
        assert_eq!(dims[0], 1, "batch dim");
        assert_eq!(dims[1], cfg.num_frames, "frame dim");
        assert_eq!(dims[2], cfg.grid_h, "H dim");
        assert_eq!(dims[3], cfg.grid_w, "W dim");
        assert_eq!(dims[4], cfg.grid_d, "D dim");

        Ok(())
    }

    #[test]
    fn test_load_nonexistent_checkpoint() {
        let cfg = small_cfg();
        let result = OccWorldCandle::load(Path::new("/no/such/checkpoint.safetensors"), cfg);
        assert!(
            matches!(result, Err(OccWorldError::CheckpointNotFound(_))),
            "expected CheckpointNotFound, got {result:?}"
        );
    }
}