gammalooprs/momentum/

sample.rs

use std::fmt::Display;

use crate::graph::{Graph, LMBext, LmbIndex, LoopMomentumBasis};
use crate::integrands::process::evaluators::SingleOrAllOrientations;
use crate::momentum::{FourMomentum, Polarization, Rotatable, Rotation, SignOrZero, ThreeMomentum};

use crate::momentum::signature::LoopSignature;
use crate::utils::hyperdual_utils::new_constant;
use crate::utils::{ArbPrec, F, FloatLike, Length, PrecisionUpgradable};
use crate::{DependentMomentaConstructor, define_index, settings::runtime::kinematic::Externals};
use bincode_trait_derive::{Decode, Encode};
use color_eyre::Result;
use derive_more::{From, Into};
use eyre::eyre;
use linnet::half_edge::HedgeGraph;
use linnet::half_edge::involution::{EdgeIndex, EdgeVec, Orientation};
use linnet::half_edge::nodestore::NodeStorageOps;
use linnet::half_edge::subgraph::subset::SubSet;
use linnet::half_edge::subgraph::{
    Inclusion, InternalSubGraph, ModifySubSet, SuBitGraph, SubSetLike, SubSetOps,
};
use linnet::half_edge::typed_vec::IndexLike;
use serde::{Deserialize, Serialize};
use std::ops::{Add, Index, IndexMut, Sub};
use symbolica::domains::dual::HyperDual;
use tabled::settings::Style;

use typed_index_collections::TiVec;

#[derive(
    From,
    Into,
    Copy,
    Clone,
    Hash,
    Eq,
    Ord,
    PartialEq,
    PartialOrd,
    bincode::Encode,
    bincode::Decode,
    Debug,
    Serialize,
    Deserialize,
)]
pub struct LoopIndex(pub usize);

// #[derive(
//     From,
//     Into,
//     Copy,
//     Clone,
//     Hash,
//     Eq,
//     Ord,
//     PartialEq,
//     PartialOrd,
//     Encode,
//     Decode,
//     Debug,
//     Serialize,
//     Deserialize,
// )]
// pub struct ExternalIndex(pub usize);

define_index!(
    pub struct ExternalIndex;
);

// define_indexed_vec!(
//     EdgeIndex;

//     pub struct ExternalThr;
// );

impl Display for LoopIndex {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.0)
    }
}

impl Display for ExternalIndex {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.0)
    }
}

#[derive(From, Into, Serialize, Deserialize, Clone, PartialEq, Debug, Encode, Decode)]
pub struct LoopMomenta<T>(pub Vec<ThreeMomentum<T>>);

impl<T: Display> Display for LoopMomenta<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let mut table_builder = tabled::builder::Builder::new();
        for (i, mom) in self.0.iter().enumerate() {
            let mut row = Vec::new();
            row.push(i.to_string());
            for m in mom {
                row.push(m.to_string());
            }
            table_builder.push_record(row);
        }

        write!(f, "{}", table_builder.build().with(Style::modern_rounded()))
    }
}

impl<T> Length for LoopMomenta<T> {
    fn is_empty(&self) -> bool {
        self.0.is_empty()
    }

    fn len(&self) -> usize {
        self.0.len()
    }
}

pub type Subspace<'a> = Option<&'a [LoopIndex]>; // None means full space

#[derive(Clone, Debug, bincode::Encode, bincode::Decode)]
pub struct SubspaceData {
    subgraph: InternalSubGraph,
    #[bincode(with_serde)]
    complement_subgraph: SuBitGraph,
    lmb: LmbIndex,
    lmb_indices: Vec<LoopIndex>,
}

impl SubspaceData {
    /// Canonical coordinates constrained by one threshold solve. Each defining edge is paired
    /// with its complete signed fundamental cycle, oriented positively along that edge. Parent
    /// basis slots and changes of the fixed complement do not change this physical identity.
    pub fn solve_signature(
        &self,
        all_lmbs: &TiVec<LmbIndex, LoopMomentumBasis>,
    ) -> Vec<(EdgeIndex, Vec<(EdgeIndex, SignOrZero)>)> {
        let lmb = self.get_lmb(all_lmbs);
        let mut signature = self
            .lmb_indices
            .iter()
            .copied()
            .map(|loop_index| {
                let defining_edge = lmb.loop_edges[loop_index];
                let orientation = lmb.edge_signatures[defining_edge].internal[loop_index];
                assert!(
                    orientation.is_sign(),
                    "a defining edge must carry its own loop momentum"
                );
                let mut cycle = lmb
                    .edge_signatures
                    .iter()
                    .filter_map(|(edge, signature)| {
                        let sign = signature.internal[loop_index] * orientation;
                        sign.is_sign().then_some((edge, sign))
                    })
                    .collect::<Vec<_>>();
                cycle.sort_by_key(|(edge, _)| *edge);
                (defining_edge, cycle)
            })
            .collect::<Vec<_>>();
        signature.sort_by_key(|(edge, _)| *edge);
        signature
    }

    #[cfg(test)]
    pub(crate) fn has_same_embedding(&self, other: &Self) -> bool {
        self.lmb == other.lmb
            && self.lmb_indices == other.lmb_indices
            && self.subgraph.filter == other.subgraph.filter
            && self.complement_subgraph == other.complement_subgraph
    }

    pub(crate) fn has_equivalent_embedding(
        &self,
        other: &Self,
        all_lmbs: &TiVec<LmbIndex, LoopMomentumBasis>,
    ) -> bool {
        // Different parent bases are equivalent when their selected edge momenta vary along
        // the same signed cycles. Runtime data may still retain either parent's coordinates.
        self.solve_signature(all_lmbs) == other.solve_signature(all_lmbs)
    }

    pub(crate) fn is_mergable_with(
        &self,
        other: &Self,
        all_lmbs: &TiVec<LmbIndex, LoopMomentumBasis>,
    ) -> bool {
        // Independent radial projections must leave each other's defining edge momenta fixed.
        // This is invariant under changes of parent basis and permits adding their physical
        // momentum displacements without copying unrelated native complement coordinates.
        [(self, other), (other, self)]
            .into_iter()
            .all(|(active, fixed)| {
                let lmb = active.get_lmb(all_lmbs);
                fixed.iter_basis_edges(all_lmbs).all(|edge| {
                    active.iter_lmb_indices().all(|loop_index| {
                        lmb.edge_signatures[edge].internal[loop_index] == SignOrZero::Zero
                    })
                })
            })
    }

    /// Build the smallest active subspace containing every supplied subspace in one parent LMB.
    ///
    /// This test scaffold checks that the overlap wrapper partitions different subspaces before
    /// solving. Runtime groups never union incompatible coordinates: each threshold retains its
    /// selected subspace and fixes its complementary coordinates at its own sampled point.
    #[cfg(test)]
    pub(crate) fn union_in_common_parent<'a>(
        subspaces: impl IntoIterator<Item = &'a Self>,
        graph: &Graph,
        all_lmbs: &TiVec<LmbIndex, LoopMomentumBasis>,
    ) -> Result<Self> {
        let mut subspaces = subspaces.into_iter();
        let first = subspaces
            .next()
            .ok_or_else(|| eyre!("Cannot build a common threshold subspace from an empty set"))?;
        let parent_lmb_index = first.lmb;
        let mut active_subgraph = first.subgraph.filter.clone();
        let mut expected_indices = first
            .lmb_indices
            .iter()
            .copied()
            .collect::<std::collections::BTreeSet<_>>();

        for subspace in subspaces {
            if subspace.lmb != parent_lmb_index {
                return Err(eyre!(
                    "Threshold subspaces use different parent LMBs {} and {}",
                    usize::from(parent_lmb_index),
                    usize::from(subspace.lmb),
                ));
            }
            active_subgraph = active_subgraph.union(&subspace.subgraph.filter);
            expected_indices.extend(subspace.lmb_indices.iter().copied());
        }

        let mut union =
            Self::new_with_user_selected_lmb(active_subgraph, parent_lmb_index, graph, all_lmbs)?;
        let resolved_indices = union
            .lmb_indices
            .iter()
            .copied()
            .collect::<std::collections::BTreeSet<_>>();
        if !expected_indices.is_subset(&resolved_indices) {
            return Err(eyre!(
                "Common threshold subspace resolved active indices {:?}, which do not contain the requested union {:?}",
                resolved_indices,
                expected_indices,
            ));
        }
        // Joining projected subgraphs can close an additional graph-theoretic cycle even though
        // none of the participating threshold variants varies that coordinate. The common
        // object is a center-coordinate frame, not another threshold surface: keep precisely the
        // union of the variants' parent-LMB coordinates and leave every other coordinate fixed.
        union.lmb_indices = expected_indices.into_iter().collect();
        Ok(union)
    }

    fn cleaned_filter_pessimist<E, V, H, N: NodeStorageOps<NodeData = V>>(
        mut filter: SuBitGraph,
        graph: &HedgeGraph<E, V, H, N>,
    ) -> InternalSubGraph {
        let mut to_remove = SuBitGraph::empty(filter.size());

        for i in filter.included_iter() {
            if !filter.includes(&graph.inv(i)) {
                to_remove.add(i);
            }
        }
        filter.subtract_with(&to_remove);

        InternalSubGraph {
            filter,
            loopcount: None,
        }
    }

    pub(crate) fn new_with_user_selected_lmb(
        subgraph: SuBitGraph,
        lmb: LmbIndex,
        graph: &Graph,
        all_lmbs: &TiVec<LmbIndex, LoopMomentumBasis>,
    ) -> Result<Self> {
        let bridges = graph.bridges_of(&subgraph);

        let mut subgraph = Self::cleaned_filter_pessimist(subgraph.subtract(&bridges), graph);
        subgraph.set_loopcount(graph);

        let edges_in_subgraph = graph
            .iter_edges_of(&subgraph.filter)
            .map(|ed| ed.1)
            .collect::<Vec<_>>();
        let complement_subgraph = graph.full_filter().subtract(&subgraph.filter);

        let parent_lmb = &all_lmbs[lmb];
        let compatible_sub_lmb =
            graph.try_compatible_sub_lmb(&subgraph, graph.full_crown(&subgraph), parent_lmb)?;
        let mut lmb_indices = compatible_sub_lmb
            .loop_edges
            .iter()
            .map(|edge_index| {
                parent_lmb
                    .loop_edges
                    .iter_enumerated()
                    .find_map(|(loop_index, parent_edge)| {
                        (parent_edge == edge_index).then_some(loop_index)
                    })
                    .ok_or_else(|| {
                        eyre!(
                            "Compatible subgraph LMB edge {:?} is not a defining edge of parent LMB {:?}",
                            edge_index,
                            parent_lmb.loop_edges,
                        )
                    })
            })
            .collect::<Result<Vec<_>>>()?;
        lmb_indices.sort();

        if lmb_indices.len() != subgraph.loopcount.unwrap()
            || compatible_sub_lmb.loop_edges.len() != subgraph.loopcount.unwrap()
        {
            return Err(color_eyre::eyre::eyre!(
                "Provided loop momentum basis is not topologically compatible with subgraph, lmb_indices: {:?}, lmb_edges: {:?}, compatible_sub_lmb_edges: {:?}, subgraph loopcount: {}, edges_in_subgraph: {:?}",
                lmb_indices,
                parent_lmb.loop_edges,
                compatible_sub_lmb.loop_edges,
                subgraph.loopcount.unwrap(),
                edges_in_subgraph,
            ));
        }

        Ok(Self {
            subgraph,
            complement_subgraph,
            lmb,
            lmb_indices,
        })
    }

    /// Build a subspace from defining edges in an already selected parent LMB.
    ///
    /// The requested order is intentionally not stored here: [`SubspaceData`] describes the
    /// geometric active-coordinate set and canonicalizes its loop indices. Callers that attach
    /// semantic meaning to the user's ordering must retain the original edge list separately.
    pub(crate) fn new_from_parent_basis_edges(
        requested_basis_edges: &[EdgeIndex],
        containing_subgraph: &SuBitGraph,
        parent_lmb_index: LmbIndex,
        graph: &Graph,
        all_lmbs: &TiVec<LmbIndex, LoopMomentumBasis>,
    ) -> Result<Self> {
        if requested_basis_edges.is_empty() {
            return Err(eyre!(
                "A threshold-counterterm subspace must contain at least one basis edge"
            ));
        }

        let parent_lmb = all_lmbs.get(parent_lmb_index).ok_or_else(|| {
            eyre!(
                "Threshold-counterterm parent LMB {} is out of range ({} generated LMBs)",
                usize::from(parent_lmb_index),
                all_lmbs.len(),
            )
        })?;
        let mut requested_indices = Vec::with_capacity(requested_basis_edges.len());
        let mut seen_basis_edges = std::collections::BTreeSet::new();
        for &edge in requested_basis_edges {
            if !seen_basis_edges.insert(edge) {
                return Err(eyre!(
                    "Threshold-counterterm subspace repeats basis edge {edge}"
                ));
            }
            let loop_index = parent_lmb
                .loop_edges
                .iter_enumerated()
                .find_map(|(loop_index, &parent_edge)| {
                    (parent_edge == edge).then_some(loop_index)
                })
                .ok_or_else(|| {
                    eyre!(
                        "Threshold-counterterm subspace edge {edge} is not a defining edge of parent LMB {:?}",
                        parent_lmb.loop_edges,
                    )
                })?;
            requested_indices.push(loop_index);
        }

        // A parent's fundamental cycle contains precisely the edges whose signature carries the
        // corresponding defining loop momentum. Restricting the union to the inferred cut side
        // produces the active topology while preserving the sampled complement outside it.
        let mut cycle_union: SuBitGraph = graph.empty_subgraph();
        for (edge_pair, edge, _) in graph.iter_edges() {
            if requested_indices.iter().any(|&loop_index| {
                parent_lmb.edge_signatures[edge].internal[loop_index] != SignOrZero::Zero
            }) {
                cycle_union.add(edge_pair);
            }
        }
        let active_subgraph = cycle_union.intersection(containing_subgraph);
        let subspace =
            Self::new_with_user_selected_lmb(active_subgraph, parent_lmb_index, graph, all_lmbs)?;

        let resolved_basis_edges = subspace.iter_basis_edges(all_lmbs).collect::<Vec<_>>();
        let requested_set = requested_basis_edges
            .iter()
            .copied()
            .collect::<std::collections::BTreeSet<_>>();
        let resolved_set = resolved_basis_edges
            .iter()
            .copied()
            .collect::<std::collections::BTreeSet<_>>();
        if subspace.loopcount() != requested_basis_edges.len() || resolved_set != requested_set {
            return Err(eyre!(
                "Requested threshold-counterterm basis {:?} resolves to basis {:?} with loop count {} in parent LMB {:?}",
                requested_basis_edges,
                resolved_basis_edges,
                subspace.loopcount(),
                parent_lmb.loop_edges,
            ));
        }

        Ok(subspace)
    }

    /// this function chooses the lmb automatically based on the subgraph
    #[allow(dead_code)]
    pub(crate) fn new(
        subgraph: SuBitGraph,
        graph: &Graph,
        all_lmbs: &TiVec<LmbIndex, LoopMomentumBasis>,
    ) -> Result<Self> {
        let mut errors = Vec::new();
        for (lmb_index, _) in all_lmbs.iter_enumerated() {
            match Self::new_with_user_selected_lmb(subgraph.clone(), lmb_index, graph, all_lmbs) {
                Ok(subspace) => return Ok(subspace),
                Err(error) => errors.push(format!("LMB {}: {error:#}", usize::from(lmb_index))),
            }
        }

        Err(eyre!(
            "No topology-compatible loop momentum basis found for subgraph:\n{}",
            errors.join("\n")
        ))
    }
    pub(crate) fn get_lmb<'a>(
        &self,
        all_lmbs: &'a TiVec<LmbIndex, LoopMomentumBasis>,
    ) -> &'a LoopMomentumBasis {
        &all_lmbs[self.lmb]
    }

    pub(crate) fn parent_lmb_index(&self) -> LmbIndex {
        self.lmb
    }

    pub(crate) fn does_not_contain<'a>(
        &'a self,
        edges: &'a [EdgeIndex],
        graph: &'a Graph,
    ) -> impl Iterator<Item = EdgeIndex> + 'a {
        graph
            .iter_edges_of(&self.complement_subgraph)
            .map(|ed| ed.1)
            .filter(|e| edges.contains(e))
    }

    pub(crate) fn contains<'a>(
        &'a self,
        edges: &'a [EdgeIndex],
        graph: &'a Graph,
    ) -> impl Iterator<Item = EdgeIndex> + 'a {
        graph
            .iter_edges_of(&self.subgraph)
            .map(|ed| ed.1)
            .filter(|e| edges.contains(e))
    }

    pub(crate) fn project_loop_signature<'a>(
        &'a self,
        signature: &'a LoopSignature,
    ) -> impl Iterator<Item = SignOrZero> + 'a {
        signature.iter_enumerated().map(|(loop_index, sign)| {
            if self.lmb_indices.contains(&loop_index) {
                *sign
            } else {
                SignOrZero::Zero
            }
        })
    }

    pub(crate) fn project_loop_signature_filtered<'a>(
        &'a self,
        signature: &'a LoopSignature,
    ) -> impl Iterator<Item = SignOrZero> + 'a {
        signature
            .iter_enumerated()
            .filter_map(move |(loop_index, sign)| {
                if self.lmb_indices.contains(&loop_index) {
                    Some(*sign)
                } else {
                    None
                }
            })
    }

    pub(crate) fn project_complement_loop_signature<'a>(
        &'a self,
        signature: &'a LoopSignature,
    ) -> impl Iterator<Item = SignOrZero> + 'a {
        signature.iter_enumerated().map(|(loop_index, sign)| {
            if !self.lmb_indices.contains(&loop_index) {
                *sign
            } else {
                SignOrZero::Zero
            }
        })
    }

    pub(crate) fn iter_lmb_indices<'a>(&'a self) -> impl Iterator<Item = LoopIndex> + 'a {
        self.lmb_indices.iter().copied()
    }

    pub(crate) fn iter_basis_edges<'a>(
        &'a self,
        all_lmbs: &'a TiVec<LmbIndex, LoopMomentumBasis>,
    ) -> impl Iterator<Item = EdgeIndex> + 'a {
        self.lmb_indices
            .iter()
            .map(|&loop_index| self.get_lmb(all_lmbs).loop_edges[loop_index])
    }

    pub(crate) fn contains_loop_index(&self, loop_index: LoopIndex) -> bool {
        self.lmb_indices.contains(&loop_index)
    }

    pub(crate) fn loopcount(&self) -> usize {
        // `lmb_indices` is the runtime definition of an active coordinate. For ordinary
        // threshold subspaces its length equals the graph-theoretic subgraph loop count. A
        // projected common-center frame may deliberately exclude incidental cycles closed by
        // the union of several variant subgraphs.
        self.lmb_indices.len()
    }

    pub(crate) fn as_subspace_simple(&self) -> Subspace<'_> {
        Some(self.lmb_indices.as_slice())
    }
}

// #[comemo::track]
// impl<T> LoopMomenta<T> {}

impl<T> LoopMomenta<T> {
    pub(crate) fn iter(&self) -> std::slice::Iter<'_, ThreeMomentum<T>> {
        self.0.iter()
    }

    pub(crate) fn first(&self) -> Option<&ThreeMomentum<T>> {
        self.0.first()
    }

    pub(crate) fn iter_enumerated(&self) -> impl Iterator<Item = (LoopIndex, &ThreeMomentum<T>)> {
        self.0.iter().enumerate().map(|(i, m)| (LoopIndex(i), m))
    }
}

impl<T> IntoIterator for LoopMomenta<T> {
    type Item = ThreeMomentum<T>;
    type IntoIter = std::vec::IntoIter<ThreeMomentum<T>>;

    fn into_iter(self) -> Self::IntoIter {
        self.0.into_iter()
    }
}

impl<T> FromIterator<ThreeMomentum<T>> for LoopMomenta<T> {
    fn from_iter<I: IntoIterator<Item = ThreeMomentum<T>>>(iter: I) -> Self {
        LoopMomenta(iter.into_iter().collect())
    }
}

impl<T> Index<LoopIndex> for LoopMomenta<T> {
    type Output = ThreeMomentum<T>;

    fn index(&self, index: LoopIndex) -> &Self::Output {
        &self.0[index.0]
    }
}

impl<T> IndexMut<LoopIndex> for LoopMomenta<T> {
    fn index_mut(&mut self, index: LoopIndex) -> &mut Self::Output {
        &mut self.0[index.0]
    }
}

impl<T: FloatLike> LoopMomenta<F<T>> {
    pub(crate) fn hyper_radius_squared(&self, subspace: Subspace) -> F<T> {
        let zero = self.0[0].px.zero();
        match subspace {
            None => self.iter().fold(zero, |acc, x| acc + x.norm_squared()),
            Some(subspace) => subspace
                .iter()
                .fold(zero, |acc, &i| acc + self[i].norm_squared()),
        }
    }

    pub(crate) fn rescale(&self, factor: &F<T>, subspace: Subspace) -> Self {
        match subspace {
            None => LoopMomenta::from_iter(self.iter().map(|k| k * factor)),
            // this branch is wrong
            Some(subspace) => LoopMomenta::from_iter(self.iter_enumerated().map(|(i, k)| {
                if subspace.contains(&i) {
                    k * factor
                } else {
                    k.clone()
                }
            })),
        }
    }

    pub(crate) fn rescale_with_hyper_dual(
        &self,
        factor: &HyperDual<F<T>>,
        subspace: Subspace,
    ) -> LoopMomenta<HyperDual<F<T>>> {
        match subspace {
            None => LoopMomenta::from_iter(self.iter().map(|k| {
                ThreeMomentum::new(
                    new_constant(factor, &k.px) * factor,
                    new_constant(factor, &k.py) * factor,
                    new_constant(factor, &k.pz) * factor,
                )
            })),
            Some(subspace) => LoopMomenta::from_iter(self.iter_enumerated().map(|(i, k)| {
                if subspace.contains(&i) {
                    ThreeMomentum::new(
                        new_constant(factor, &k.px) * factor,
                        new_constant(factor, &k.py) * factor,
                        new_constant(factor, &k.pz) * factor,
                    )
                } else {
                    ThreeMomentum::new(
                        new_constant(factor, &k.px),
                        new_constant(factor, &k.py),
                        new_constant(factor, &k.pz),
                    )
                }
            })),
        }
    }

    pub(crate) fn rotate(&self, rotation: &Rotation) -> Self {
        LoopMomenta::from_iter(self.iter().map(|k| k.rotate(rotation)))
    }

    pub(crate) fn lmb_transform(
        &self,
        from: &LoopMomentumBasis,
        to: &LoopMomentumBasis,
        externals: &ExternalThreeMomenta<F<T>>,
    ) -> Self {
        LoopMomenta::from_iter(
            to.loop_edges
                .iter()
                .map(|e_id| from.edge_signatures[*e_id].compute_momentum(self, externals)),
        )
    }
}

impl<T: FloatLike> LoopMomenta<HyperDual<F<T>>> {
    pub fn lmb_transform(
        &self,
        from: &LoopMomentumBasis,
        to: &LoopMomentumBasis,
        externals: &ExternalThreeMomenta<HyperDual<F<T>>>,
    ) -> Self {
        LoopMomenta::from_iter(to.loop_edges.iter().map(|e_id| {
            from.edge_signatures[*e_id]
                .try_compute_momentum(&self.0, &externals.raw)
                .unwrap()
        }))
    }

    pub fn rescale(&self, factor: &HyperDual<F<T>>, subspace: Subspace) -> Self {
        match subspace {
            None => LoopMomenta::from_iter(self.iter().map(|k| k * factor)),
            // this branch is wrong
            Some(subspace) => LoopMomenta::from_iter(self.iter_enumerated().map(|(i, k)| {
                if subspace.contains(&i) {
                    k * factor
                } else {
                    k.clone()
                }
            })),
        }
    }
}

impl LoopMomenta<F<f64>> {
    pub(crate) fn cast<T: FloatLike>(&self) -> LoopMomenta<F<T>> {
        LoopMomenta::from_iter(self.iter().map(|m| m.map(&|x| F::from_ff64(x))))
    }
}

impl<T: FloatLike> Sub<&LoopMomenta<F<T>>> for &LoopMomenta<F<T>> {
    type Output = LoopMomenta<F<T>>;

    fn sub(self, rhs: &LoopMomenta<F<T>>) -> Self::Output {
        LoopMomenta::from_iter(self.iter().zip(rhs.iter()).map(|(l, r)| l - r))
    }
}

impl<T: FloatLike> Add<&LoopMomenta<F<T>>> for &LoopMomenta<F<T>> {
    type Output = LoopMomenta<F<T>>;

    fn add(self, rhs: &LoopMomenta<F<T>>) -> Self::Output {
        LoopMomenta::from_iter(self.iter().zip(rhs.iter()).map(|(l, r)| l + r))
    }
}

// define_indexed_vec!(
//     MyIdx;
//     pub struct ExternalMomentumtwo<ThreeMomentum<T>>;
// );
pub type ExternalThreeMomenta<T> = TiVec<ExternalIndex, ThreeMomentum<T>>;
// define_indexed_vec!()
pub type ExternalFourMomenta<T> = TiVec<ExternalIndex, FourMomentum<T>>;
// impl<T: Display> Display for ExternalFourMomenta<T> {
//     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
//         let mut table_builder = tabled::builder::Builder::new();
//         for (i, mom) in self.iter_enumerated() {
//             let mut row = Vec::new();
//             row.push(i.to_string());
//             for m in mom {
//                 row.push(m.to_string());
//             }
//             table_builder.push_record(row);
//         }

//         write!(f, "{}", table_builder.build().with(Style::modern_rounded()))
//     }
// }
pub type PolarizationVectors<T> = TiVec<ExternalIndex, Polarization<T>>; // should be the same length as #externals

fn extract_external_spatial<T: Clone>(
    external_four_momenta: &ExternalFourMomenta<T>,
) -> ExternalThreeMomenta<T> {
    ExternalThreeMomenta::from_iter(external_four_momenta.iter().map(|fm| fm.spatial.clone()))
}

#[derive(Debug, Clone)]
pub struct BareMomentumSample<T: FloatLike> {
    pub loop_moms: LoopMomenta<F<T>>,
    pub dual_loop_moms: Option<LoopMomenta<HyperDual<F<T>>>>,
    pub loop_mom_cache_id: usize,
    /// Base cache ID for the fundamental loop momentum configuration (before transformations)
    pub loop_mom_base_cache_id: usize,
    pub external_moms: ExternalFourMomenta<F<T>>,
    pub external_mom_cache_id: usize,
    /// Base cache ID for the fundamental external momentum configuration (before transformations)
    pub external_mom_base_cache_id: usize,
    pub jacobian: F<T>,
    pub orientation: Option<usize>,
    pub parameterization_branch: Option<usize>,
}

#[derive(Debug, Clone)]
pub struct MomentumSample<T: FloatLike> {
    pub sample: BareMomentumSample<T>,
    // pub uuid: Uuid,
}

impl MomentumSample<ArbPrec> {
    /// Materialize the retained canonical point directly in a physical lane.
    /// This is not a cast of a previously evaluated lower-precision sample.
    pub(crate) fn materialize<T: FloatLike>(&self) -> Result<MomentumSample<T>> {
        let sample = &self.sample;
        if sample.dual_loop_moms.is_some() {
            return Err(eyre!(
                "canonical draw materialization requires scalar source momenta"
            ));
        }
        let spatial = |p: &ThreeMomentum<F<ArbPrec>>| -> Result<ThreeMomentum<F<T>>> {
            Ok(ThreeMomentum::new(
                F::from_arb(&p.px.0)?,
                F::from_arb(&p.py.0)?,
                F::from_arb(&p.pz.0)?,
            ))
        };
        Ok(MomentumSample {
            sample: BareMomentumSample {
                loop_moms: sample
                    .loop_moms
                    .iter()
                    .map(spatial)
                    .collect::<Result<_>>()?,
                dual_loop_moms: None,
                loop_mom_cache_id: sample.loop_mom_cache_id,
                loop_mom_base_cache_id: sample.loop_mom_base_cache_id,
                external_moms: sample
                    .external_moms
                    .iter()
                    .map(|p| {
                        Ok(FourMomentum {
                            temporal: crate::momentum::Energy {
                                value: F::from_arb(&p.temporal.value.0)?,
                            },
                            spatial: spatial(&p.spatial)?,
                        })
                    })
                    .collect::<Result<_>>()?,
                external_mom_cache_id: sample.external_mom_cache_id,
                external_mom_base_cache_id: sample.external_mom_base_cache_id,
                jacobian: F::from_arb(&sample.jacobian.0)?,
                orientation: sample.orientation,
                parameterization_branch: sample.parameterization_branch,
            },
        })
    }
}

impl<T: FloatLike> Display for MomentumSample<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let mut table = tabled::builder::Builder::new();

        table.push_record(["Sample"]);

        table.push_record(["Loop Momenta", "p_x", "p_y", "p_z"]);
        // table

        for (index, loop_mom) in self.loop_moms().0.iter().enumerate() {
            table.push_record([
                index.to_string(),
                loop_mom.px.to_string(),
                loop_mom.py.to_string(),
                loop_mom.pz.to_string(),
            ]);
        }

        table.push_record(["External Momenta", "E", "p_x", "p_y", "p_z"]);
        for (index, external_mom) in self.external_moms().iter_enumerated() {
            table.push_record([
                index.to_string(),
                external_mom.temporal.to_string(),
                external_mom.spatial.px.to_string(),
                external_mom.spatial.py.to_string(),
                external_mom.spatial.pz.to_string(),
            ]);
        }

        table.push_record(["Jacobian".into(), format!("{:+e}", self.sample.jacobian)]);
        table.build().with(Style::rounded()).fmt(f)
    }
}

impl<T: FloatLike> BareMomentumSample<T> {
    #[inline(never)]
    pub(crate) fn new(
        loop_moms: LoopMomenta<F<T>>,
        loop_mom_cache_id: usize,
        external_moms: &Externals,
        external_mom_cache_id: usize,
        jacobian: F<T>,
        dependent_momenta_constructor: DependentMomentaConstructor,
        orientation: Option<usize>,
    ) -> Result<Self> {
        let external_moms = external_moms.get_dependent_externals(dependent_momenta_constructor)?;

        Ok(Self {
            loop_moms,
            dual_loop_moms: None,
            loop_mom_cache_id,
            loop_mom_base_cache_id: loop_mom_cache_id, // Initially same as cache_id
            external_mom_cache_id,
            external_mom_base_cache_id: external_mom_cache_id, // Initially same as cache_id
            external_moms,
            jacobian,
            orientation,
            parameterization_branch: None,
        })
    }

    pub(crate) fn one(&self) -> F<T> {
        if let Some(f) = self.loop_moms.first() {
            f.px.one()
        } else if let Some(f) = self.external_moms.first() {
            f.spatial.px.one()
        } else {
            panic!("No momenta in sample")
        }
    }

    pub(crate) fn zero(&self) -> F<T> {
        if let Some(f) = self.loop_moms.first() {
            f.px.zero()
        } else if let Some(f) = self.external_moms.first() {
            f.spatial.px.zero()
        } else {
            panic!("No momenta in sample")
        }
    }

    /// Cast the sample to a different precision
    #[inline]
    fn cast_sample<T2: FloatLike>(&self) -> BareMomentumSample<T2>
    where
        F<T2>: From<F<T>>,
    {
        BareMomentumSample {
            loop_mom_cache_id: self.loop_mom_cache_id,
            loop_mom_base_cache_id: self.loop_mom_base_cache_id,
            external_mom_cache_id: self.external_mom_cache_id,
            external_mom_base_cache_id: self.external_mom_base_cache_id,
            loop_moms: self.loop_moms.iter().map(ThreeMomentum::cast).collect(),
            dual_loop_moms: self
                .dual_loop_moms
                .as_ref()
                .map(|_dlm| todo!("make sure the cast works if there are hyperdual momenta")),
            external_moms: self.external_moms.iter().map(FourMomentum::cast).collect(),
            jacobian: self.jacobian.clone().into(),
            orientation: self.orientation,
            parameterization_branch: self.parameterization_branch,
        }
    }

    pub(crate) fn higher_precision(&self) -> BareMomentumSample<T::Higher>
    where
        T::Higher: FloatLike,
        T::Lower: FloatLike,
    {
        BareMomentumSample {
            loop_moms: self.loop_moms.iter().map(ThreeMomentum::higher).collect(),
            dual_loop_moms: self.dual_loop_moms.as_ref().map(|dlm| {
                LoopMomenta::from_iter(dlm.iter().map(|m| m.clone().map(&|x| x.higher())))
            }),
            external_moms: self
                .external_moms
                .iter()
                .map(FourMomentum::higher)
                .collect(),
            loop_mom_cache_id: self.loop_mom_cache_id,
            loop_mom_base_cache_id: self.loop_mom_base_cache_id,
            external_mom_cache_id: self.external_mom_cache_id,
            external_mom_base_cache_id: self.external_mom_base_cache_id,
            jacobian: self.jacobian.higher(),
            orientation: self.orientation,
            parameterization_branch: self.parameterization_branch,
        }
    }

    pub(crate) fn lower_precision(&self) -> BareMomentumSample<T::Lower>
    where
        T::Higher: FloatLike,
        T::Lower: FloatLike,
    {
        BareMomentumSample {
            loop_moms: self.loop_moms.iter().map(ThreeMomentum::lower).collect(),
            dual_loop_moms: self.dual_loop_moms.as_ref().map(|dlm| {
                LoopMomenta::from_iter(dlm.iter().map(|m| m.clone().map(&|x| x.lower())))
            }),
            external_moms: self.external_moms.iter().map(FourMomentum::lower).collect(),
            jacobian: self.jacobian.lower(),
            orientation: self.orientation,
            loop_mom_cache_id: self.loop_mom_cache_id,
            loop_mom_base_cache_id: self.loop_mom_base_cache_id,
            external_mom_cache_id: self.external_mom_cache_id,
            external_mom_base_cache_id: self.external_mom_base_cache_id,
            parameterization_branch: self.parameterization_branch,
        }
    }

    #[inline]
    pub(crate) fn rotate(
        &self,
        rotation: &Rotation,
        loop_mom_cache_id: usize,
        external_mom_cache_id: usize,
    ) -> Self {
        Self {
            loop_moms: self.loop_moms.iter().map(|l| l.rotate(rotation)).collect(),
            dual_loop_moms: self
                .dual_loop_moms
                .as_ref()
                .map(|dlm| LoopMomenta::from_iter(dlm.iter().map(|l| l.rotate(rotation)))),
            external_moms: self
                .external_moms
                .iter()
                .map(|l| l.rotate(rotation))
                .collect(),
            loop_mom_cache_id,
            loop_mom_base_cache_id: self.loop_mom_base_cache_id, // Preserve base cache ID
            external_mom_cache_id,
            external_mom_base_cache_id: self.external_mom_base_cache_id, // Preserve base cache ID
            jacobian: self.jacobian.clone(),
            orientation: self.orientation,
            parameterization_branch: self.parameterization_branch,
        }
    }

    #[inline]
    #[allow(dead_code)]
    pub(crate) fn rescaled_loop_momenta(&self, factor: &F<T>, subspace: Subspace) -> Self {
        Self {
            loop_moms: self.loop_moms.rescale(factor, subspace),
            dual_loop_moms: self
                .dual_loop_moms
                .as_ref()
                .map(|dlm| dlm.rescale(&new_constant(&dlm[LoopIndex(0)].px, factor), subspace)),
            loop_mom_cache_id: self.loop_mom_cache_id + 1,
            loop_mom_base_cache_id: self.loop_mom_base_cache_id, // Preserve base cache ID
            external_moms: self.external_moms.clone(),
            external_mom_cache_id: self.external_mom_cache_id,
            external_mom_base_cache_id: self.external_mom_base_cache_id, // Preserve base cache ID
            jacobian: self.jacobian.clone(),
            orientation: self.orientation,
            parameterization_branch: self.parameterization_branch,
        }
    }

    #[inline]
    pub(crate) fn lmb_transform(&self, from: &LoopMomentumBasis, to: &LoopMomentumBasis) -> Self {
        Self {
            loop_moms: self.loop_moms.lmb_transform(
                from,
                to,
                &extract_external_spatial(&self.external_moms),
            ),
            dual_loop_moms: self.dual_loop_moms.as_ref().map(|dlm| {
                let dual_externals = self
                    .external_moms
                    .iter()
                    .map(|four_mom| {
                        ThreeMomentum::new(
                            new_constant(&dlm[LoopIndex(0)].px, &four_mom.spatial.px),
                            new_constant(&dlm[LoopIndex(0)].px, &four_mom.spatial.py),
                            new_constant(&dlm[LoopIndex(0)].px, &four_mom.spatial.pz),
                        )
                    })
                    .collect();
                dlm.lmb_transform(from, to, &dual_externals)
            }),
            loop_mom_cache_id: self.loop_mom_cache_id + 1,
            loop_mom_base_cache_id: self.loop_mom_base_cache_id, // Preserve base cache ID
            external_moms: self.external_moms.clone(),
            external_mom_cache_id: self.external_mom_cache_id,
            external_mom_base_cache_id: self.external_mom_base_cache_id, // Preserve base cache ID
            jacobian: self.jacobian.clone(),
            orientation: self.orientation,
            parameterization_branch: self.parameterization_branch,
        }
    }
}

impl<T: FloatLike> MomentumSample<T> {
    pub(crate) fn orientations<'a, OID: From<usize> + Copy + IndexLike>(
        &self,
        filter: &'a SubSet<OID>,
        orientations: &'a TiVec<OID, EdgeVec<Orientation>>,
    ) -> SingleOrAllOrientations<'a, OID>
    where
        usize: From<OID>,
    {
        if let Some(o) = self.sample.orientation {
            let id = if filter.is_full() {
                OID::from(o)
            } else {
                filter
                    .included_iter()
                    .nth(o)
                    .unwrap_or_else(|| {
                        panic!(
                            "sampled orientation index {o} must resolve within the filtered orientation subset"
                        )
                    })
            };
            SingleOrAllOrientations::Single {
                id,
                orientation: &orientations[id],
            }
        } else {
            SingleOrAllOrientations::All {
                all: orientations,
                filter,
            }
        }
    }

    // pub(crate) fn uuid(&self) -> Option<Uuid> {
    //     if self.rotated_sample.is_some() {
    //         None
    //     } else {
    //         Some(self.uuid)
    //     }
    // }

    pub(crate) fn loop_moms(&self) -> &LoopMomenta<F<T>> {
        &self.sample.loop_moms
    }

    pub(crate) fn external_moms(&self) -> &ExternalFourMomenta<F<T>> {
        &self.sample.external_moms
    }

    pub(crate) fn jacobian(&self) -> F<T> {
        self.sample.jacobian.clone()
    }

    pub(crate) fn new(
        loop_moms: LoopMomenta<F<T>>,
        loop_mom_cache_id: usize,
        external_moms: &Externals,
        external_mom_cache_id: usize,
        jacobian: F<T>,
        dependent_momenta_constructor: DependentMomentaConstructor,
        orientation: Option<usize>,
    ) -> Result<Self> {
        // info!("New with ext cache id:{external_mom_cache_id}");
        Ok(Self {
            sample: BareMomentumSample::new(
                loop_moms,
                loop_mom_cache_id,
                external_moms,
                external_mom_cache_id,
                jacobian,
                dependent_momenta_constructor,
                orientation,
            )?,
            // uuid: Uuid::new_v4(),
        })
    }

    pub(crate) fn one(&self) -> F<T> {
        self.sample.one()
    }

    pub(crate) fn zero(&self) -> F<T> {
        self.sample.zero()
    }

    /// Cast the sample to a different precision
    #[inline]
    pub(crate) fn cast_sample<T2: FloatLike>(&self) -> MomentumSample<T2>
    where
        F<T2>: From<F<T>>,
    {
        MomentumSample {
            sample: self.sample.cast_sample(),
        }
    }

    pub(crate) fn higher_precision(&self) -> MomentumSample<T::Higher>
    where
        T::Higher: FloatLike + Default,
        T::Lower: FloatLike + Default,
    {
        MomentumSample {
            sample: self.sample.higher_precision(),
        }
    }

    pub(crate) fn lower_precision(&self) -> MomentumSample<T::Lower>
    where
        T::Lower: FloatLike + Default,
        T::Higher: FloatLike + Default,
    {
        MomentumSample {
            sample: self.sample.lower_precision(),
        }
    }

    #[inline]
    pub(crate) fn rotate(
        &self,
        rotation: &Rotation,
        loop_mom_cache_id: usize,
        external_mom_cache_id: usize,
    ) -> Self {
        Self {
            sample: self
                .sample
                .rotate(rotation, loop_mom_cache_id, external_mom_cache_id),
        }
    }

    #[allow(dead_code)]
    #[inline]
    pub(crate) fn rescaled_loop_momenta(&self, factor: &F<T>, subspace: Subspace) -> Self {
        Self {
            sample: self.sample.rescaled_loop_momenta(factor, subspace),
        }
    }

    #[inline]
    pub(crate) fn lmb_transform(&self, from: &LoopMomentumBasis, to: &LoopMomentumBasis) -> Self {
        Self {
            sample: self.sample.lmb_transform(from, to),
        }
    }
}

#[cfg(test)]
mod subspace_tests {
    use super::*;
    use crate::{
        dot,
        graph::{lmb::LMBwithEdges, parse::from_dot::IntoGraph},
        initialisation::test_initialise,
    };
    use typed_index_collections::ti_vec;

    #[test]
    fn canonical_momentum_materialization_preserves_source_metadata() {
        let one = F::<ArbPrec>::default().one();
        let mut source = MomentumSample {
            sample: BareMomentumSample {
                loop_moms: LoopMomenta::from_iter([ThreeMomentum::new(
                    one.clone(),
                    -&one,
                    one.zero(),
                )]),
                dual_loop_moms: None,
                loop_mom_cache_id: 17,
                loop_mom_base_cache_id: 11,
                external_moms: ti_vec![FourMomentum::from_args(
                    one.from_i64(4),
                    one.zero(),
                    one.clone(),
                    one.zero()
                )],
                external_mom_cache_id: 23,
                external_mom_base_cache_id: 19,
                jacobian: one.from_i64(3),
                orientation: Some(7),
                parameterization_branch: Some(2),
            },
        };
        fn check<T: FloatLike>(source: &MomentumSample<ArbPrec>) {
            let native = source.materialize::<T>().unwrap();
            assert_eq!(native.sample.loop_mom_cache_id, 17);
            assert_eq!(native.sample.loop_mom_base_cache_id, 11);
            assert_eq!(native.sample.external_mom_cache_id, 23);
            assert_eq!(native.sample.external_mom_base_cache_id, 19);
            assert_eq!(native.sample.orientation, Some(7));
            assert_eq!(native.sample.parameterization_branch, Some(2));
            assert_eq!(native.sample.jacobian, native.one().from_i64(3));
            assert_eq!(native.sample.loop_moms[LoopIndex(0)].py, -native.one());
        }
        check::<f64>(&source);
        check::<crate::utils::QuadFloat>(&source);
        check::<ArbPrec>(&source);
        source.sample.dual_loop_moms = Some(LoopMomenta(vec![]));
        assert!(
            source
                .materialize::<f64>()
                .unwrap_err()
                .to_string()
                .contains("scalar source")
        );
    }

    #[test]
    fn selected_cycle_identity_and_radial_projection_survive_parent_changes() {
        test_initialise().unwrap();
        let graph: Graph = dot!(digraph selected_cycles {
            ext [style=invis]
            edge [num=1 mass=0]
            node [num=1]
            ext->a:0 [id=0]
            a->b [id=1 name=a1]
            b->a [id=2 name=a2]
            a->b [id=3 name=a3]
            b->c [id=4 name=b1]
            c->b [id=5 name=b2]
            b->c [id=6 name=b3]
            ext->c:1 [id=7]
        })
        .unwrap();
        let edges = ["a1", "a2", "a3", "b1", "b2", "b3"]
            .map(|name| graph.edge_name_to_index(name).unwrap());
        let parent = graph
            .lmb_with_loop_edges([edges[0], edges[2], edges[3], edges[5]].as_slice())
            .unwrap();
        let mut alternate = graph
            .lmb_with_loop_edges([edges[0], edges[2], edges[4], edges[5]].as_slice())
            .unwrap();
        let old_slot = alternate
            .loop_edges
            .iter_enumerated()
            .find_map(|(index, edge)| (*edge == edges[0]).then_some(index))
            .unwrap();
        alternate.swap_loops(
            old_slot,
            LoopIndex((old_slot.0 + 1) % alternate.loop_edges.len()),
        );
        let different_cycle = graph
            .lmb_with_loop_edges([edges[0], edges[1], edges[3], edges[5]].as_slice())
            .unwrap();
        let lmbs = ti_vec![parent, alternate, different_cycle];
        let subspaces = (0..3)
            .map(|index| {
                SubspaceData::new_from_parent_basis_edges(
                    &[edges[0]],
                    &graph.full_filter(),
                    LmbIndex::from(index),
                    &graph,
                    &lmbs,
                )
                .unwrap()
            })
            .collect::<Vec<_>>();
        assert_ne!(subspaces[0].lmb_indices, subspaces[1].lmb_indices);
        assert_eq!(
            subspaces[0].solve_signature(&lmbs),
            subspaces[1].solve_signature(&lmbs)
        );
        assert!(subspaces[0].has_equivalent_embedding(&subspaces[1], &lmbs));
        assert_ne!(
            subspaces[0].solve_signature(&lmbs),
            subspaces[2].solve_signature(&lmbs)
        );
        assert!(!subspaces[0].has_equivalent_embedding(&subspaces[2], &lmbs));
        let independent = SubspaceData::new_from_parent_basis_edges(
            &[edges[5]],
            &graph.full_filter(),
            LmbIndex::from(1),
            &graph,
            &lmbs,
        )
        .unwrap();
        assert!(subspaces[0].is_mergable_with(&independent, &lmbs));
        assert!(!subspaces[0].is_mergable_with(&subspaces[2], &lmbs));

        let momenta = LoopMomenta::from_iter(
            (0..4).map(|i| ThreeMomentum::new(F(i as f64 + 1.0), F(2.0), F(-3.0))),
        );
        let externals = [
            ThreeMomentum::new(F(3.0), F(2.0), F(1.0)),
            ThreeMomentum::new(F(-3.0), F(-2.0), F(-1.0)),
        ]
        .into_iter()
        .collect();
        let transformed = momenta.lmb_transform(
            &lmbs[LmbIndex::from(0)],
            &lmbs[LmbIndex::from(1)],
            &externals,
        );
        let mut centers = [momenta.clone(), momenta.clone()];
        for (center, subspace) in centers.iter_mut().zip(&subspaces) {
            for momentum in &mut center.0 {
                *momentum = ThreeMomentum::new(F(0.0), F(0.0), F(0.0));
            }
            center[subspace.lmb_indices[0]] = ThreeMomentum::new(F(4.0), F(-5.0), F(6.0));
        }
        let shifted = &momenta - &centers[0];
        let alternate_shifted = &transformed - &centers[1];
        assert_eq!(
            shifted.hyper_radius_squared(subspaces[0].as_subspace_simple()),
            alternate_shifted.hyper_radius_squared(subspaces[1].as_subspace_simple()),
        );
        let projected = &shifted.rescale(&F(0.37), subspaces[0].as_subspace_simple()) + &centers[0];
        let alternate_projected =
            &alternate_shifted.rescale(&F(0.37), subspaces[1].as_subspace_simple()) + &centers[1];
        let expected = projected.lmb_transform(
            &lmbs[LmbIndex::from(0)],
            &lmbs[LmbIndex::from(1)],
            &externals,
        );
        for (actual, expected) in alternate_projected.iter().zip(expected.iter()) {
            assert!((actual - expected).norm_squared() < F(1e-24));
        }
    }

    #[test]
    fn explicit_parent_basis_edges_recover_their_fundamental_cycle_subspace() {
        test_initialise().unwrap();
        let graph: Graph = dot!(digraph explicit_basis_subspace {
            ext [style=invis]
            edge [num=1 mass=0]
            node [num=1]
            ext->v1:0 [id=0]
            v1->v2 [id=1]
            v2->v1 [id=2]
            v1->v2 [id=3]
            v2->v1 [id=4]
            ext->v2:1 [id=5]
        })
        .unwrap();
        let all_lmbs = ti_vec![graph.loop_momentum_basis.clone()];
        let requested = [
            graph.loop_momentum_basis.loop_edges[LoopIndex(0)],
            graph.loop_momentum_basis.loop_edges[LoopIndex(2)],
        ];
        let containing_subgraph: SuBitGraph = graph.full_filter();

        let subspace = SubspaceData::new_from_parent_basis_edges(
            &requested,
            &containing_subgraph,
            LmbIndex::from(0),
            &graph,
            &all_lmbs,
        )
        .unwrap();

        assert_eq!(subspace.loopcount(), requested.len());
        assert_eq!(
            subspace
                .iter_basis_edges(&all_lmbs)
                .collect::<std::collections::BTreeSet<_>>(),
            requested
                .into_iter()
                .collect::<std::collections::BTreeSet<_>>(),
        );
        let nested_subspace = SubspaceData::new_from_parent_basis_edges(
            &[requested[0]],
            &containing_subgraph,
            LmbIndex::from(0),
            &graph,
            &all_lmbs,
        )
        .unwrap();
        let common_subspace =
            SubspaceData::union_in_common_parent([&nested_subspace, &subspace], &graph, &all_lmbs)
                .unwrap();
        assert_eq!(
            common_subspace
                .iter_basis_edges(&all_lmbs)
                .collect::<std::collections::BTreeSet<_>>(),
            requested
                .into_iter()
                .collect::<std::collections::BTreeSet<_>>(),
        );
        let disjoint_subspace = SubspaceData::new_from_parent_basis_edges(
            &[requested[1]],
            &containing_subgraph,
            LmbIndex::from(0),
            &graph,
            &all_lmbs,
        )
        .unwrap();
        let sparse_common_subspace = SubspaceData::union_in_common_parent(
            [&nested_subspace, &disjoint_subspace],
            &graph,
            &all_lmbs,
        )
        .unwrap();
        assert_eq!(sparse_common_subspace.loopcount(), 2);
        assert_eq!(
            sparse_common_subspace
                .iter_basis_edges(&all_lmbs)
                .collect::<std::collections::BTreeSet<_>>(),
            requested
                .into_iter()
                .collect::<std::collections::BTreeSet<_>>(),
        );
        assert!(
            SubspaceData::new_from_parent_basis_edges(
                &[requested[0], requested[0]],
                &containing_subgraph,
                LmbIndex::from(0),
                &graph,
                &all_lmbs,
            )
            .unwrap_err()
            .to_string()
            .contains("repeats basis edge")
        );
    }
}