//! Qwen SparseMoeBlock as ordinary Array nodes, with the original two
//! CustomKernel nodes only where MTPLX selects fused decode/verify.
use super::array::{Array, Dtype, Layout};
use super::binary::{Binary, binary};
use super::mlp::{self, QuantizedLinear};
use super::stream::{Device, Stream, Streams};
use super::{indexing, linear, normalization, ops, reduce, sort, unary};

pub(super) enum GateUp {
    Fused(QuantizedLinear),
    Separate {
        gate: QuantizedLinear,
        up: QuantizedLinear,
    },
}

pub(super) struct FeedForward {
    pub(super) input: GateUp,
    pub(super) down: QuantizedLinear,
}

fn gather(
    input: &Array,
    layer: &QuantizedLinear,
    ids: &Array,
    sorted: bool,
    streams: &Streams,
    stream: Stream,
) -> Result<Array, String> {
    linear::guarded_gather_qmm(
        input,
        [&layer.weight, &layer.scales, &layer.biases],
        [None, Some(ids)],
        [layer.bits, layer.group],
        sorted,
        streams,
        stream,
    )
}

impl FeedForward {
    fn shared(&self, input: &Array, stream: Stream) -> Result<Array, String> {
        match &self.input {
            GateUp::Fused(gu) => mlp::fused_gate_up_mlp(input, gu, &self.down, stream),
            GateUp::Separate { gate, up } => {
                // Qwen3NextMLP evaluates gate first, unlike SwitchGLU below.
                let gate = gate.apply(input, stream)?;
                let up = up.apply(input, stream)?;
                self.down.apply(&mlp::swiglu(&gate, &up, stream)?, stream)
            }
        }
    }

    fn switched(
        &self,
        input: &Array,
        indices: &Array,
        sorted_gather: bool,
        streams: &Streams,
        stream: Stream,
    ) -> Result<Array, String> {
        let do_sort = indices.layout().size() >= 64;
        let mut x = if do_sort {
            input.clone()
        } else {
            ops::expand_dims(input, &[-2, -3], stream)?
        };
        let mut ids = indices.clone();
        let mut inverse = None;
        let mut row_map = None;
        if do_sort {
            let size =
                i32::try_from(indices.layout().size()).map_err(|_| "MoE index size overflow")?;
            ids = ops::reshape(&ids, &[size], stream)?;
            let order = sort::sort(&ids, Some(-1), true, None, stream)?;
            inverse = Some(sort::sort(&order, Some(-1), true, None, stream)?);
            let topk = *indices.layout().shape().last().ok_or("MoE indices rank")? as u32;
            let buffer = super::super::gpu::Buffer::mtplx_bytes(4)?;
            buffer.write(0, &topk.to_le_bytes())?;
            let divisor = Array::new(&[], Dtype::U32, buffer)?;
            let rows = super::binary::floor_divide_integer(&order, &divisor, stream)?;
            let k = *input.layout().shape().last().ok_or("MoE input rank")?;
            let flat_rows = i32::try_from(input.layout().size() / k as usize)
                .map_err(|_| "MoE row count overflow")?;
            x = ops::reshape(&x, &[flat_rows, 1, k], stream)?;
            row_map = Some(rows);
            ids = indexing::take(&ids, &order, 0, stream)?;
        }
        let hidden = if let Some(map) = row_map {
            let (gate, up) = match &self.input {
                GateUp::Fused(gu) => (gu, None),
                GateUp::Separate { gate, up } => (gate, Some(up)),
            };
            super::moe_sorted::swiglu_rows(
                &x,
                &map,
                gate,
                up,
                &ids,
                sorted_gather,
                streams,
                stream,
            )?
        } else {
            match &self.input {
                GateUp::Fused(gu) => {
                    let gu = gather(&x, gu, &ids, do_sort, streams, stream)?;
                    let n = *gu.layout().shape().last().ok_or("MoE gate/up rank")?;
                    if n % 2 != 0 {
                        return Err("MoE fused gate/up width must be even".into());
                    }
                    let pair = ops::split(&gu, &[n / 2], -1, stream)?;
                    binary(
                        &mlp::silu(&pair[0], stream)?,
                        &pair[1],
                        Binary::Multiply,
                        stream,
                    )?
                }
                GateUp::Separate { gate, up } => {
                    let up = gather(&x, up, &ids, do_sort, streams, stream)?;
                    let gate = gather(&x, gate, &ids, do_sort, streams, stream)?;
                    mlp::swiglu(&gate, &up, stream)?
                }
            }
        };
        let mut output = gather(&hidden, &self.down, &ids, do_sort, streams, stream)?;
        if let Some(inverse) = inverse {
            output = indexing::take(&output, &inverse, 0, stream)?;
            let mut shape = indices.layout().shape().to_vec();
            shape.extend_from_slice(&output.layout().shape()[1..]);
            output = ops::reshape(&output, &shape, stream)?;
        }
        ops::squeeze(&output, Some(&[-2]), stream)
    }
}

pub(super) struct SparseMoe {
    pub(super) experts: FeedForward,
    pub(super) shared: FeedForward,
    pub(super) router: QuantizedLinear,
    pub(super) shared_gate: QuantizedLinear,
    pub(super) topk: i32,
    pub(super) normalize: bool,
    pub(super) sorted_gather: bool,
    pub(super) fused_decode: bool,
    pub(super) fused_verify: bool,
}

impl SparseMoe {
    fn fused_rows(&self, input: &Array) -> Option<u32> {
        let GateUp::Fused(gu) = &self.experts.input else {
            return None;
        };
        if gu.bits != 4
            || !matches!(gu.group, 32 | 64)
            || self.experts.down.bits != 4
            || !matches!(self.experts.down.group, 32 | 64)
        {
            return None;
        }
        let l = input.layout();
        let shape = l.shape();
        let rows = shape[shape.len() - 2];
        let k = shape[shape.len() - 1];
        if l.size() != rows as usize * k as usize {
            return None;
        }
        ((rows == 1 && self.fused_decode) || ((2..=4).contains(&rows) && self.fused_verify))
            .then_some(rows as u32)
    }

    pub(super) fn apply(
        &self,
        input: &Array,
        streams: &Streams,
        stream: Stream,
    ) -> Result<Array, String> {
        if stream.device() != Device::Gpu
            || input.layout().dtype() != Dtype::BF16
            || input.layout().shape().len() < 2
            || input.layout().size() == 0
        {
            return Err("SparseMoe requires nonempty BF16 GPU rows".into());
        }
        let fused = self.fused_rows(input);
        let gates = normalization::softmax_last(&self.router.apply(input, stream)?, true, stream)?;
        let shape = gates.layout().shape().to_vec();
        let rank = shape.len();
        let experts = shape[rank - 1];
        if self.topk < 1 || self.topk > experts {
            return Err("invalid MoE top-k".into());
        }
        let partition = sort::sort(&gates, Some(-1), true, Some(-self.topk), stream)?;
        let mut starts = vec![0; rank];
        starts[rank - 1] = experts - self.topk;
        let indices = ops::slice(&partition, &starts, &shape, &vec![1; rank], stream)?;
        let mut scores = indexing::take_along_axis(&gates, &indices, -1, stream)?;
        if self.normalize {
            scores = binary(
                &scores,
                &reduce::sum(&scores, &[-1], true, stream)?,
                Binary::Divide,
                stream,
            )?;
        }
        let y = if let Some(rows) = fused {
            let GateUp::Fused(gu) = &self.experts.input else {
                unreachable!()
            };
            fused_experts(
                input,
                gu,
                &self.experts.down,
                &indices,
                &scores,
                rows,
                self.topk,
                stream,
            )?
        } else {
            let switched =
                self.experts
                    .switched(input, &indices, self.sorted_gather, streams, stream)?;
            let scores = ops::expand_dims(&scores, &[-1], stream)?;
            reduce::sum(
                &binary(&switched, &scores, Binary::Multiply, stream)?,
                &[-2],
                false,
                stream,
            )?
        };
        let gate = || {
            unary::unary(
                &self.shared_gate.apply(input, stream)?,
                unary::Unary::Sigmoid,
                stream,
            )
        };
        let (gate, shared) = if fused.is_some() {
            (gate()?, self.shared.shared(input, stream)?)
        } else {
            let shared = self.shared.shared(input, stream)?;
            (gate()?, shared)
        };
        let y = binary(
            &y,
            &binary(&gate, &shared, Binary::Multiply, stream)?,
            Binary::Add,
            stream,
        )?;
        if fused.is_some() {
            ops::astype(&y, input.layout().dtype(), false, stream)
        } else {
            Ok(y)
        }
    }
}

#[allow(clippy::too_many_arguments)] // Original two CustomKernel call inputs and template dimensions.
fn fused_experts(
    input: &Array,
    gu: &QuantizedLinear,
    down: &QuantizedLinear,
    indices: &Array,
    scores: &Array,
    rows: u32,
    topk: i32,
    stream: Stream,
) -> Result<Array, String> {
    if topk != 10
        || gu.weight.layout().shape().get(1) != Some(&1280)
        || down.weight.layout().shape().get(1) != Some(&2560)
        || input.layout().shape().last() != Some(&2560)
    {
        return Err("fused MoE exports require the installed top10/2560/640 shape".into());
    }
    let flatten = |a: &Array| ops::reshape(a, &[a.layout().size() as i32], stream);
    let x = if rows == 1 {
        flatten(input)?
    } else {
        ops::reshape(input, &[rows as i32, 2560], stream)?
    };
    let ids = ops::astype(&flatten(indices)?, Dtype::U32, false, stream)?;
    let scores = ops::astype(&flatten(scores)?, Dtype::F32, false, stream)?;
    let hshape = if rows == 1 {
        vec![10, 640]
    } else {
        vec![rows as i32, 10, 640]
    };
    let h = Array::make_operation_with_inputs(
        stream,
        &[
            x,
            gu.weight.clone(),
            gu.scales.clone(),
            gu.biases.clone(),
            ids.clone(),
        ],
        Layout::new(&hshape, Dtype::BF16)?,
        ops::Operation::MoeGluH {
            rows,
            group: gu.group,
        },
    )?;
    let yshape = if rows == 1 {
        vec![2560]
    } else {
        vec![rows as i32, 2560]
    };
    let y = Array::make_operation_with_inputs(
        stream,
        &[
            h,
            down.weight.clone(),
            down.scales.clone(),
            down.biases.clone(),
            ids,
            scores,
        ],
        Layout::new(&yshape, Dtype::BF16)?,
        ops::Operation::MoeDownY {
            rows,
            group: down.group,
        },
    )?;
    ops::reshape(&y, input.layout().shape(), stream)
}

pub(super) fn fused_evaluate(
    inputs: &[Array],
    outputs: &[Array],
    rows: u32,
    group: u32,
    down: bool,
) -> Result<(), String> {
    if !(1..=4).contains(&rows) || !matches!(group, 32 | 64) {
        return Err("invalid fused MoE specialization".into());
    }
    let suffix = if rows == 1 {
        format!("g{group}_bf16")
    } else {
        format!("m_m{rows}_g{group}_bf16")
    };
    let name = format!(
        "kernel_qwen_mtplx_moe_{}_{suffix}",
        if down { "down_y" } else { "glu_h" }
    );
    let count = rows * if down { 2560 } else { 10 * 640 };
    super::gdn::custom_evaluate(inputs, outputs, |bindings| {
        super::dispatch_geometry(
            &name,
            bindings,
            &[],
            [count.div_ceil(32) * 1024, 1, 1],
            [1024, 1, 1],
            true,
        )
    })
}

#[test]
#[ignore = "requires Apple M5 Max; full MoE Array graph, no model weights; use test-supervisor"]
fn mtplx_graph_sparse_moe_matches_reference() {
    use super::super::{
        gpu::Context,
        qwen_mtplx_tests::{capture_dispatches, pattern, verify},
    };
    use super::{allocator, configure_sources};
    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();
    let fixtures = include_str!("../../../../tests/fixtures/mtplx-custom-kernels.jsonl")
        .lines()
        .map(|line| serde_json::from_str::<serde_json::Value>(line).unwrap())
        .filter(|row| row["kernel"].as_str().unwrap().starts_with("moe_"))
        .collect::<Vec<_>>();
    assert_eq!(fixtures.len(), 220);
    for prepared in [false, true] {
        for row in &fixtures {
            let [tokens, experts, k, intermediate, gu_group, down_group] = [
                "tokens",
                "experts",
                "k",
                "intermediate",
                "gu_group",
                "down_group",
            ]
            .map(|key| row[key].as_u64().unwrap() as u32);
            eprintln!("MoE Array graph prepared={prepared} {}", row["kernel"]);
            let pack = |e: u32, n: u32, k: u32, group: u32, salt: u32, expert: bool| {
                let mut q = mlp::fixture_linear(e * n, k, 4, group, salt);
                if expert {
                    for a in [&mut q.weight, &mut q.scales, &mut q.biases] {
                        let width = *a.layout().shape().last().unwrap();
                        *a = ops::reshape(a, &[e as i32, n as i32, width], gpu).unwrap();
                    }
                }
                q
            };
            let ff = |expert: bool, separate: bool| {
                let e = if expert { experts } else { 1 };
                let salt = if expert { 39 } else { 46 };
                FeedForward {
                    input: if separate {
                        GateUp::Separate {
                            gate: pack(e, intermediate, k, gu_group, salt, expert),
                            up: pack(
                                e,
                                intermediate,
                                k,
                                down_group,
                                if expert { 45 } else { 54 },
                                expert,
                            ),
                        }
                    } else {
                        GateUp::Fused(pack(e, 2 * intermediate, k, gu_group, salt, expert))
                    },
                    down: pack(
                        e,
                        k,
                        intermediate,
                        down_group,
                        if expert { 41 } else { 48 },
                        expert,
                    ),
                }
            };
            let model = SparseMoe {
                experts: ff(true, row["separate_experts"].as_bool().unwrap_or(false)),
                shared: ff(false, row["separate_shared"].as_bool().unwrap_or(false)),
                router: mlp::fixture_linear(experts, k, 8, 64, 50),
                shared_gate: mlp::fixture_linear(1, k, 4, 64, 52),
                topk: 10,
                normalize: row["normalize"].as_bool().unwrap(),
                sorted_gather: true,
                fused_decode: row["fused"].as_bool().unwrap_or(false),
                fused_verify: row["fused"].as_bool().unwrap_or(false),
            };
            if prepared {
                let mut parameters = Vec::new();
                let mut add = |q: &QuantizedLinear| {
                    parameters.extend([q.weight.clone(), q.scales.clone(), q.biases.clone()]);
                };
                for ff in [&model.experts, &model.shared] {
                    match &ff.input {
                        GateUp::Fused(gu) => add(gu),
                        GateUp::Separate { gate, up } => {
                            add(gate);
                            add(up);
                        }
                    }
                    add(&ff.down);
                }
                add(&model.router);
                add(&model.shared_gate);
                ops::evaluate(&streams, &parameters, gpu, false).unwrap();
            }
            let input = Array::new(
                &[1, tokens as i32, k as i32],
                Dtype::BF16,
                pattern(tokens * k, 44),
            )
            .unwrap();
            let output = model.apply(&input, &streams, gpu).unwrap();
            let fused = row["fused"].as_bool().unwrap_or(false)
                && !row["separate_experts"].as_bool().unwrap_or(false)
                && tokens <= 4;
            assert_eq!(model.fused_rows(&input).is_some(), fused);
            // Graph edges, not an external model/parameter table, keep inputs alive.
            drop((input, model));
            let (_, dispatches) = capture_dispatches(|| {
                ops::evaluate(&streams, std::slice::from_ref(&output), gpu, false).unwrap();
            });
            let actual = dispatches
                .into_iter()
                .filter(|call| call.0.starts_with("kernel_qwen_mtplx_moe_"))
                .collect::<Vec<_>>();
            let expected = if fused {
                [
                    ("glu_h", gu_group, tokens * 10 * 640),
                    ("down_y", down_group, tokens * 2560),
                ]
                .map(|(name, group, rows)| {
                    let suffix = if tokens == 1 {
                        format!("g{group}_bf16")
                    } else {
                        format!("m_m{tokens}_g{group}_bf16")
                    };
                    (
                        format!("kernel_qwen_mtplx_moe_{name}_{suffix}"),
                        [rows.div_ceil(32) * 1024, 1, 1],
                        [1024, 1, 1],
                        true,
                    )
                })
                .to_vec()
            } else {
                Vec::new()
            };
            assert_eq!(
                actual, expected,
                "original CustomKernel dispatch geometry/order"
            );
            verify(
                row["kernel"].as_str().unwrap(),
                &[(&output.buffer(), tokens * k)],
            );
        }
    }
    streams.clear_streams().unwrap();
}