//! Qwen4ExpTextModel/TextModel graph orchestration from the pinned MTPLX.
//! GDN runs use the reference simplify/fuse passes. Synthetic run graphs are
//! checked; real-weight and serving integration remain outside this test port.
use super::array::{Array, Dtype};
use super::binary::{Binary, binary};
use super::decoder::{Attention, AttentionCache, Cache, DecoderLayer, Options};
use super::mlp::Linear;
use super::stream::{Stream, Streams};
use super::{cache::QsaCache, compiled::Replay, gdn, hyper, indexing, linear, ops, ple};
use std::collections::HashMap;
use std::ops::Range;

#[derive(Clone, Copy, Default)]
pub(super) struct CompilePolicy {
    explicit_off: bool,
    environment: bool,
}

impl CompilePolicy {
    pub(super) fn from_values(gdn: Option<&str>, family: Option<&str>) -> Self {
        let truthy = |s: &str| {
            matches!(
                s.trim().to_ascii_lowercase().as_str(),
                "1" | "true" | "yes" | "on"
            )
        };
        let explicit_off = [gdn, family].into_iter().flatten().any(|s| !truthy(s));
        Self {
            explicit_off,
            environment: !explicit_off && [gdn, family].into_iter().flatten().any(truthy),
        }
    }

    #[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))]
    pub(super) fn from_env() -> Self {
        Self::from_values(
            std::env::var("MTPLX_COMPILED_GDN").ok().as_deref(),
            std::env::var("MTPLX_QWEN4EXP_COMPILE").ok().as_deref(),
        )
    }

    fn enabled(self) -> bool {
        !self.explicit_off && self.environment
    }
}

#[derive(Debug, PartialEq, Eq)]
enum Run {
    Eager(usize),
    Gdn(Range<usize>),
}

fn runs(layers: impl IntoIterator<Item = (bool, bool)>) -> Vec<Run> {
    let mut result = Vec::new();
    let mut start = None;
    let mut end = 0;
    for (i, (linear, ple)) in layers.into_iter().enumerate() {
        end = i + 1;
        if linear && !ple {
            start.get_or_insert(i);
        } else {
            if let Some(start) = start.take() {
                result.push(Run::Gdn(start..i));
            }
            result.push(Run::Eager(i));
        }
    }
    if let Some(start) = start {
        result.push(Run::Gdn(start..end));
    }
    result
}

#[derive(Hash, PartialEq, Eq)]
struct RunKey {
    range: Range<usize>,
    capture: bool,
    stream: Stream,
    signature: Vec<(Vec<i32>, Dtype)>,
}

pub(super) struct TextModel<'a> {
    /// Original Module.parameters tree order, populated by checkpoint binding.
    pub(super) parameter_order: Vec<String>,
    pub(super) embedding: Linear<'a>,
    pub(super) layers: Vec<DecoderLayer<'a>>,
    pub(super) mixer: hyper::GatedResidual<'a>,
    pub(super) head: Option<Linear<'a>>, // None means tied embedding weights.
    pub(super) mtp: Option<super::mtp::Mtp<'a>>,
    pub(super) vision: Option<super::vision::Input>,
    pub(super) last_widened: Option<Array>,
    policy: CompilePolicy,
    decode_runs: Option<Vec<Run>>,
    replays: HashMap<RunKey, Replay>,
    commit_refusals: usize,
}

/// ArraysCache.state before verification; QSA entries need no snapshot here.
pub(super) type ArraySnapshot = Vec<Option<Array>>;

#[derive(Clone, Copy)]
pub(super) struct OutputOptions {
    pub(super) emit_logits: bool,
    pub(super) logits_keep: i32,
    pub(super) return_hidden: bool,
}

pub(super) struct Output {
    pub(super) logits: Option<Array>,
    pub(super) hidden: Option<Array>,
}

pub(super) fn embedding(layer: Linear<'_>, ids: &Array, stream: Stream) -> Result<Array, String> {
    match layer {
        Linear::Dense(weight) => indexing::take_rows(weight, ids, stream),
        Linear::Quantized(layer) => linear::dequantize(
            &indexing::take_rows(&layer.weight, ids, stream)?,
            &indexing::take_rows(&layer.scales, ids, stream)?,
            &indexing::take_rows(&layer.biases, ids, stream)?,
            layer.group,
            layer.bits,
            stream,
        ),
    }
}

fn mask(cache: Option<&Cache>, rows: i32, stream: Stream) -> Result<Option<Array>, String> {
    let Some(Cache {
        attention: AttentionCache::Linear(cache),
        ..
    }) = cache
    else {
        return Ok(None);
    };
    let (metadata, op) = if let Some(padding) = &cache.left_padding {
        (padding, Binary::GreaterEqual)
    } else if let Some(lengths) = &cache.lengths {
        (lengths, Binary::Less)
    } else {
        return Ok(None);
    };
    binary(
        &indexing::arange_i32(rows, stream)?,
        &ops::expand_dims(metadata, &[-1], stream)?,
        op,
        stream,
    )
    .map(Some)
}

impl<'a> TextModel<'a> {
    #[cfg(test)]
    pub(super) fn replay_count(&self) -> usize {
        self.replays.len()
    }
    pub(super) fn new(
        embedding: Linear<'a>,
        layers: Vec<DecoderLayer<'a>>,
        mixer: hyper::GatedResidual<'a>,
        head: Option<Linear<'a>>,
        policy: CompilePolicy,
    ) -> Result<Self, String> {
        if layers.is_empty() || mixer.inject.is_some() {
            return Err("Qwen requires decoder layers and a non-combining final mixer".into());
        }
        Ok(Self {
            embedding,
            layers,
            mixer,
            head,
            parameter_order: Vec::new(),
            mtp: None,
            vision: None,
            last_widened: None,
            policy,
            decode_runs: None,
            replays: HashMap::new(),
            commit_refusals: 0,
        })
    }

    pub(super) fn make_cache(&self, ratio: i32) -> Vec<Option<Cache>> {
        self.layers
            .iter()
            .map(|layer| {
                Some(Cache {
                    attention: match layer.attention {
                        Attention::Linear(_) => AttentionCache::Linear(gdn::Cache::default()),
                        Attention::Qsa(_) => {
                            AttentionCache::Qsa(QsaCache::new(if ratio == 0 { 4 } else { ratio }))
                        }
                    },
                    ple: ple::Cache::default(),
                    capture_ple: None,
                })
            })
            .collect()
    }

    pub(super) fn make_mtp_cache(&self, ratio: i32) -> [Option<Cache>; 1] {
        super::mtp::make_cache(ratio)
    }

    #[allow(clippy::too_many_arguments)]
    pub(super) fn mtp_forward(
        &mut self,
        hidden: &Array,
        ids: &Array,
        cache: Option<&mut [Option<Cache>]>,
        return_hidden: bool,
        options: &Options,
        vision: Option<(Option<&Array>, i32)>,
        streams: &Streams,
        stream: Stream,
    ) -> Result<Output, String> {
        let image = self.vision.clone();
        let vision = vision.or_else(|| image.as_ref().map(|v| (Some(&v.table), v.delta)));
        let emb = embedding(self.embedding, ids, stream)?;
        let mtp = self
            .mtp
            .as_mut()
            .ok_or("Qwen MTP weights are not attached")?;
        let h = mtp.fuse_and_run(hidden, &emb, cache, options, vision, streams, stream)?;
        let mixed = mtp
            .mixer
            .forward(&h, options.hyper_fused, options.hyper_v3, streams, stream)?
            .mixed;
        let logits = self.head.unwrap_or(self.embedding).apply(&mixed, stream)?;
        Ok(Output {
            logits: Some(logits),
            hidden: return_hidden.then_some(h),
        })
    }

    #[allow(clippy::too_many_arguments)]
    pub(super) fn mtp_update_cache(
        &mut self,
        hidden: &Array,
        ids: &Array,
        cache: Option<&mut [Option<Cache>]>,
        input_embeddings: Option<&Array>,
        options: &Options,
        vision: Option<(Option<&Array>, i32)>,
        streams: &Streams,
        stream: Stream,
    ) -> Result<Array, String> {
        let image = self.vision.clone();
        let vision = vision.or_else(|| image.as_ref().map(|v| (Some(&v.table), v.delta)));
        let spliced = match (&image, input_embeddings) {
            (Some(v), None) => v.embeddings(
                self.embedding,
                ids,
                super::vision::cache_position(cache.as_deref()) + 1,
                stream,
            )?,
            _ => None,
        };
        let emb = match input_embeddings.or(spliced.as_ref()) {
            Some(emb) => emb.clone(),
            None => embedding(self.embedding, ids, stream)?,
        };
        self.mtp
            .as_mut()
            .ok_or("Qwen MTP weights are not attached")?
            .fuse_and_run(hidden, &emb, cache, options, vision, streams, stream)
    }

    /// Model.post_weight_load: attach each PLE table after trunk weight loading.
    pub(super) fn post_weight_load(
        &mut self,
        root: &std::path::Path,
        streams: &Streams,
        stream: Stream,
    ) -> Result<(), String> {
        let path = root.join("ngram-table.safetensors");
        if !path.exists() {
            return Ok(());
        }
        let resident = super::ngram::resident_policy();
        let mut sidecar = None;
        for layer in &mut self.layers {
            let Some(ple) = &mut layer.ple else { continue };
            let table = &mut ple.embedding.table;
            table.attach_sidecar(&path)?;
            if resident {
                // Original attach_resident reports failure and retains the SSD fallback.
                table.attach_resident(&path, streams, stream);
            }
            sidecar = table.sidecar.as_ref();
        }
        if !resident && let Some(sidecar) = sidecar {
            println!(
                "[qwen4_exp] ngram table: streamed from SSD (file-backed pages, hot cache {}M, prefetch {}) — resident only at >=160G RAM or MTPLX_NGRAM_RESIDENT=1",
                sidecar.hot_mebibytes(),
                if sidecar.prefetch_enabled() {
                    "on"
                } else {
                    "off"
                }
            );
            let _ = std::io::Write::flush(&mut std::io::stdout());
        }
        Ok(())
    }

    #[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))]
    pub(super) fn clear_verify_capture(cache: &mut [Option<Cache>]) {
        for cache in cache.iter_mut().flatten() {
            cache.capture_ple = None;
            if let AttentionCache::Linear(c) = &mut cache.attention {
                c.capture = None;
            }
        }
    }

    fn refuse_commit(&mut self, layer: usize, reason: &str) -> bool {
        if self.commit_refusals < 3 {
            use std::io::Write;
            println!(
                "[qwen4_exp] capture-commit refused (layer {layer}: {reason}) — falling back to re-forward"
            );
            let _ = std::io::stdout().flush();
        }
        self.commit_refusals += 1;
        false
    }

    #[allow(clippy::too_many_arguments)]
    pub(super) fn commit_verified_window(
        &mut self,
        cache: &mut [Option<Cache>],
        snapshots: Option<&[Option<ArraySnapshot>]>,
        keep: i32,
        verified: i32,
        options: &Options,
        vision: Option<(Option<&Array>, i32)>,
        streams: &Streams,
        stream: Stream,
    ) -> Result<bool, String> {
        if keep < 1 || keep > verified || cache.len() != self.layers.len() {
            return Ok(false);
        }
        // First validate EVERY layer, without cache mutations or GPU work.
        // A refusal must leave the full verify window available for fallback.
        for (i, (layer, entry)) in self.layers.iter().zip(cache.iter()).enumerate() {
            let Some(entry) = entry else {
                return Ok(self.refuse_commit(i, "entry_missing"));
            };
            let AttentionCache::Linear(c) = &entry.attention else {
                continue;
            };
            let snapshot = snapshots.and_then(|s| s.get(i)).and_then(Option::as_ref);
            let validation = if layer.ple.is_some() {
                match (snapshot, &entry.capture_ple) {
                    (None, _) => Err("snapshot_missing".to_owned()),
                    (_, None) => Err("ple_rows_missing".to_owned()),
                    (Some(_), Some((h, _))) if h.layout().shape()[1] != verified => Err(format!(
                        "ple_rows_width_{}_vs_{verified}",
                        h.layout().shape()[1]
                    )),
                    (Some(pre), _) if pre.len() < 4 => Err("ple_snapshot_short".to_owned()),
                    _ => Ok(()),
                }
            } else {
                c.validate_verified_window(snapshot.map(Vec::as_slice), verified)
            };
            if let Err(reason) = validation {
                return Ok(self.refuse_commit(i, &reason));
            }
        }
        for (i, (layer, entry)) in self.layers.iter_mut().zip(cache).enumerate() {
            let entry = entry.as_mut().unwrap();
            if let AttentionCache::Qsa(c) = &mut entry.attention {
                if keep != verified {
                    c.trim(verified - keep);
                }
                continue;
            }
            let snapshot = snapshots.unwrap()[i].as_ref().unwrap();
            if layer.ple.is_some() {
                let (hidden, ids) = entry.capture_ple.as_ref().unwrap();
                let prefix = |a: &Array| {
                    let mut end = a.layout().shape().to_vec();
                    end[1] = keep;
                    ops::slice(a, &vec![0; end.len()], &end, &vec![1; end.len()], stream)
                };
                let hidden = prefix(hidden)?;
                let ids = prefix(ids)?;
                let AttentionCache::Linear(c) = &mut entry.attention else {
                    unreachable!()
                };
                c.conv = snapshot[0].clone();
                c.delta = snapshot[1].clone();
                entry.ple.conv = snapshot[2].clone();
                entry.ple.history = snapshot[3].clone();
                layer.ple.as_mut().unwrap().embedding.stage(
                    &ids,
                    entry.ple.history.as_ref(),
                    &options.stage,
                    streams,
                    stream,
                );
                layer.call(
                    &hidden,
                    Some(&ids),
                    None,
                    Some(entry),
                    options,
                    vision,
                    streams,
                    stream,
                )?;
                entry.capture_ple = None;
            } else {
                let AttentionCache::Linear(c) = &mut entry.attention else {
                    unreachable!()
                };
                let Attention::Linear(weights) = &layer.attention else {
                    return Err("GDN capture-commit layer/cache mismatch".into());
                };
                c.replay_verified_prefix(snapshot, keep, weights.a_log, weights.dt_bias, stream)?;
            }
        }
        Ok(true)
    }

    #[allow(clippy::too_many_arguments)]
    pub(super) fn call(
        &mut self,
        ids: &Array,
        cache: Option<&mut [Option<Cache>]>,
        input_embeddings: Option<&Array>,
        output: OutputOptions,
        options: &Options,
        vision: Option<(Option<&Array>, i32)>,
        streams: &Streams,
        stream: Stream,
    ) -> Result<Output, String> {
        let image = self.vision.clone();
        let vision = vision.or_else(|| image.as_ref().map(|v| (Some(&v.table), v.delta)));
        let spliced = match (&image, input_embeddings) {
            (Some(v), None) => v.embeddings(
                self.embedding,
                ids,
                super::vision::cache_position(cache.as_deref()),
                stream,
            )?,
            _ => None,
        };
        let h = match input_embeddings.or(spliced.as_ref()) {
            Some(h) => h.clone(),
            None => embedding(self.embedding, ids, stream)?,
        };
        let shape = h.layout().shape().to_vec();
        if shape.len() != 3
            || shape[0] < 1
            || shape[1] < 1
            || shape[2] != 2560
            || ids.layout().shape() != &shape[..2]
        {
            return Err(
                "installed Qwen requires [B,S,2560] embeddings and matching token IDs".into(),
            );
        }
        let mut empty = Vec::new();
        let cache = match cache {
            Some(cache) => cache,
            None => {
                empty.resize_with(self.layers.len(), || None);
                &mut empty
            }
        };
        if cache.len() != self.layers.len() {
            return Err("Qwen cache/layer count mismatch".into());
        }
        let ssm = self
            .layers
            .iter()
            .position(|l| matches!(l.attention, Attention::Linear(_)))
            .unwrap_or(0);
        let ssm_mask = mask(cache[ssm].as_ref(), shape[1], stream)?;
        if let Some(i) = self.layers.iter().position(|l| l.ple.is_some()) {
            self.layers[i].ple.as_mut().unwrap().embedding.stage(
                ids,
                cache[i].as_ref().and_then(|c| c.ple.history.as_ref()),
                &options.stage,
                streams,
                stream,
            );
        }
        // tile(h, (1,1,4)): original reshape -> broadcast -> reshape.
        let h = ops::reshape(&h, &[shape[0], shape[1], 1, 2560], stream)?;
        let h = ops::broadcast_to(&h, &[shape[0], shape[1], 4, 2560], stream)?;
        let mut h = ops::reshape(&h, &[shape[0], shape[1], 10240], stream)?;
        if (1..=4).contains(&shape[1])
            && ssm_mask.is_none()
            && self.policy.enabled()
            && cache[ssm].is_some()
        {
            h = self.decode(h, ids, cache, options, vision, streams, stream)?;
        } else {
            for (layer, cache) in self.layers.iter_mut().zip(cache) {
                h = eager(
                    layer,
                    &h,
                    ids,
                    ssm_mask.as_ref(),
                    cache.as_mut(),
                    options,
                    vision,
                    streams,
                    stream,
                )?;
            }
        }
        self.last_widened = Some(h.clone());
        let mixed = self
            .mixer
            .forward(&h, options.hyper_fused, options.hyper_v3, streams, stream)?
            .mixed;
        let logits = if output.emit_logits {
            let mixed = if output.logits_keep != 0 {
                let keep = output.logits_keep.max(1).min(shape[1]);
                ops::slice(
                    &mixed,
                    &[0, shape[1] - keep, 0],
                    &[shape[0], shape[1], 2560],
                    &[1; 3],
                    stream,
                )?
            } else {
                mixed
            };
            Some(self.head.unwrap_or(self.embedding).apply(&mixed, stream)?)
        } else {
            None
        };
        Ok(Output {
            logits,
            hidden: output.return_hidden.then_some(h),
        })
    }

    #[allow(clippy::too_many_arguments)]
    fn decode(
        &mut self,
        mut h: Array,
        ids: &Array,
        cache: &mut [Option<Cache>],
        options: &Options,
        vision: Option<(Option<&Array>, i32)>,
        streams: &Streams,
        stream: Stream,
    ) -> Result<Array, String> {
        let runs = self.decode_runs.get_or_insert_with(|| {
            runs(
                self.layers
                    .iter()
                    .map(|l| (matches!(l.attention, Attention::Linear(_)), l.ple.is_some())),
            )
        });
        for run in runs {
            let range = match run {
                Run::Eager(i) => {
                    h = eager(
                        &mut self.layers[*i],
                        &h,
                        ids,
                        None,
                        cache[*i].as_mut(),
                        options,
                        vision,
                        streams,
                        stream,
                    )?;
                    continue;
                }
                Run::Gdn(range) => range.clone(),
            };
            let mut flat = vec![h.clone()];
            for c in &cache[range.clone()] {
                let Some(Cache {
                    attention: AttentionCache::Linear(c),
                    ..
                }) = c
                else {
                    break;
                };
                let (Some(conv), Some(delta)) = (&c.conv, &c.delta) else {
                    break;
                };
                flat.extend([conv.clone(), delta.clone()]);
            }
            if flat.len() != 1 + 2 * range.len() {
                for i in range {
                    h = eager(
                        &mut self.layers[i],
                        &h,
                        ids,
                        None,
                        cache[i].as_mut(),
                        options,
                        vision,
                        streams,
                        stream,
                    )?;
                }
                continue;
            }
            let key = RunKey {
                range: range.clone(),
                capture: options.gdn.capture,
                stream,
                signature: flat
                    .iter()
                    .map(|a| (a.layout().shape().to_vec(), a.layout().dtype()))
                    .collect(),
            };
            let replay = match self.replays.entry(key) {
                std::collections::hash_map::Entry::Occupied(entry) => entry.into_mut(),
                std::collections::hash_map::Entry::Vacant(entry) => {
                    let replay = Replay::trace_compiled(&flat, |args| {
                        let mut h = args[0].clone();
                        let mut states = Vec::new();
                        let mut captured = Vec::new();
                        for (n, i) in range.clone().enumerate() {
                            let mut c = Cache {
                                attention: AttentionCache::Linear(gdn::Cache {
                                    conv: Some(args[1 + 2 * n].clone()),
                                    delta: Some(args[2 + 2 * n].clone()),
                                    ..Default::default()
                                }),
                                ple: ple::Cache::default(),
                                capture_ple: None,
                            };
                            h = self.layers[i].call(
                                &h,
                                None,
                                None,
                                Some(&mut c),
                                options,
                                vision,
                                streams,
                                stream,
                            )?;
                            let AttentionCache::Linear(c) = c.attention else {
                                unreachable!()
                            };
                            states.extend([
                                c.conv.ok_or("compiled GDN conv missing")?,
                                c.delta.ok_or("compiled GDN state missing")?,
                            ]);
                            if options.gdn.capture {
                                captured.extend(c.capture.ok_or("compiled GDN capture missing")?);
                            }
                        }
                        Ok(std::iter::once(h).chain(states).chain(captured).collect())
                    })?;
                    entry.insert(replay)
                }
            };
            let out = replay.run(&flat)?;
            h = out[0].clone();
            for (n, i) in range.clone().enumerate() {
                let AttentionCache::Linear(c) = &mut cache[i].as_mut().unwrap().attention else {
                    unreachable!()
                };
                c.conv = Some(out[1 + 2 * n].clone());
                c.delta = Some(out[2 + 2 * n].clone());
                if options.gdn.capture {
                    let start = 1 + 2 * range.len() + 6 * n;
                    c.capture = Some(std::array::from_fn(|j| out[start + j].clone()));
                }
            }
        }
        Ok(h)
    }
}

#[cfg(test)]
impl TextModel<'_> {
    pub(super) fn run_contract(&self, rows: i32, capture: bool) -> serde_json::Value {
        use super::ops::Operation as O;
        let replay = &self
            .replays
            .iter()
            .find(|(key, _)| key.capture == capture && key.signature[0].0 == [1, rows, 10240])
            .expect("compiled fixture run missing")
            .1;
        let mut counts = std::collections::BTreeMap::<String, usize>::new();
        let contract = replay.contract(
            &["hidden", "conv", "delta"],
            |op| {
                let name = match op {
                    O::Compiled(expression) => expression.name(),
                    _ => match op {
                        O::Reshape => "Reshape",
                        O::Flatten { .. } => "Flatten",
                        O::Unflatten { .. } => "Unflatten",
                        O::Transpose(_) => "Transpose",
                        O::ExpandDims(_) => "ExpandDims",
                        O::Squeeze(_) => "Squeeze",
                        O::Slice { .. } => "Slice",
                        O::Contiguous(_) => "Contiguous",
                        O::RmsNorm(_) => "RMSNorm",
                        O::Softmax(_) => "Softmax",
                        O::QuantizedLinearBf16 { .. } => "QuantizedMatmul",
                        O::GatherQmmBf16 { .. } => "GatherQMM",
                        O::DenseMatmul => "Matmul",
                        O::Sum(_) => "Sum",
                        O::Sort {
                            indices: true,
                            kth: Some(_),
                            ..
                        } => "ArgPartition",
                        O::Split { .. } => "Split",
                        O::HyperRead(_)
                        | O::GdnStepBf16
                        | O::GdnConvNorm(_)
                        | O::GatedDeltaBf16(_) => "CustomKernel",
                        O::Concatenate(_) => "Concatenate",
                        O::DepthwiseConv1d => "Convolution",
                        O::GatherAxis(_) => "GatherAxis",
                        O::ArangeI32(_) | O::ArangeU32 => "Arange",
                        _ => super::fused::operator(op).expect("uncaptured GDN primitive"),
                    }
                    .to_owned(),
                };
                *counts.entry(name.clone()).or_default() += 1;
                name
            },
            if capture { 9 } else { 3 },
        );
        serde_json::json!({"contract":contract,"ops":counts})
    }
}

#[test]
fn mtplx_model_compile_policy_and_run_boundaries() {
    assert_eq!(
        runs([
            (true, false),
            (true, false),
            (true, true),
            (true, false),
            (false, false),
            (true, false)
        ]),
        vec![
            Run::Gdn(0..2),
            Run::Eager(2),
            Run::Gdn(3..4),
            Run::Eager(4),
            Run::Gdn(5..6)
        ]
    );
    for on in ["1", " true ", "YES", "On"] {
        assert!(CompilePolicy::from_values(Some(on), None).enabled());
        assert!(CompilePolicy::from_values(None, Some(on)).enabled());
        for off in ["0", "", "false", "off", "anything"] {
            for policy in [
                CompilePolicy::from_values(Some(on), Some(off)),
                CompilePolicy::from_values(Some(off), Some(on)),
            ] {
                assert!(!policy.enabled());
            }
        }
    }
    let automatic = CompilePolicy::from_values(None, None);
    assert!(!automatic.enabled());
    let environment = CompilePolicy::from_env();
    assert_eq!(
        environment.enabled(),
        CompilePolicy::from_values(
            std::env::var("MTPLX_COMPILED_GDN").ok().as_deref(),
            std::env::var("MTPLX_QWEN4EXP_COMPILE").ok().as_deref(),
        )
        .enabled()
    );
}

#[allow(clippy::too_many_arguments)]
fn eager(
    layer: &mut DecoderLayer<'_>,
    hidden: &Array,
    ids: &Array,
    mask: Option<&Array>,
    mut cache: Option<&mut Cache>,
    options: &Options,
    vision: Option<(Option<&Array>, i32)>,
    streams: &Streams,
    stream: Stream,
) -> Result<Array, String> {
    if options.gdn.capture
        && layer.ple.is_some()
        && let Some(c) = cache.as_deref_mut()
    {
        c.capture_ple = Some((hidden.clone(), ids.clone()));
    }
    layer.call(
        hidden,
        Some(ids),
        mask,
        cache,
        options,
        vision,
        streams,
        stream,
    )
}