//! Multi-horizon sigmoid heads (128 → 5) and projection (128 → 8 + layer-norm).
//!
//! Per spec Section 2 + 2026-05-16 stacked amendment, the heads sit on
//! the CfC hidden output and emit 5 sigmoid probabilities of "up" at
//! horizons {30, 100, 300, 1000, 6000} snapshots forward.

use std::sync::Arc;

use anyhow::{Context, Result};
use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut, LaunchConfig, PushKernelArg};
use ml_core::device::MlDevice;

use crate::pinned_mem::MappedF32Buffer;

const HEADS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/multi_horizon_heads.cubin"));
const PROJ_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/projection.cubin"));

pub const N_HORIZONS: usize = 5;
pub const HORIZONS: [usize; N_HORIZONS] = [30, 100, 300, 1000, 6000];
pub const HIDDEN_DIM: usize = 128;
pub const PROJ_DIM: usize = 8;
/// GRN intermediate dimension (Phase 1.7). Two linear layers + GLU gate
/// + skip projection all operate at this width per horizon. Matches the
/// HEAD_MID_H define in `cuda/multi_horizon_heads.cu`.
pub const HEAD_MID_DIM: usize = 64;

#[derive(Clone, Debug)]
pub struct HeadsWeights {
    pub w: Vec<f32>, // [5, 128] row-major
    pub b: Vec<f32>, // [5]
}

#[derive(Clone, Debug)]
pub struct ProjectionWeights {
    pub w: Vec<f32>,        // [8, 128]
    pub b: Vec<f32>,        // [8]
    pub ln_gain: Vec<f32>,  // [8]
    pub ln_bias: Vec<f32>,  // [8]
}

/// Placeholder marker — concrete trunk integration in Task 9.
pub struct MultiHorizonHeads;

pub struct Projection;

/// Backward through `multi_horizon_heads`. Returns (grad_w, grad_b, grad_h).
pub fn multi_horizon_heads_backward_gpu(
    dev: &MlDevice,
    w: &HeadsWeights,
    probs: &[f32; N_HORIZONS],
    h: &[f32],
    grad_probs: &[f32; N_HORIZONS],
) -> Result<(Vec<f32>, Vec<f32>, Vec<f32>)> {
    assert_eq!(h.len(), HIDDEN_DIM);

    let stream: &Arc<CudaStream> = dev.cuda_stream().context("heads-bwd stream")?;
    let ctx = dev.cuda_context().context("heads-bwd ctx")?;
    let module = ctx.load_cubin(HEADS_CUBIN.to_vec()).context("load heads cubin")?;
    let func = module.load_function("multi_horizon_heads_backward").context("heads_bwd fn")?;

    let w_d = upload(stream, &w.w)?;
    let probs_d = upload(stream, probs.as_slice())?;
    let h_d = upload(stream, h)?;
    let grad_p_d = upload(stream, grad_probs.as_slice())?;
    // Zero-filled carry buffer — semantically identical to passing
    // nullptr (kernel adds 0.0 in either case). Helper is for tests
    // that don't exercise the recurrent backward path.
    let zero_carry = stream.alloc_zeros::<f32>(HIDDEN_DIM).context("zero carry alloc")?;
    let mut grad_w_d = stream.alloc_zeros::<f32>(N_HORIZONS * HIDDEN_DIM).context("grad_w alloc")?;
    let mut grad_b_d = stream.alloc_zeros::<f32>(N_HORIZONS).context("grad_b alloc")?;
    let mut grad_h_d = stream.alloc_zeros::<f32>(HIDDEN_DIM).context("grad_h alloc")?;

    let cfg = LaunchConfig {
        grid_dim: (1, 1, 1),
        block_dim: (HIDDEN_DIM as u32, 1, 1),
        shared_mem_bytes: 0,
    };
    let mut launch = stream.launch_builder(&func);
    launch
        .arg(&w_d).arg(&probs_d).arg(&h_d).arg(&grad_p_d).arg(&zero_carry)
        .arg(&mut grad_w_d).arg(&mut grad_b_d).arg(&mut grad_h_d);
    unsafe { launch.launch(cfg).context("heads_bwd launch")?; }

    Ok((
        download(stream, &grad_w_d)?,
        download(stream, &grad_b_d)?,
        download(stream, &grad_h_d)?,
    ))
}

pub fn projection_gpu(
    dev: &MlDevice,
    w: &ProjectionWeights,
    h: &[f32],
) -> Result<[f32; PROJ_DIM]> {
    assert_eq!(h.len(), HIDDEN_DIM);
    assert_eq!(w.w.len(), PROJ_DIM * HIDDEN_DIM);
    assert_eq!(w.b.len(), PROJ_DIM);
    assert_eq!(w.ln_gain.len(), PROJ_DIM);
    assert_eq!(w.ln_bias.len(), PROJ_DIM);

    let stream: &Arc<CudaStream> = dev.cuda_stream().context("proj stream")?;
    let ctx = dev.cuda_context().context("proj ctx")?;
    let module = ctx.load_cubin(PROJ_CUBIN.to_vec()).context("load proj cubin")?;
    let func = module.load_function("projection_kernel").context("load function")?;

    let w_d = upload(stream, &w.w)?;
    let b_d = upload(stream, &w.b)?;
    let g_d = upload(stream, &w.ln_gain)?;
    let n_d = upload(stream, &w.ln_bias)?;
    let h_d = upload(stream, h)?;
    let mut out_d = stream.alloc_zeros::<f32>(PROJ_DIM).context("proj out alloc")?;

    let cfg = LaunchConfig {
        grid_dim: (1, 1, 1),
        block_dim: (PROJ_DIM as u32, 1, 1),
        shared_mem_bytes: 0,
    };
    let mut launch = stream.launch_builder(&func);
    launch
        .arg(&w_d)
        .arg(&b_d)
        .arg(&g_d)
        .arg(&n_d)
        .arg(&h_d)
        .arg(&mut out_d);
    unsafe { launch.launch(cfg).context("proj launch")?; }

    let v = download(stream, &out_d)?;
    let mut out = [0f32; PROJ_DIM];
    out.copy_from_slice(&v);
    Ok(out)
}

pub fn multi_horizon_heads_gpu(
    dev: &MlDevice,
    w: &HeadsWeights,
    h: &[f32],
) -> Result<[f32; N_HORIZONS]> {
    assert_eq!(h.len(), HIDDEN_DIM);
    assert_eq!(w.w.len(), N_HORIZONS * HIDDEN_DIM);
    assert_eq!(w.b.len(), N_HORIZONS);

    let stream: &Arc<CudaStream> = dev.cuda_stream().context("heads stream")?;
    let ctx = dev.cuda_context().context("heads ctx")?;
    let module = ctx.load_cubin(HEADS_CUBIN.to_vec()).context("load heads cubin")?;
    let func = module.load_function("multi_horizon_heads").context("load function")?;

    let w_d = upload(stream, &w.w)?;
    let b_d = upload(stream, &w.b)?;
    let h_d = upload(stream, h)?;
    let mut probs_d = stream.alloc_zeros::<f32>(N_HORIZONS).context("probs alloc")?;

    let cfg = LaunchConfig {
        grid_dim: (1, 1, 1),
        block_dim: (N_HORIZONS as u32, 1, 1),
        shared_mem_bytes: 0,
    };
    let mut launch = stream.launch_builder(&func);
    launch.arg(&w_d).arg(&b_d).arg(&h_d).arg(&mut probs_d);
    unsafe { launch.launch(cfg).context("heads launch")?; }

    let v = download(stream, &probs_d)?;
    let mut out = [0f32; N_HORIZONS];
    out.copy_from_slice(&v);
    Ok(out)
}

fn upload(stream: &Arc<CudaStream>, host: &[f32]) -> Result<CudaSlice<f32>> {
    let n = host.len();
    let staging = unsafe { MappedF32Buffer::new(n) }
        .map_err(|e| anyhow::anyhow!("heads upload staging: {e}"))?;
    staging.write_from_slice(host);
    let mut dst = stream.alloc_zeros::<f32>(n).context("heads upload alloc")?;
    if n > 0 {
        let nbytes = n * std::mem::size_of::<f32>();
        unsafe {
            let (dst_ptr, _g) = dst.device_ptr_mut(stream);
            cudarc::driver::result::memcpy_dtod_async(dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream())
                .context("heads upload DtoD")?;
        }
    }
    Ok(dst)
}

fn download(stream: &Arc<CudaStream>, src: &CudaSlice<f32>) -> Result<Vec<f32>> {
    let n = src.len();
    let staging = unsafe { MappedF32Buffer::new(n) }
        .map_err(|e| anyhow::anyhow!("heads download staging: {e}"))?;
    let nbytes = n * std::mem::size_of::<f32>();
    unsafe {
        let (src_ptr, _g) = src.device_ptr(stream);
        cudarc::driver::result::memcpy_dtod_async(staging.dev_ptr, src_ptr, nbytes, stream.cu_stream())
            .context("heads download DtoD")?;
    }
    stream.synchronize().context("heads download sync")?;
    Ok(staging.read_all())
}
