//! Graph unary primitives: ops.cpp and metal/unary.cpp from the pinned runtime.
use super::array::{Array, Dtype, Layout};
use super::stream::{Device, Stream};
use super::{bytes, dispatch_geometry, ops, tensor, views};

#[derive(Clone, Copy, Debug)]
pub(super) enum Unary {
    LogicalNot,
    Abs,
    Sign,
    Negative,
    #[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))]
    Square,
    Exp,
    #[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))]
    Log,
    Sigmoid,
    Tanh,
    Sqrt,
    Rsqrt,
    Sin,
    Cos,
    Erf,
}

pub(super) fn unary(input: &Array, op: Unary, stream: Stream) -> Result<Array, String> {
    let dtype = input.layout().dtype();
    if stream.device() != Device::Gpu {
        return Err("CPU unary arithmetic is not ported".into());
    }
    if matches!(
        (op, dtype),
        (Unary::Negative, Dtype::Bool) | (Unary::Sigmoid, Dtype::C64)
    ) {
        return Err(format!("{op:?} does not support {dtype:?}"));
    }
    let dtype = if matches!(op, Unary::LogicalNot) {
        Dtype::Bool
    } else if matches!(
        op,
        Unary::Negative | Unary::Square | Unary::Abs | Unary::Sign
    ) {
        dtype
    } else {
        dtype.at_least_float()
    };
    let input = ops::astype(input, dtype, false, stream)?;
    let output = Array::make_operation(
        stream,
        &input,
        Layout::new(input.layout().shape(), dtype)?,
        ops::Operation::Unary(op),
    )?;
    if matches!(op, Unary::Abs) && dtype == Dtype::C64 {
        ops::astype(&output, Dtype::F32, false, stream)
    } else {
        Ok(output)
    }
}

pub(super) fn evaluate(input: &Array, output: &Array, op: Unary) -> Result<(), String> {
    views::set_unary_output_data(input, output)?;
    let src = input.layout();
    let dst = output.layout();
    if src.size() == 0 {
        return Ok(());
    }
    let contiguous = src.flags().contiguous;
    let size = input.data_size();
    let large = if contiguous {
        size > u32::MAX as usize
    } else {
        size > i32::MAX as usize || dst.size() > i32::MAX as usize
    };
    let work = if contiguous {
        if size < 1 << 16 {
            1
        } else {
            (8 / src.dtype().itemsize()).max(1)
        }
    } else if large {
        4
    } else {
        1
    };
    let prefix = if contiguous {
        if large {
            "v2"
        } else if work > 1 {
            "vn"
        } else {
            "v"
        }
    } else if large {
        "gn4large"
    } else {
        "gn1"
    };
    let name = format!(
        "{prefix}_{op:?}{}{}",
        src.dtype().kernel_name(),
        dst.dtype().kernel_name()
    );
    let max_threads = views::kernel_max_threads(&name)?;
    let source = input.buffer();
    let target = output.buffer();
    let mut bindings = vec![
        tensor(0, &source)
            .at_byte_offset(input.offset())
            .input(size),
        tensor(1, &target)
            .at_byte_offset(output.offset())
            .output(output.data_size()),
    ];
    if contiguous {
        let n32 = size as i32;
        let n64 = size as i64;
        bindings.push(if large {
            bytes(2, &n64)
        } else {
            bytes(2, &n32)
        });
        let grid = if large {
            views::large_copy_grid(&dst, work)?
        } else {
            [size.div_ceil(work) as u32, 1, 1]
        };
        let group = [(max_threads as usize).min(size.div_ceil(work)) as u32, 1, 1];
        dispatch_geometry(&name, &bindings, &[], grid, group, true)
    } else {
        let (shape, strides) = views::collapse_dims(&src);
        let ndim = i32::try_from(shape.len()).map_err(|_| "unary rank overflow")?;
        let dim0 = shape.last().copied().unwrap_or(1) as usize;
        let dim1 = shape.iter().rev().nth(1).copied().unwrap_or(1) as usize;
        let rest = dst.size() / (dim0 * dim1);
        let grid = [dim0.div_ceil(work), dim1, rest]
            .map(|d| u32::try_from(d).map_err(|_| "unary grid overflow"));
        let grid = [grid[0]?, grid[1]?, grid[2]?];
        if max_threads != 1024 {
            return Err("[Metal::unary] Must use 1024 sized block".into());
        }
        bindings.extend([
            bytes(2, shape.as_slice()),
            bytes(3, strides.as_slice()),
            bytes(4, &ndim),
        ]);
        dispatch_geometry(&name, &bindings, &[], grid, super::block_dims(grid), true)
    }
}

#[test]
#[ignore = "requires Apple Metal; graph unary layout, donation and kernel boundaries"]
fn mtplx_graph_unary_preserves_reference_layout_and_values() {
    super::eval::execution_tests::observed(|| {
        use super::super::gpu::{Buffer, Context};
        use super::{allocator, configure_sources, stream::Streams};
        configure_sources().unwrap();
        let _context = Context::open_qwen(0).unwrap();
        let streams = Streams::new(Some(allocator::Allocator::new().unwrap()));
        let gpu = streams.default_stream(Device::Gpu).unwrap();
        // Exact dyadic values only; no host transcendental approximation.
        let encode = |dtype: Dtype, values: &[f32]| -> Vec<u8> {
            values
                .iter()
                .flat_map(|&v| {
                    let f = v.to_bits();
                    let bits = match dtype {
                        Dtype::BF16 => f >> 16,
                        Dtype::F16 => {
                            (f >> 16 & 0x8000)
                                | if f & 0x7fffffff == 0 {
                                    0
                                } else {
                                    (((f >> 23 & 255) - 112) << 10) | (f >> 13 & 1023)
                                }
                        }
                        Dtype::I32 => v as i32 as u32,
                        _ => f,
                    };
                    bits.to_le_bytes()[..dtype.itemsize()].to_vec()
                })
                .collect()
        };
        for dtype in [Dtype::F16, Dtype::BF16, Dtype::F32] {
            for (op, value, expected) in [
                (Unary::Negative, 4., -4.),
                (Unary::Abs, -4., 4.),
                (Unary::Sign, -4., -1.),
                (Unary::Square, 4., 16.),
                (Unary::Exp, 0., 1.),
                (Unary::Sigmoid, 0., 0.5),
                (Unary::Tanh, 0., 0.),
                (Unary::Sqrt, 4., 2.),
                (Unary::Rsqrt, 4., 0.5),
            ] {
                for count in [0, 1, 65535, 65536, 65537] {
                    for keep_alias in [false, true] {
                        let data = encode(dtype, &vec![value; count]);
                        let buffer = Buffer::mtplx_bytes(data.len() as u64).unwrap();
                        buffer.write(0, &data).unwrap();
                        let input = Array::new(&[count as i32], dtype, buffer).unwrap();
                        let identity = input.buffer().data_identity();
                        let out = unary(&input, op, gpu).unwrap();
                        let alias = keep_alias.then(|| input.clone());
                        drop(input);
                        ops::evaluate(&streams, std::slice::from_ref(&out), gpu, false).unwrap();
                        assert_eq!(out.buffer().data_identity() == identity, !keep_alias);
                        let mut actual = vec![0; count * dtype.itemsize()];
                        out.buffer().read(out.offset(), &mut actual).unwrap();
                        assert_eq!(
                            actual,
                            encode(dtype, &vec![expected; count]),
                            "{op:?} {dtype:?}"
                        );
                        drop(alias);
                    }
                }
            }
            for mode in 0..5 {
                let values = [1., 4., 1., 4., 1., 4., 1., 4.];
                let bytes = encode(dtype, &values);
                let buffer = Buffer::mtplx_bytes(bytes.len() as u64).unwrap();
                buffer.write(0, &bytes).unwrap();
                let base = Array::new(&[2, 4], dtype, buffer).unwrap();
                let view = match mode {
                    0 => ops::transpose(&base, &[1, 0], gpu).unwrap(),
                    1 => ops::slice(&base, &[0, 1], &[2, 4], &[1, 2], gpu).unwrap(),
                    2 => ops::slice(&base, &[1, 3], &[-3, -5], &[-1, -1], gpu).unwrap(),
                    3 => {
                        let row = ops::slice(&base, &[1, 0], &[2, 4], &[1, 1], gpu).unwrap();
                        ops::broadcast_to(&row, &[3, 4], gpu).unwrap()
                    }
                    _ => ops::slice(&base, &[1, 0], &[2, 4], &[1, 1], gpu).unwrap(),
                };
                let identity = base.buffer().data_identity();
                let out = unary(&view, Unary::Square, gpu).unwrap();
                drop((view, base));
                ops::evaluate(&streams, std::slice::from_ref(&out), gpu, false).unwrap();
                let expected: &[f32] = match mode {
                    0 => &[1., 16., 1., 16., 1., 16., 1., 16.],
                    1 => &[16., 16., 16., 16.],
                    2 => &[16., 1., 16., 1., 16., 1., 16., 1.],
                    _ => &[1., 16., 1., 16.],
                };
                let mut actual = vec![0; out.data_size() * dtype.itemsize()];
                out.buffer().read(out.offset(), &mut actual).unwrap();
                assert_eq!(actual, encode(dtype, expected), "layout {mode}");
                // check_contiguity counts positive strides only: reversed
                // and stepped inputs both require a fresh dense output.
                assert_eq!(
                    out.buffer().data_identity() == identity,
                    !matches!(mode, 1 | 2)
                );
                if mode == 3 {
                    assert_eq!(out.layout().strides(), [0, 1]);
                }
            }
            let oversized = Buffer::mtplx_bytes(20000).unwrap();
            oversized.write(0, &encode(dtype, &[1., 4.])).unwrap();
            let identity = oversized.data_identity();
            let input = Array::new(&[2], dtype, oversized).unwrap();
            let out = unary(&input, Unary::Square, gpu).unwrap();
            drop(input);
            ops::evaluate(&streams, std::slice::from_ref(&out), gpu, false).unwrap();
            assert_ne!(
                out.buffer().data_identity(),
                identity,
                "16KiB donation slack"
            );
            let mut actual = vec![0; 2 * dtype.itemsize()];
            out.buffer().read(0, &mut actual).unwrap();
            assert_eq!(actual, encode(dtype, &[1., 16.]));
        }
        let buffer = Buffer::mtplx_bytes(4).unwrap();
        buffer.write(0, &4_i32.to_le_bytes()).unwrap();
        let input = Array::new(&[], Dtype::I32, buffer).unwrap();
        let out = unary(&input, Unary::Sqrt, gpu).unwrap();
        ops::evaluate(&streams, std::slice::from_ref(&out), gpu, false).unwrap();
        assert_eq!(out.layout().dtype(), Dtype::F32);
        let mut actual = [0; 4];
        out.buffer().read(0, &mut actual).unwrap();
        assert_eq!(actual, 2_f32.to_le_bytes());
        let squared = unary(&input, Unary::Square, gpu).unwrap();
        ops::evaluate(&streams, std::slice::from_ref(&squared), gpu, false).unwrap();
        assert_eq!(squared.layout().dtype(), Dtype::I32);
        squared.buffer().read(0, &mut actual).unwrap();
        assert_eq!(actual, 16_i32.to_le_bytes());
        for (dtype, op) in [(Dtype::Bool, Unary::Negative), (Dtype::C64, Unary::Sigmoid)] {
            let input = Array::new(&[0], dtype, Buffer::mtplx_bytes(0).unwrap()).unwrap();
            assert!(unary(&input, op, gpu).is_err());
        }
        streams.clear_streams().unwrap();
    });
}