//! Quantile Output Layer for TFT
//!
//! Implements quantile regression for uncertainty estimation in multi-horizon
//! forecasting with proper quantile loss and monotonicity constraints.

use std::sync::Arc;

use cudarc::driver::CudaStream;
use ml_core::MLError;

use crate::gpu_tensor::{
    gpu_add, gpu_exp, gpu_log, gpu_add_scalar, gpu_mean_all, gpu_narrow_2d, gpu_scale,
    gpu_select_dim1, gpu_stack_2d, gpu_sub, gpu_abs,
    GpuLinear, GpuTensor,
};

/// Quantile output layer for uncertainty estimation
#[derive(Debug, Clone)]
pub struct QuantileLayer {
    pub hidden_dim: usize,
    pub prediction_horizon: usize,
    pub num_quantiles: usize,
    pub quantile_levels: Vec<f64>,
    // Separate linear layers for each quantile
    quantile_projections: Vec<GpuLinear>,
    // Monotonicity constraint layers
    monotonicity_weights: Vec<GpuLinear>,
    #[allow(dead_code)]
    stream: Arc<CudaStream>,
}

impl QuantileLayer {
    pub fn new(
        hidden_dim: usize,
        prediction_horizon: usize,
        num_quantiles: usize,
        stream: &Arc<CudaStream>,
    ) -> Result<Self, MLError> {
        let quantile_levels = (1..=num_quantiles)
            .map(|i| i as f64 / (num_quantiles + 1) as f64)
            .collect::<Vec<_>>();

        let mut quantile_projections = Vec::new();
        let mut monotonicity_weights = Vec::new();

        for i in 0..num_quantiles {
            let projection = GpuLinear::new(hidden_dim, prediction_horizon, stream)?;
            quantile_projections.push(projection);

            if i > 0 {
                let mono_weight = GpuLinear::new(hidden_dim, prediction_horizon, stream)?;
                monotonicity_weights.push(mono_weight);
            }
        }

        Ok(Self {
            hidden_dim,
            prediction_horizon,
            num_quantiles,
            quantile_levels,
            quantile_projections,
            monotonicity_weights,
            stream: Arc::clone(stream),
        })
    }

    /// Forward pass: input [batch, `hidden_dim`] -> [batch, `prediction_horizon`, `num_quantiles`]
    pub fn forward(&self, x: &GpuTensor) -> Result<GpuTensor, MLError> {
        // Handle 3D input: use last timestep via GPU op
        if x.shape.len() == 3 {
            let seq = x.dim(1)?;
            let last_t = seq.saturating_sub(1);
            // gpu_select_dim1 extracts [batch, features] from [batch, seq, features]
            let last_step = gpu_select_dim1(x, last_t)?;
            return self.forward(&last_step);
        }

        if x.shape.len() != 2 {
            return Err(MLError::InvalidInput(format!(
                "QuantileLayer expects 2D or 3D input, got {:?}",
                x.shape
            )));
        }

        let batch_size = x.dim(0)?;
        let final_hidden_dim = x.dim(1)?;

        if final_hidden_dim != self.hidden_dim {
            return Err(MLError::InvalidInput(format!(
                "Expected hidden dimension {}, got {}",
                self.hidden_dim, final_hidden_dim
            )));
        }

        // Compute base quantile predictions — all on GPU
        let mut quantile_outputs: Vec<GpuTensor> = Vec::new();

        // First quantile (no monotonicity constraint)
        let first_proj = self.quantile_projections.first().ok_or_else(|| {
            MLError::InvalidInput("No quantile projections".to_owned())
        })?;
        let q0 = first_proj.forward(x)?;
        quantile_outputs.push(q0);

        // Remaining quantiles with monotonicity constraints
        for i in 1..self.num_quantiles {
            let proj = self.quantile_projections.get(i).ok_or_else(|| {
                MLError::InvalidInput(format!("Missing quantile projection {}", i))
            })?;
            let raw_output = proj.forward(x)?;
            let mono_weight = self.monotonicity_weights.get(i - 1).ok_or_else(|| {
                MLError::InvalidInput(format!("Missing monotonicity weight {}", i))
            })?;
            let mono_adjustment = mono_weight.forward(x)?;
            let combined = gpu_add(&raw_output, &mono_adjustment)?;

            // Softplus: log(1 + exp(x))
            let exp_x = gpu_exp(&combined)?;
            let one_plus_exp = gpu_add_scalar(&exp_x, 1.0)?;
            let softplus_out = gpu_log(&one_plus_exp)?;

            let prev_quantile = quantile_outputs.get(i - 1).ok_or_else(|| {
                MLError::InvalidInput("Missing previous quantile".to_owned())
            })?;
            let current_quantile = gpu_add(prev_quantile, &softplus_out)?;
            quantile_outputs.push(current_quantile);
        }

        // Build output on GPU: [batch, prediction_horizon, num_quantiles]
        // Each quantile_output is [batch, prediction_horizon]
        // Stack them along a new last dimension.
        //
        // gpu_stack_2d stacks N tensors of [R, C] -> [N, R, C].
        // We want [batch, horizon, num_quantiles], so we stack to get
        // [num_quantiles, batch, horizon] then transpose dims 0 and 2.
        //
        // Alternative: for each batch, stack the quantile slices per horizon.
        // Most direct: stack quantile outputs to get [num_quantiles, batch, horizon],
        // then we need a dim permutation. Since we don't have a 3D transpose,
        // build the result by extracting per-batch rows.
        //
        // Since quantile_outputs[q] is [batch, horizon], we can extract row b
        // from each quantile and concatenate to get [num_quantiles, horizon],
        // then stack across batches.
        let mut batch_results: Vec<GpuTensor> = Vec::with_capacity(batch_size);
        for b in 0..batch_size {
            let mut q_rows: Vec<GpuTensor> = Vec::with_capacity(self.num_quantiles);
            for q_out in &quantile_outputs {
                // Extract row b: [1, horizon]
                let row = gpu_narrow_2d(q_out, 0, b, 1)?;
                q_rows.push(row);
            }
            // Stack [num_quantiles] items of [1, horizon] -> [num_quantiles, 1, horizon]
            let stacked = gpu_stack_2d(&q_rows)?;
            // stacked is [num_quantiles, 1, horizon], reshape to [horizon, num_quantiles]
            // by first reshaping to [num_quantiles, horizon] then transposing
            let flat = stacked.reshape(&[self.num_quantiles, self.prediction_horizon])?;
            let transposed = crate::gpu_tensor::gpu_transpose(&flat)?; // [horizon, num_quantiles]
            batch_results.push(transposed);
        }
        // Stack [batch] items of [horizon, num_quantiles] -> [batch, horizon, num_quantiles]
        gpu_stack_2d(&batch_results)
    }

    /// Quantile loss computation on GPU (returns scalar f32)
    ///
    /// Uses GPU ops for element-wise computation, only a single scalar
    /// is downloaded at the end via `gpu_mean_all`.
    pub fn quantile_loss(
        &self,
        predictions: &GpuTensor,
        targets: &GpuTensor,
    ) -> Result<f32, MLError> {
        // predictions: [batch, horizon, quantiles]
        // targets: [batch, horizon]
        if predictions.shape.len() != 3 || targets.shape.len() != 2 {
            return Err(MLError::InvalidInput(format!(
                "quantile_loss expects pred=[B,H,Q] target=[B,H], got pred={:?} target={:?}",
                predictions.shape, targets.shape
            )));
        }

        let batch = predictions.dim(0)?;
        let horizon = predictions.dim(1)?;
        let n_q = predictions.dim(2)?;

        // Flatten predictions to [batch * horizon, num_quantiles]
        let pred_flat = predictions.reshape(&[batch * horizon, n_q])?;
        // Flatten targets to [batch * horizon, 1] and broadcast for each quantile
        let target_flat = targets.reshape(&[batch * horizon, 1])?;

        // Compute quantile loss for each quantile level on GPU
        let mut total_loss_sum = 0.0_f64;
        let mut total_count = 0_usize;

        for (qi, &tau) in self.quantile_levels.iter().enumerate() {
            if qi >= n_q {
                break;
            }
            // Extract quantile qi predictions: [batch * horizon, 1]
            let pred_q = gpu_narrow_2d(&pred_flat, 1, qi, 1)?;
            // residual = target - pred
            let residual = gpu_sub(&target_flat, &pred_q)?;
            let abs_residual = gpu_abs(&residual)?;

            // Compute asymmetric loss:
            // loss = tau * residual where residual >= 0
            // loss = (tau - 1) * residual where residual < 0
            // Equivalent to: loss = |tau - I(residual < 0)| * |residual|
            //               = max(tau * residual, (tau - 1) * residual)
            // Which equals: 0.5 * |residual| + (tau - 0.5) * residual
            let half_abs = gpu_scale(&abs_residual, 0.5)?;
            let tau_shifted = (tau - 0.5) as f32;
            let shifted_residual = gpu_scale(&residual, tau_shifted)?;
            let loss_per_elem = gpu_add(&half_abs, &shifted_residual)?;

            // Mean across batch * horizon
            let mean_loss = gpu_mean_all(&loss_per_elem)?;
            total_loss_sum += mean_loss as f64;
            total_count += 1;
        }

        if total_count > 0 {
            Ok((total_loss_sum / total_count as f64) as f32)
        } else {
            Ok(0.0)
        }
    }

    pub fn get_quantile_levels(&self) -> Vec<f64> {
        self.quantile_levels.clone()
    }
}

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

    fn test_stream() -> Arc<CudaStream> {
        let ctx = CudaContext::new(0).expect("CUDA context required");
        ctx.new_stream().expect("Failed to create CUDA stream")
    }

    #[test]
    fn test_quantile_layer_creation() -> Result<(), MLError> {
        let stream = test_stream();
        let ql = QuantileLayer::new(64, 10, 9, &stream)?;
        assert_eq!(ql.hidden_dim, 64);
        assert_eq!(ql.prediction_horizon, 10);
        assert_eq!(ql.num_quantiles, 9);
        assert_eq!(ql.quantile_levels.len(), 9);
        Ok(())
    }

    #[test]
    fn test_quantile_levels() -> Result<(), MLError> {
        let stream = test_stream();
        let ql = QuantileLayer::new(32, 5, 9, &stream)?;
        let levels = ql.get_quantile_levels();
        assert_eq!(levels.len(), 9);
        assert!((levels[0] - 0.1).abs() < 0.01);
        assert!((levels[8] - 0.9).abs() < 0.01);
        for i in 1..levels.len() {
            assert!(levels[i] > levels[i - 1]);
        }
        Ok(())
    }

    #[test]
    fn test_quantile_forward_2d() -> Result<(), MLError> {
        let stream = test_stream();
        let ql = QuantileLayer::new(32, 5, 7, &stream)?;
        let input = GpuTensor::from_vec(vec![1.0_f32; 64], &[2, 32], &stream)?;
        let output = ql.forward(&input)?;
        assert_eq!(output.shape, vec![2, 5, 7]);
        Ok(())
    }
}
