#![allow(unsafe_code)] // Required for CUDA kernel launch

//! Signal adapter utilities for converting between model outputs and action scores.
//!
//! Pure cudarc `CudaSlice<f32>` operations -- zero Candle Tensor usage.
//! Three CUDA kernels bridge the gap between model predictions and the 5-action
//! DQN exposure space (Short100, Short50, Flat, Long50, Long100).

use cudarc::driver::{CudaContext, CudaFunction, CudaSlice, CudaStream, LaunchConfig, PushKernelArg};
use cudarc::nvrtc::Ptx;
use std::sync::{Arc, OnceLock};

use crate::MLError;

// ── Kernel compilation ──────────────────────────────────────────────────

static SIGNAL_ADAPTER_PTX: OnceLock<Result<Ptx, String>> = OnceLock::new();

fn compile_signal_adapter_ptx(context: &Arc<CudaContext>) -> Result<Ptx, String> {
    let kernel_src = include_str!("signal_adapter_kernel.cu");
    crate::cuda_pipeline::compile_ptx_for_device(kernel_src, context)
}

/// Lazily compiled kernel handle set for signal adapter operations.
struct KernelSet {
    ppo_exposure: CudaFunction,
    signal_to_action: CudaFunction,
    tft_extract: CudaFunction,
}

fn load_kernels(context: &Arc<CudaContext>) -> Result<KernelSet, MLError> {
    let ptx = SIGNAL_ADAPTER_PTX
        .get_or_init(|| compile_signal_adapter_ptx(context))
        .as_ref()
        .map_err(|e| MLError::ModelError(format!("signal_adapter PTX: {e}")))?;
    let module = context
        .load_module(ptx.clone())
        .map_err(|e| MLError::ModelError(format!("signal_adapter module load: {e}")))?;

    let ppo_exposure = module
        .load_function("ppo_to_exposure_scores_kernel")
        .map_err(|e| MLError::ModelError(format!("ppo_to_exposure_scores_kernel load: {e}")))?;
    let signal_to_action = module
        .load_function("signal_to_action_scores_kernel")
        .map_err(|e| MLError::ModelError(format!("signal_to_action_scores_kernel load: {e}")))?;
    let tft_extract = module
        .load_function("tft_quantile_extract_kernel")
        .map_err(|e| MLError::ModelError(format!("tft_quantile_extract_kernel load: {e}")))?;

    Ok(KernelSet { ppo_exposure, signal_to_action, tft_extract })
}

// ── Public API ──────────────────────────────────────────────────────────

/// Aggregate PPO 45-action softmax probabilities into 5 exposure scores.
///
/// PPO uses a 45-action factored space (5 exposure x 3 order x 3 urgency).
/// This collapses the order/urgency dimensions to produce a `[batch * 5]`
/// `CudaSlice` whose per-row argmax gives the dominant exposure bucket.
///
/// # Arguments
/// * `probs` - GPU-resident `[batch * 45]` softmax probabilities
/// * `batch` - number of rows
/// * `stream` - CUDA stream for kernel launch and allocation
///
/// # Returns
/// GPU-resident `CudaSlice<f32>` of shape `[batch * 5]`.
///
/// # Errors
/// Returns `MLError::ModelError` on kernel compilation or launch failure.
pub fn ppo_to_exposure_scores(
    probs: &CudaSlice<f32>,
    batch: usize,
    stream: &Arc<CudaStream>,
) -> Result<CudaSlice<f32>, MLError> {
    let context = stream.context();
    let kernels = load_kernels(&context)?;

    let mut out = stream
        .alloc_zeros::<f32>(batch * 5)
        .map_err(|e| MLError::ModelError(format!("alloc exposure scores: {e}")))?;

    let batch_i32 = batch as i32;
    let grid = ((batch as u32 + 255) / 256).max(1);
    let cfg = LaunchConfig {
        grid_dim: (grid, 1, 1),
        block_dim: (256, 1, 1),
        shared_mem_bytes: 0,
    };

    // Safety: argument order matches ppo_to_exposure_scores_kernel(probs, out, batch)
    unsafe {
        stream
            .launch_builder(&kernels.ppo_exposure)
            .arg(probs)
            .arg(&mut out)
            .arg(&batch_i32)
            .launch(cfg)
            .map_err(|e| MLError::ModelError(format!("ppo_to_exposure_scores launch: {e}")))?;
    }

    Ok(out)
}

/// Map scalar return predictions (bps) to 5-action one-hot-ish scores.
///
/// Thresholds partition the real line into five bins matching the DQN
/// exposure actions:
/// - `pred < -high`           => Short100 (action 0)
/// - `-high <= pred < -low`   => Short50  (action 1)
/// - `-low  <= pred <= +low`  => Flat     (action 2)
/// - `+low  <  pred <= +high` => Long50   (action 3)
/// - `pred  >  +high`         => Long100  (action 4)
///
/// # Arguments
/// * `predictions` - GPU-resident `[batch]` scalar signals
/// * `batch` - number of elements
/// * `high_threshold_bps` - strong signal threshold
/// * `low_threshold_bps` - mild signal threshold
/// * `stream` - CUDA stream for kernel launch and allocation
///
/// # Returns
/// GPU-resident `CudaSlice<f32>` of shape `[batch * 5]` (one-hot rows).
///
/// # Errors
/// Returns `MLError::ModelError` on kernel compilation or launch failure.
pub fn signal_to_action_scores(
    predictions: &CudaSlice<f32>,
    batch: usize,
    high_threshold_bps: f32,
    low_threshold_bps: f32,
    stream: &Arc<CudaStream>,
) -> Result<CudaSlice<f32>, MLError> {
    let context = stream.context();
    let kernels = load_kernels(&context)?;

    let mut out = stream
        .alloc_zeros::<f32>(batch * 5)
        .map_err(|e| MLError::ModelError(format!("alloc action scores: {e}")))?;

    let batch_i32 = batch as i32;
    let grid = ((batch as u32 + 255) / 256).max(1);
    let cfg = LaunchConfig {
        grid_dim: (grid, 1, 1),
        block_dim: (256, 1, 1),
        shared_mem_bytes: 0,
    };

    // Safety: argument order matches signal_to_action_scores_kernel(
    //   predictions, out_scores, high_threshold_bps, low_threshold_bps, batch)
    unsafe {
        stream
            .launch_builder(&kernels.signal_to_action)
            .arg(predictions)
            .arg(&mut out)
            .arg(&high_threshold_bps)
            .arg(&low_threshold_bps)
            .arg(&batch_i32)
            .launch(cfg)
            .map_err(|e| MLError::ModelError(format!("signal_to_action_scores launch: {e}")))?;
    }

    Ok(out)
}

/// Extract the median signal from TFT quantile predictions.
///
/// Input layout: `[batch * horizon * num_quantiles]` (typically `[batch * 1 * 3]`).
/// Returns `[batch]` -- the median (quantile index 1) from the first horizon step.
///
/// # Arguments
/// * `quantiles` - GPU-resident `[batch * horizon * num_quantiles]` data
/// * `batch` - batch size
/// * `horizon` - number of horizon steps (>= 1)
/// * `num_quantiles` - number of quantiles (>= 2, median at index 1)
/// * `stream` - CUDA stream for kernel launch and allocation
///
/// # Returns
/// GPU-resident `CudaSlice<f32>` of shape `[batch]`.
///
/// # Errors
/// Returns `MLError::ModelError` if dimensions are invalid or kernel fails.
pub fn tft_quantile_to_signal(
    quantiles: &CudaSlice<f32>,
    batch: usize,
    horizon: usize,
    num_quantiles: usize,
    stream: &Arc<CudaStream>,
) -> Result<CudaSlice<f32>, MLError> {
    if horizon < 1 {
        return Err(MLError::ModelError(
            "tft_quantile_to_signal: horizon must be >= 1".to_owned(),
        ));
    }
    if num_quantiles < 2 {
        return Err(MLError::ModelError(format!(
            "tft_quantile_to_signal: need >= 2 quantiles, got {num_quantiles}"
        )));
    }

    let context = stream.context();
    let kernels = load_kernels(&context)?;

    let mut out = stream
        .alloc_zeros::<f32>(batch)
        .map_err(|e| MLError::ModelError(format!("alloc tft signal: {e}")))?;

    let batch_i32 = batch as i32;
    let horizon_i32 = horizon as i32;
    let num_q_i32 = num_quantiles as i32;
    let grid = ((batch as u32 + 255) / 256).max(1);
    let cfg = LaunchConfig {
        grid_dim: (grid, 1, 1),
        block_dim: (256, 1, 1),
        shared_mem_bytes: 0,
    };

    // Safety: argument order matches tft_quantile_extract_kernel(
    //   quantiles, out_signal, horizon, num_quantiles, batch)
    unsafe {
        stream
            .launch_builder(&kernels.tft_extract)
            .arg(quantiles)
            .arg(&mut out)
            .arg(&horizon_i32)
            .arg(&num_q_i32)
            .arg(&batch_i32)
            .launch(cfg)
            .map_err(|e| MLError::ModelError(format!("tft_quantile_extract launch: {e}")))?;
    }

    Ok(out)
}

/// Compute walk-forward backtest fitness for hyperopt (minimization).
///
/// - `trades >= min_trades`: return `-sharpe` (lower is better for minimizers)
/// - `0 < trades < min_trades`: return `-sharpe * (trades / min_trades)` (penalty)
/// - `trades == 0`: return `1.0 / (1.0 + threshold_distance)` (always positive,
///   guides the optimizer toward viable threshold regions)
pub fn backtest_fitness(
    sharpe: f32,
    total_trades: u32,
    min_trades: u32,
    threshold_distance: f32,
) -> f64 {
    if total_trades == 0 {
        return 1.0 / (1.0 + threshold_distance as f64);
    }
    let s = -(sharpe as f64);
    if total_trades >= min_trades {
        s
    } else {
        s * (total_trades as f64 / min_trades as f64)
    }
}

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

    fn cuda_stream() -> Arc<CudaStream> {
        let context = cudarc::driver::CudaContext::new(0).expect("CUDA context required");
        context.default_stream()
    }

    // ── ppo_to_exposure_scores ──────────────────────────────────────────

    #[test]
    fn test_ppo_to_exposure_scores_shape() {
        let stream = cuda_stream();
        let batch = 4;
        let uniform = vec![1.0_f32 / 45.0; batch * 45];
        let mut probs_buf = stream.alloc_zeros::<f32>(batch * 45).unwrap();
        stream.memcpy_htod(&uniform, &mut probs_buf).unwrap();

        let scores = ppo_to_exposure_scores(&probs_buf, batch, &stream).unwrap();

        // Download and check shape + values
        let mut host = vec![0.0_f32; batch * 5];
        stream.memcpy_dtoh(&scores, &mut host).unwrap(); // test readback
        stream.synchronize().unwrap();

        assert_eq!(host.len(), batch * 5);
        // Each of 5 bins sums 9 cells of 1/45 = 0.2
        for val in &host {
            assert!(
                (*val - 0.2).abs() < 1e-5,
                "expected ~0.2 everywhere, got {val}"
            );
        }
    }

    #[test]
    fn test_ppo_to_exposure_scores_argmax() {
        let stream = cuda_stream();
        // Put all mass on actions 36..44 => exposure bin 4 (Long100)
        let batch = 2;
        let mut raw = vec![0.0_f32; batch * 45];
        for b in 0..batch {
            for a in 36..45 {
                if let Some(slot) = raw.get_mut(b * 45 + a) {
                    *slot = 1.0 / 9.0;
                }
            }
        }
        let mut probs_buf = stream.alloc_zeros::<f32>(batch * 45).unwrap();
        stream.memcpy_htod(&raw, &mut probs_buf).unwrap();

        let scores = ppo_to_exposure_scores(&probs_buf, batch, &stream).unwrap();
        let mut host = vec![0.0_f32; batch * 5];
        stream.memcpy_dtoh(&scores, &mut host).unwrap(); // test readback
        stream.synchronize().unwrap();

        // Each row should have argmax at index 4 (Long100)
        for b in 0..batch {
            let row_start = b * 5;
            let mut best_idx = 0;
            let mut best_val = f32::NEG_INFINITY;
            for i in 0..5 {
                let v = host[row_start + i];
                if v > best_val {
                    best_val = v;
                    best_idx = i;
                }
            }
            assert_eq!(best_idx, 4, "batch {b}: expected argmax=4, got {best_idx}");
        }
    }

    // ── signal_to_action_scores ─────────────────────────────────────────

    /// Host-side argmax over a [1, 5] row downloaded from GPU.
    fn host_argmax(host: &[f32]) -> usize {
        let mut best_idx = 0;
        let mut best_val = f32::NEG_INFINITY;
        for (i, &v) in host.iter().enumerate() {
            if v > best_val {
                best_val = v;
                best_idx = i;
            }
        }
        best_idx
    }

    fn run_signal_test(pred_val: f32, high: f32, low: f32) -> Vec<f32> {
        let stream = cuda_stream();
        let mut pred_buf = stream.alloc_zeros::<f32>(1).unwrap();
        stream.memcpy_htod(&[pred_val], &mut pred_buf).unwrap();
        let scores = signal_to_action_scores(&pred_buf, 1, high, low, &stream).unwrap();
        let mut host = vec![0.0_f32; 5];
        stream.memcpy_dtoh(&scores, &mut host).unwrap(); // test readback
        stream.synchronize().unwrap();
        host
    }

    #[test]
    fn test_signal_to_action_scores_strong_long() {
        let host = run_signal_test(20.0, 10.0, 5.0);
        assert_eq!(host_argmax(&host), 4); // Long100
    }

    #[test]
    fn test_signal_to_action_scores_strong_short() {
        let host = run_signal_test(-20.0, 10.0, 5.0);
        assert_eq!(host_argmax(&host), 0); // Short100
    }

    #[test]
    fn test_signal_to_action_scores_flat() {
        let host = run_signal_test(0.0, 10.0, 5.0);
        assert_eq!(host_argmax(&host), 2); // Flat
    }

    #[test]
    fn test_signal_to_action_scores_mild_long() {
        let host = run_signal_test(7.0, 10.0, 5.0);
        assert_eq!(host_argmax(&host), 3); // Long50
    }

    #[test]
    fn test_signal_to_action_scores_mild_short() {
        let host = run_signal_test(-7.0, 10.0, 5.0);
        assert_eq!(host_argmax(&host), 1); // Short50
    }

    // ── tft_quantile_to_signal ──────────────────────────────────────────

    #[test]
    fn test_tft_quantile_to_signal() {
        let stream = cuda_stream();
        // [batch=2, horizon=1, quantiles=3]
        let data: Vec<f32> = vec![
            -1.0, 0.5, 2.0,   // batch 0: median = 0.5
            -3.0, -1.5, 0.0,  // batch 1: median = -1.5
        ];
        let mut q_buf = stream.alloc_zeros::<f32>(data.len()).unwrap();
        stream.memcpy_htod(&data, &mut q_buf).unwrap();

        let signal = tft_quantile_to_signal(&q_buf, 2, 1, 3, &stream).unwrap();
        let mut host = vec![0.0_f32; 2];
        stream.memcpy_dtoh(&signal, &mut host).unwrap(); // test readback
        stream.synchronize().unwrap();

        assert!(
            (host[0] - 0.5).abs() < 1e-6,
            "batch 0: expected 0.5, got {}",
            host[0]
        );
        assert!(
            (host[1] - (-1.5)).abs() < 1e-6,
            "batch 1: expected -1.5, got {}",
            host[1]
        );
    }

    #[test]
    fn test_tft_quantile_rejects_bad_dims() {
        let stream = cuda_stream();
        let mut q_buf = stream.alloc_zeros::<f32>(6).unwrap();
        stream.memcpy_htod(&[0.0_f32; 6], &mut q_buf).unwrap();

        // horizon=0 should fail
        let err = tft_quantile_to_signal(&q_buf, 2, 0, 3, &stream);
        assert!(err.is_err(), "should reject horizon=0");

        // num_quantiles=1 should fail
        let err = tft_quantile_to_signal(&q_buf, 2, 1, 1, &stream);
        assert!(err.is_err(), "should reject num_quantiles=1");
    }

    // ── backtest_fitness ────────────────────────────────────────────────

    #[test]
    fn test_backtest_fitness_normal_trades() {
        let f = backtest_fitness(1.5, 50, 30, 0.0);
        assert!((f - (-1.5)).abs() < 1e-9);
    }

    #[test]
    fn test_backtest_fitness_few_trades() {
        let f = backtest_fitness(1.5, 15, 30, 0.0);
        // -1.5 * (15/30) = -0.75
        assert!((f - (-0.75)).abs() < 1e-9);
    }

    #[test]
    fn test_backtest_fitness_zero_trades() {
        let f = backtest_fitness(1.5, 0, 30, 2.0);
        // 1 / (1 + 2) = 1/3
        assert!((f - (1.0 / 3.0)).abs() < 1e-9);
    }

    #[test]
    fn test_backtest_fitness_zero_trades_sane_thresholds() {
        let f = backtest_fitness(1.5, 0, 30, 0.0);
        // 1 / (1 + 0) = 1.0
        assert!((f - 1.0).abs() < 1e-9);
    }
}
