//! Implicit Q-Learning (IQL) -- Kostrikov et al. 2021
//!
//! IQL avoids querying out-of-distribution actions by using expectile regression
//! on the value function. This makes it ideal for offline RL with historical
//! market data where we cannot explore.
//!
//! Key components:
//! - Value network V(s): trained with asymmetric (expectile) loss
//! - Q-network Q(s,a): trained with standard Bellman backup using V(s') as target
//! - Policy extraction: advantage-weighted regression
//!
//! # References
//!
//! Kostrikov, I., Nair, A., & Levine, S. (2021). Offline Reinforcement Learning
//! with Implicit Q-Learning. *arXiv preprint arXiv:2110.06169*.

use std::sync::Arc;

use cudarc::cublas::CudaBlas;
use cudarc::driver::CudaStream;

use ml_core::cuda_autograd::{
    ActivationKernels, GpuLinear, GpuTensor, GpuVarStore,
};
use ml_core::MLError;

/// IQL configuration parameters.
///
/// Controls the expectile regression asymmetry and advantage-weighted policy
/// extraction temperature. Higher `expectile_tau` produces more optimistic
/// value estimates (extracting the value of the best in-distribution action),
/// while higher `advantage_temperature` makes the extracted policy more greedy.
#[derive(Debug, Clone)]
pub struct IqlConfig {
    /// Expectile parameter tau (0.5 = symmetric, 0.7-0.9 = asymmetric toward high values)
    pub expectile_tau: f32,
    /// Temperature beta for advantage-weighted policy extraction
    pub advantage_temperature: f32,
    /// Learning rate for value network
    pub value_lr: f64,
    /// Discount factor
    pub gamma: f32,
    /// Whether IQL mode is active
    pub enabled: bool,
}

impl Default for IqlConfig {
    fn default() -> Self {
        Self {
            expectile_tau: 0.7,
            advantage_temperature: 3.0,
            value_lr: 3e-4,
            gamma: 0.99,
            enabled: false,
        }
    }
}

/// Value network V(s) for IQL.
///
/// Same architecture as Q-network but outputs a single scalar per state.
/// Trained with expectile regression loss to approximate the value of the
/// best in-distribution action without explicit maximization.
pub struct ValueNetwork {
    layer1: GpuLinear,
    layer2: GpuLinear,
    output: GpuLinear,
    store: GpuVarStore,
    cublas: CudaBlas,
    activations: ActivationKernels,
    stream: Arc<CudaStream>,
}

impl std::fmt::Debug for ValueNetwork {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ValueNetwork")
            .field("layers", &"[GpuLinear, GpuLinear, GpuLinear(1)]")
            .finish()
    }
}

impl ValueNetwork {
    /// Create a new value network with the given dimensions.
    ///
    /// # Errors
    ///
    /// Returns an error if weight initialization fails.
    pub fn new(
        state_dim: usize,
        hidden_dim: usize,
        stream: Arc<CudaStream>,
    ) -> Result<Self, MLError> {
        let mut store = GpuVarStore::new(stream.clone());
        let layer1 = store.linear("v_layer1", state_dim, hidden_dim)?;
        let layer2 = store.linear("v_layer2", hidden_dim, hidden_dim)?;
        let output = store.linear("v_output", hidden_dim, 1)?;

        let cublas = CudaBlas::new(stream.clone()).map_err(|e| {
            MLError::ModelError(format!("cuBLAS init failed: {e}"))
        })?;
        let activations = ActivationKernels::new(&stream)?;

        Ok(Self {
            layer1,
            layer2,
            output,
            store,
            cublas,
            activations,
            stream,
        })
    }

    /// Forward pass: state -> scalar value V(s).
    ///
    /// Architecture: `Linear` -> `ReLU` -> `Linear` -> `ReLU` -> `Linear`.
    /// Output shape: `[batch, 1]`.
    ///
    /// # Errors
    ///
    /// Returns an error if any tensor operation fails.
    pub fn forward(&self, state: &GpuTensor) -> Result<GpuTensor, MLError> {
        let (x, _) = self.layer1.forward(state, &self.store, &self.cublas, &self.stream)?;
        let (x, _) = self.activations.relu_fwd(&x, &self.stream)?;
        let (x, _) = self.layer2.forward(&x, &self.store, &self.cublas, &self.stream)?;
        let (x, _) = self.activations.relu_fwd(&x, &self.stream)?;
        let (out, _) = self.output.forward(&x, &self.store, &self.cublas, &self.stream)?;
        Ok(out) // [batch, 1]
    }
}

/// Compute expectile regression loss for the value function (CPU fallback).
///
/// The asymmetric loss function:
///
/// ```text
/// L_tau(u) = |tau - 1(u < 0)| * u^2
/// ```
///
/// When tau > 0.5, this penalizes underestimation more than overestimation,
/// effectively extracting the value of the best in-distribution action
/// without querying out-of-distribution actions.
///
/// NOTE: This is a CPU-side computation for small tensors. The hot-path
/// expectile loss runs inside a fused CUDA kernel.
///
/// # Arguments
///
/// * `predicted_v` - V(s) predictions, slice of length `batch`
/// * `target_q` - Q(s,a) targets, slice of length `batch`
/// * `tau` - Expectile parameter (0.5 = MSE, >0.5 = optimistic)
///
/// # Errors
///
/// Returns an error if inputs have different lengths.
pub fn expectile_loss(
    predicted_v: &[f32],
    target_q: &[f32],
    tau: f32,
) -> Result<f32, MLError> {
    if predicted_v.len() != target_q.len() {
        return Err(MLError::DimensionMismatch {
            expected: predicted_v.len(),
            actual: target_q.len(),
        });
    }
    if predicted_v.is_empty() {
        return Err(MLError::InvalidInput("empty inputs".into()));
    }

    let mut total = 0.0_f32;
    for (p, t) in predicted_v.iter().zip(target_q.iter()) {
        let diff = t - p;
        let weight = if diff >= 0.0 { tau } else { 1.0 - tau };
        total += weight * diff * diff;
    }
    Ok(total / predicted_v.len() as f32)
}

/// Compute advantage-weighted action probabilities (CPU fallback).
///
/// Extracts a policy via advantage-weighted regression:
///
/// ```text
/// pi(a|s) proportional to exp(beta * A(s,a))
/// ```
///
/// where `A(s,a) = Q(s,a) - V(s)`.
///
/// # Arguments
///
/// * `q_values` - Q-values for all actions, flat `[batch * num_actions]` row-major
/// * `v_values` - Value estimates, `[batch]`
/// * `num_actions` - Number of actions per state
/// * `temperature` - Inverse temperature beta (higher = more greedy)
///
/// # Returns
///
/// Action probabilities, flat `[batch * num_actions]` row-major.
///
/// # Errors
///
/// Returns an error if dimensions are inconsistent.
pub fn advantage_weighted_action(
    q_values: &[f32],
    v_values: &[f32],
    num_actions: usize,
    temperature: f32,
) -> Result<Vec<f32>, MLError> {
    let batch = v_values.len();
    if q_values.len() != batch * num_actions {
        return Err(MLError::DimensionMismatch {
            expected: batch * num_actions,
            actual: q_values.len(),
        });
    }

    let mut probs = Vec::with_capacity(batch * num_actions);
    for b in 0..batch {
        let v = v_values.get(b).copied().ok_or_else(|| {
            MLError::ModelError("v_values index out of bounds".into())
        })?;

        // Compute scaled advantages and find max for numerical stability
        let base = b * num_actions;
        let mut max_adv = f32::NEG_INFINITY;
        for a in 0..num_actions {
            let q = q_values.get(base + a).copied().ok_or_else(|| {
                MLError::ModelError("q_values index out of bounds".into())
            })?;
            let adv = (q - v) * temperature;
            let clamped = adv.clamp(-10.0, 10.0);
            if clamped > max_adv {
                max_adv = clamped;
            }
        }

        // Numerically stable softmax
        let mut exp_sum = 0.0_f32;
        let mut exp_vals = Vec::with_capacity(num_actions);
        for a in 0..num_actions {
            let q = q_values.get(base + a).copied().unwrap_or(0.0);
            let adv = ((q - v) * temperature).clamp(-10.0, 10.0);
            let e = (adv - max_adv).exp();
            exp_vals.push(e);
            exp_sum += e;
        }

        for e in &exp_vals {
            probs.push(e / exp_sum);
        }
    }

    Ok(probs)
}

#[cfg(test)]
#[allow(
    clippy::manual_let_else,
    clippy::assertions_on_result_states
)]
mod tests {
    use super::*;

    #[test]
    fn test_iql_config_default() {
        let config = IqlConfig::default();
        assert!((config.expectile_tau - 0.7).abs() < 1e-6);
        assert!((config.advantage_temperature - 3.0).abs() < 1e-6);
        assert!(!config.enabled);
    }

    #[test]
    fn test_value_network_forward() {
        let device = ml_core::device::MlDevice::cuda(0).expect("CUDA required");
        let stream = device.cuda_stream().expect("stream").clone();

        let net = ValueNetwork::new(4, 16, stream.clone());
        assert!(net.is_ok());
        let net = match net {
            Ok(n) => n,
            Err(_) => return,
        };

        // Batch of 3 states with dim 4
        let host_data = vec![0.1_f32; 3 * 4];
        let states = GpuTensor::from_host(&host_data, vec![3, 4], &stream);
        let states = match states {
            Ok(s) => s,
            Err(_) => return,
        };

        let values = net.forward(&states);
        assert!(values.is_ok());
        let values = match values {
            Ok(v) => v,
            Err(_) => return,
        };

        // Output should be [3, 1]
        assert_eq!(values.shape(), &[3, 1]);
    }

    #[test]
    fn test_expectile_loss_symmetric() {
        // With tau=0.5, expectile loss should equal MSE
        let predicted = vec![1.0_f32, 2.0, 3.0];
        let target = vec![1.5_f32, 2.5, 3.5];

        let loss = expectile_loss(&predicted, &target, 0.5);
        assert!(loss.is_ok());
        let loss_val = loss.unwrap_or(f32::NAN);

        // MSE of diff=[0.5, 0.5, 0.5] should be 0.5 * 0.25 = 0.125
        // With symmetric weighting (tau=0.5), weight = 0.5 everywhere
        // weighted_loss = 0.5 * 0.25 = 0.125 per element, mean = 0.125
        assert!(
            (loss_val - 0.125).abs() < 1e-4,
            "Expected ~0.125 for symmetric expectile, got {}",
            loss_val,
        );
    }

    #[test]
    fn test_expectile_loss_asymmetric() {
        // With tau=0.9, underestimation (diff>0) weighted 0.9, overestimation weighted 0.1
        let predicted = vec![1.0_f32, 3.0];
        let target = vec![2.0_f32, 2.0];

        let loss_asym = expectile_loss(&predicted, &target, 0.9);
        let loss_sym = expectile_loss(&predicted, &target, 0.5);

        // Asymmetric should weight the underestimation (first elem) more
        assert!(loss_asym.is_ok());
        assert!(loss_sym.is_ok());
    }

    #[test]
    fn test_advantage_weighted_action() {
        // 2 states, 3 actions
        let q_values = vec![1.0_f32, 2.0, 3.0, 3.0, 1.0, 2.0];
        let v_values = vec![2.0_f32, 2.0];

        let probs = advantage_weighted_action(&q_values, &v_values, 3, 3.0);
        assert!(probs.is_ok(), "advantage_weighted_action should succeed");
        let probs = match probs {
            Ok(p) => p,
            Err(_) => return,
        };

        assert_eq!(probs.len(), 6); // 2 * 3

        // Probabilities should sum to ~1.0 per row
        for b in 0..2 {
            let base = b * 3;
            let sum: f32 = probs.get(base..base + 3)
                .map(|s| s.iter().sum())
                .unwrap_or(0.0);
            assert!(
                (sum - 1.0).abs() < 1e-4,
                "Row sum should be ~1.0, got {}",
                sum,
            );
        }
    }
}
