gammalooprs/momentum/

mod.rs

#![allow(dead_code)]

pub mod sample;
pub mod signature;

use std::{
    borrow::Borrow,
    fmt::{Display, LowerExp},
    ops::{Add, AddAssign, Index, Mul, MulAssign, Neg, Sub, SubAssign},
    str::FromStr,
};

use bincode_trait_derive::{Decode, Encode};
use eyre::Context;
use linnet::half_edge::involution::EdgeIndex;
pub use linnet::num_traits::{Pow, Sign, SignError, SignOrZero};
use momtrop::vector::Vector;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use spenso::{
    algebra::{algebraic_traits::RefZero, complex::Complex, upgrading_arithmetic::FallibleAdd},
    contraction::Contract,
    iterators::IteratableTensor,
    shadowing::Shadowable,
    structure::{
        Canonicalized, CastStructure, IndexLess, NamedStructure, OrderedStructure,
        TensorDataLayout, TensorStructure, ToSymbolic,
        abstract_index::AbstractIndex,
        representation::{BaseRepName, Euclidean, LibraryRep, Minkowski, RepName},
        slot::{DualSlotTo, Slot},
    },
    tensors::{
        complex::RealOrComplexTensor,
        data::{
            DataIterator, DataTensor, DenseTensor, HasTensorData, SetTensorData, SparseTensor,
            StorageTensor,
        },
        parametric::{
            EvalTensor, FlatCoefficent, MixedTensor, ParamOrConcrete, atomcore::TensorAtomOps,
        },
    },
};
use symbolica::{
    domains::{
        dual::HyperDual,
        float::{Complex as SymComplex, FloatLike as SymFloatLike},
        rational::RationalField,
    },
    prelude::*,
};
use symbolica_utils::NoArgs;
use thiserror::Error;

use crate::{
    GammaLoopContext,
    settings::runtime::RotationSetting,
    utils::{
        ApproxEq, F, FloatLike, RefDefault, hyperdual_utils::new_constant, representations::GR,
    },
};

#[derive(Debug, PartialEq, Eq, Clone, Copy, Serialize, Deserialize, Encode, Decode)]
pub struct Energy<T> {
    pub value: T,
}

impl<T> Energy<T> {
    pub(crate) fn map_ref<U>(&self, f: &impl Fn(&T) -> U) -> Energy<U> {
        Energy {
            value: f(&self.value),
        }
    }

    pub(crate) fn map<U>(self, f: &impl Fn(T) -> U) -> Energy<U> {
        Energy {
            value: f(self.value),
        }
    }
}

impl<T: FloatLike> ApproxEq<Energy<F<T>>, F<T>> for Energy<F<T>> {
    fn approx_eq(&self, other: &Energy<F<T>>, threshold: &F<T>) -> bool {
        self.value.approx_eq(&other.value, threshold)
    }
}

impl<T: FloatLike> Energy<F<T>> {
    pub(crate) fn higher(&self) -> Energy<F<T::Higher>>
    where
        T::Higher: FloatLike,
    {
        Energy {
            value: self.value.higher(),
        }
    }

    pub(crate) fn lower(&self) -> Energy<F<T::Lower>>
    where
        T::Lower: FloatLike,
    {
        Energy {
            value: self.value.lower(),
        }
    }

    pub(crate) fn from_ff64(energy: Energy<F<f64>>) -> Self {
        Energy {
            value: F::from_ff64(energy.value),
        }
    }
}

impl<T: FloatLike> From<Energy<T>> for Energy<F<T>> {
    fn from(value: Energy<T>) -> Self {
        Energy {
            value: F(value.value),
        }
    }
}

impl<T: Real> Energy<T> {
    pub(crate) fn zero(&self) -> Self {
        Energy {
            value: self.value.zero(),
        }
    }
}

impl<T> Add<Energy<T>> for Energy<T>
where
    T: Add<T, Output = T>,
{
    type Output = Energy<T>;
    fn add(self, rhs: Energy<T>) -> Self::Output {
        Energy {
            value: self.value + rhs.value,
        }
    }
}

impl<T> Add<&Energy<T>> for Energy<T>
where
    T: for<'a> Add<&'a T, Output = T>,
{
    type Output = Energy<T>;
    fn add(self, rhs: &Energy<T>) -> Self::Output {
        Energy {
            value: self.value + &rhs.value,
        }
    }
}

impl<'b, T> Add<&Energy<T>> for &'b Energy<T>
where
    &'b T: for<'a> Add<&'a T, Output = T>,
{
    type Output = Energy<T>;
    fn add(self, rhs: &Energy<T>) -> Self::Output {
        Energy {
            value: &self.value + &rhs.value,
        }
    }
}

impl<T> Add<Energy<T>> for &Energy<T>
where
    T: for<'a> Add<&'a T, Output = T>,
{
    type Output = Energy<T>;
    fn add(self, rhs: Energy<T>) -> Self::Output {
        rhs + self
    }
}

impl<T> AddAssign<Energy<T>> for Energy<T>
where
    T: AddAssign<T>,
{
    fn add_assign(&mut self, rhs: Energy<T>) {
        self.value += rhs.value;
    }
}

impl<'a, T> AddAssign<&'a Energy<T>> for Energy<T>
where
    T: AddAssign<&'a T>,
{
    fn add_assign(&mut self, rhs: &'a Energy<T>) {
        self.value += &rhs.value;
    }
}

impl<T> Sub<Energy<T>> for Energy<T>
where
    T: Sub<T, Output = T>,
{
    type Output = Energy<T>;
    fn sub(self, rhs: Energy<T>) -> Self::Output {
        Energy {
            value: self.value - rhs.value,
        }
    }
}

impl<T> SubAssign<Energy<T>> for Energy<T>
where
    T: SubAssign<T>,
{
    fn sub_assign(&mut self, rhs: Energy<T>) {
        self.value -= rhs.value;
    }
}

impl<'a, T> SubAssign<&'a Energy<T>> for Energy<T>
where
    T: SubAssign<&'a T>,
{
    fn sub_assign(&mut self, rhs: &'a Energy<T>) {
        self.value -= &rhs.value;
    }
}
impl<T> Mul<Energy<T>> for Energy<T>
where
    T: Mul<T, Output = T>,
{
    type Output = T;
    fn mul(self, rhs: Energy<T>) -> Self::Output {
        self.value * rhs.value
    }
}

impl<T> MulAssign<Energy<T>> for Energy<T>
where
    T: MulAssign<T>,
{
    fn mul_assign(&mut self, rhs: Energy<T>) {
        self.value *= rhs.value;
    }
}

impl<T> Neg for Energy<T>
where
    T: Neg<Output = T>,
{
    type Output = Energy<T>;
    fn neg(self) -> Self::Output {
        Energy { value: -self.value }
    }
}

impl<T> Energy<T> {
    pub(crate) fn new(value: T) -> Self {
        Energy { value }
    }

    // pub(crate) fn from_three_momentum(three_momentum: &ThreeMomentum<T>) -> Self
    // where
    //     T: std::ops::Mul<Output = T> + std::ops::Add<Output = T> + Copy,
    // {
    //     let px2 = three_momentum.px * three_momentum.px;
    //     let py2 = three_momentum.py * three_momentum.py;
    //     let pz2 = three_momentum.pz * three_momentum.pz;
    //     let p2 = px2 + py2 + pz2;
    //     let value = (p2 + T::default()).sqrt();
    //     Energy { value }
    // }
}

// impl<U: Borrow<T>, T> Borrow<Energy<T>> for Energy<U> {
//     fn borrow(&self) -> &Energy<T> {

//     }
// }

impl<U, T: RefZero<U>> RefZero<Energy<U>> for Energy<T>
where
    Energy<T>: Borrow<Energy<U>>,
{
    fn ref_zero(&self) -> Energy<U> {
        Energy {
            value: self.value.ref_zero(),
        }
    }
}

impl Energy<Atom> {
    pub(crate) fn new_parametric(id: usize) -> Self {
        let value = parse!(&format!("E_{}", id));
        Energy { value }
    }
}

impl<T: Display> Display for Energy<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "E: {}", self.value)
    }
}

impl<T: LowerExp> LowerExp for Energy<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "E: {:e}", self.value)
    }
}

impl<T: Default> Default for Energy<T> {
    fn default() -> Self {
        Energy {
            value: T::default(),
        }
    }
}

impl<T> From<T> for Energy<T> {
    fn from(value: T) -> Self {
        Energy { value }
    }
}

#[derive(Debug, PartialEq, Eq, Clone, Copy, Serialize, Deserialize, Encode, Decode)]
pub struct ThreeMomentum<T> {
    pub px: T,
    pub py: T,
    pub pz: T,
}

impl<T> ThreeMomentum<T> {
    pub(crate) fn map_ref<U>(&self, f: &impl Fn(&T) -> U) -> ThreeMomentum<U> {
        ThreeMomentum {
            px: f(&self.px),
            py: f(&self.py),
            pz: f(&self.pz),
        }
    }

    pub(crate) fn map<U>(self, f: &impl Fn(T) -> U) -> ThreeMomentum<U> {
        ThreeMomentum {
            px: f(self.px),
            py: f(self.py),
            pz: f(self.pz),
        }
    }
}

impl<T: FloatLike> ApproxEq<ThreeMomentum<F<T>>, F<T>> for ThreeMomentum<F<T>> {
    fn approx_eq(&self, other: &ThreeMomentum<F<T>>, threshold: &F<T>) -> bool {
        F::approx_eq_iterator(
            [&self.px, &self.py, &self.pz],
            [&other.px, &other.py, &other.pz],
            threshold,
        )
    }
}

pub struct ThreeRotation<T> {
    pub map: fn(ThreeMomentum<T>) -> ThreeMomentum<T>,
    pub inv_map: fn(ThreeMomentum<T>) -> ThreeMomentum<T>,
}

impl<T: Neg<Output = T>> ThreeRotation<T> {
    pub(crate) fn half_pi_x() -> Self {
        let map = |mut momentum: ThreeMomentum<T>| -> ThreeMomentum<T> {
            std::mem::swap(&mut momentum.py, &mut momentum.pz);
            momentum.pz = -momentum.pz;
            momentum
        };

        let inv_map = |mut momentum: ThreeMomentum<T>| -> ThreeMomentum<T> {
            momentum.pz = -momentum.pz;
            std::mem::swap(&mut momentum.py, &mut momentum.pz);
            momentum
        };

        ThreeRotation { map, inv_map }
    }

    pub(crate) fn half_pi_y() -> Self {
        let map = |mut momentum: ThreeMomentum<T>| -> ThreeMomentum<T> {
            std::mem::swap(&mut momentum.px, &mut momentum.pz);
            momentum.px = -momentum.px;
            momentum
        };

        let inv_map = |mut momentum: ThreeMomentum<T>| -> ThreeMomentum<T> {
            momentum.px = -momentum.px;
            std::mem::swap(&mut momentum.px, &mut momentum.pz);
            momentum
        };

        ThreeRotation { map, inv_map }
    }

    pub(crate) fn half_pi_z() -> Self {
        let map = |mut momentum: ThreeMomentum<T>| -> ThreeMomentum<T> {
            std::mem::swap(&mut momentum.px, &mut momentum.py);
            momentum.py = -momentum.py;
            momentum
        };

        let inv_map = |mut momentum: ThreeMomentum<T>| -> ThreeMomentum<T> {
            momentum.py = -momentum.py;
            std::mem::swap(&mut momentum.px, &mut momentum.py);
            momentum
        };

        ThreeRotation { map, inv_map }
    }
}

impl<T: FloatLike> ThreeMomentum<F<T>> {
    pub(crate) fn higher(&self) -> ThreeMomentum<F<T::Higher>>
    where
        T::Higher: FloatLike,
    {
        ThreeMomentum {
            px: self.px.higher(),
            py: self.py.higher(),
            pz: self.pz.higher(),
        }
    }

    pub(crate) fn lower(&self) -> ThreeMomentum<F<T::Lower>>
    where
        T::Lower: FloatLike,
    {
        ThreeMomentum {
            px: self.px.lower(),
            py: self.py.lower(),
            pz: self.pz.lower(),
        }
    }

    pub(crate) fn from_ff64(three_mom: ThreeMomentum<F<f64>>) -> Self {
        ThreeMomentum {
            px: F::from_ff64(three_mom.px),
            py: F::from_ff64(three_mom.py),
            pz: F::from_ff64(three_mom.pz),
        }
    }
}

impl<T: FloatLike> From<ThreeMomentum<T>> for ThreeMomentum<F<T>> {
    fn from(value: ThreeMomentum<T>) -> Self {
        ThreeMomentum {
            px: F(value.px),
            py: F(value.py),
            pz: F(value.pz),
        }
    }
}

impl<T> IntoIterator for ThreeMomentum<T> {
    type Item = T;
    type IntoIter = std::array::IntoIter<T, 3>;

    fn into_iter(self) -> Self::IntoIter {
        let [px, py, pz] = [self.px, self.py, self.pz];
        [px, py, pz].into_iter()
    }
}

impl<'a, T> IntoIterator for &'a ThreeMomentum<T> {
    type Item = &'a T;
    type IntoIter = std::array::IntoIter<&'a T, 3>;

    fn into_iter(self) -> Self::IntoIter {
        let [px, py, pz] = [&self.px, &self.py, &self.pz];
        [px, py, pz].into_iter()
    }
}

impl<T: Real> RefDefault for ThreeMomentum<T> {
    fn default(&self) -> Self {
        let zero = self.px.zero();
        ThreeMomentum {
            px: zero.clone(),
            py: zero.clone(),
            pz: zero.clone(),
        }
    }
}

impl<T: Real> ThreeMomentum<T> {
    pub(crate) fn zero(&self) -> Self {
        let zero = self.px.zero();
        ThreeMomentum {
            px: zero.clone(),
            py: zero.clone(),
            pz: zero.clone(),
        }
    }
}

impl<T: RefZero> RefZero for ThreeMomentum<T> {
    fn ref_zero(&self) -> Self {
        ThreeMomentum {
            px: self.px.ref_zero(),
            py: self.py.ref_zero(),
            pz: self.pz.ref_zero(),
        }
    }
}

impl<T> ThreeMomentum<T> {
    pub(crate) fn new(px: T, py: T, pz: T) -> Self {
        ThreeMomentum { px, py, pz }
    }

    pub(crate) fn into_dense(self, index: AbstractIndex) -> DenseTensor<T, OrderedStructure>
    where
        T: Clone,
    {
        let structure =
            Canonicalized::from_iter(vec![Euclidean {}.new_slot(3, index)]).into_canonical();
        DenseTensor::from_storage_data(vec![self.px, self.py, self.pz], structure).unwrap()
    }

    pub(crate) fn into_dense_param(self, index: AbstractIndex) -> DenseTensor<T, OrderedStructure>
    where
        T: Clone,
    {
        let structure =
            Canonicalized::from_iter(vec![Euclidean {}.new_slot(3, index)]).into_canonical();
        DenseTensor::from_storage_data(vec![self.px, self.py, self.pz], structure).unwrap()
    }
}

impl<T: FloatLike> ThreeMomentum<F<T>> {
    pub(crate) fn to_f64(&self) -> ThreeMomentum<F<f64>> {
        ThreeMomentum {
            px: F(self.px.to_f64()),
            py: F(self.py.to_f64()),
            pz: F(self.pz.to_f64()),
        }
    }
    /// Compute the phi-angle separation with p2.
    pub(crate) fn getdelphi(&self, p2: &ThreeMomentum<F<T>>) -> F<T> {
        let pt1 = self.pt();
        let pt2 = p2.pt();
        if pt1.is_zero() {
            return pt1.max_value();
        }
        if pt2.is_zero() {
            return pt2.max_value();
        }

        let mut tmp = self.px.clone() * &p2.px + self.py.clone() * &p2.py;

        tmp /= pt1 * pt2;
        if tmp.norm() > tmp.one() + tmp.epsilon() {
            panic!("Cosine larger than 1. in phase-space cuts.")
        }
        if tmp.norm() > tmp.one() {
            (tmp.clone() / tmp.norm()).acos()
        } else {
            tmp.acos()
        }
    }

    /// Compute the deltaR separation with momentum p2.
    #[inline]
    pub(crate) fn delta_r(&self, p2: &ThreeMomentum<F<T>>) -> F<T>
    where
        T: Real,
    {
        let delta_eta = self.pseudo_rap() - p2.pseudo_rap();
        let delta_phi = self.getdelphi(p2);
        (delta_eta.square() + delta_phi.square()).sqrt()
    }

    pub(crate) fn rotate_mut(&mut self, alpha: &F<T>, beta: &F<T>, gamma: &F<T>) {
        let sin_alpha = alpha.sin();
        let cos_alpha = alpha.cos();
        let sin_beta = beta.sin();
        let cos_beta = beta.cos();
        let sin_gamma = gamma.sin();
        let cos_gamma = gamma.cos();

        let px = self.px.clone();
        let py = self.py.clone();
        let pz = self.pz.clone();

        self.px = cos_gamma.clone() * &cos_beta * &px
            + (-(cos_alpha.clone()) * &sin_gamma + sin_alpha.clone() * &sin_beta * &cos_gamma)
                * &py
            + (sin_alpha.clone() * &sin_gamma + cos_alpha.clone() * &sin_beta * &cos_gamma) * &pz;

        self.py = sin_gamma.clone() * &cos_beta * &px
            + (cos_alpha.clone() * &cos_gamma + sin_alpha.clone() * &sin_beta * &sin_gamma) * &py
            + (-sin_alpha.clone() * &cos_gamma + cos_alpha.clone() * &sin_beta * &sin_gamma) * &pz;

        self.pz =
            -sin_beta * &px + cos_beta.clone() * &sin_alpha * &py + cos_alpha * &cos_beta * &pz;
    }

    /// Compute transverse momentum.
    #[inline]
    pub(crate) fn pt(&self) -> F<T> {
        (self.px.square() + self.py.square()).sqrt()
    }

    /// Compute pseudorapidity.
    #[inline]
    pub(crate) fn pseudo_rap(&self) -> F<T> {
        let pt = self.pt();
        if pt.less_than_epsilon() && self.pz.norm().less_than_epsilon() {
            if self.pz.positive() {
                return pt.max_value();
            } else {
                return pt.min_value();
            }
        }
        let th = pt.atan2(&self.pz);
        let two = pt.from_i64(2);
        -(th / two).tan().ln()
    }

    pub(crate) fn cross_product(&self, rhs: &ThreeMomentum<F<T>>) -> ThreeMomentum<F<T>> {
        ThreeMomentum {
            px: &self.py * &rhs.pz - &self.pz * &rhs.py,
            py: &self.pz * &rhs.px - &self.px * &rhs.pz,
            pz: &self.px * &rhs.py - &self.py * &rhs.px,
        }
    }

    // Rodriguez rotation formula
    pub(crate) fn axis_angle_rotation(
        &self,
        cos_theta: &F<T>,
        axis: &ThreeMomentum<F<T>>,
    ) -> ThreeMomentum<F<T>> {
        // ensure unit vector
        let axis = axis * &axis.norm().inv();
        let sin_theta = if cos_theta >= &cos_theta.one() {
            cos_theta.zero()
        } else {
            (cos_theta.one() - cos_theta.square()).sqrt()
        };
        let k_cross_v = axis.cross_product(self);
        let k_dot_v = self * axis.clone();

        self * cos_theta + k_cross_v * &sin_theta + &axis * &k_dot_v * &(self.px.one() - cos_theta)
    }

    pub(crate) fn get_cos_theta_with(&self, rhs: &ThreeMomentum<F<T>>) -> F<T> {
        let self_norm = self.norm();
        let rhs_norm = rhs.norm();
        if self_norm < self_norm.epsilon() || rhs_norm < rhs_norm.epsilon() {
            return self_norm.one();
        }

        let dot = self * rhs.clone();
        dot / (self.norm() * rhs.norm())
    }
}

impl<T: Neg<Output = T> + Clone> ThreeMomentum<T> {
    pub(crate) fn perform_pi2_rotation_x_mut(&mut self) {
        self.pz = -self.pz.clone();
        std::mem::swap(&mut self.pz, &mut self.py);
    }

    pub(crate) fn perform_pi2_rotation_x(&self) -> Self {
        // println!("X rotation");
        Self {
            px: self.px.clone(),
            py: -self.pz.clone(),
            pz: self.py.clone(),
        }
    }

    pub(crate) fn perform_pi2_rotation_y_mut(&mut self) {
        self.px = -self.px.clone();
        std::mem::swap(&mut self.px, &mut self.pz);
    }

    pub(crate) fn perform_pi2_rotation_y(&self) -> Self {
        // println!("Y rotation");
        Self {
            px: self.pz.clone(),
            py: self.py.clone(),
            pz: -self.px.clone(),
        }
    }

    pub(crate) fn perform_pi2_rotation_z_mut(&mut self) {
        self.py = -self.py.clone();
        std::mem::swap(&mut self.px, &mut self.py);
    }

    pub(crate) fn perform_pi2_rotation_z(&self) -> Self {
        // println!("Z rotation");
        Self {
            px: -self.py.clone(),
            py: self.px.clone(),
            pz: self.pz.clone(),
        }
    }
}

impl<T> Add<ThreeMomentum<T>> for ThreeMomentum<T>
where
    T: Add<T, Output = T>,
{
    type Output = ThreeMomentum<T>;
    fn add(self, rhs: ThreeMomentum<T>) -> Self::Output {
        ThreeMomentum {
            px: self.px + rhs.px,
            py: self.py + rhs.py,
            pz: self.pz + rhs.pz,
        }
    }
}

impl<T> Add<&ThreeMomentum<T>> for ThreeMomentum<T>
where
    T: for<'a> Add<&'a T, Output = T>,
{
    type Output = ThreeMomentum<T>;
    fn add(self, rhs: &ThreeMomentum<T>) -> Self::Output {
        ThreeMomentum {
            px: self.px + &rhs.px,
            py: self.py + &rhs.py,
            pz: self.pz + &rhs.pz,
        }
    }
}

impl<'b, T> Add<&ThreeMomentum<T>> for &'b ThreeMomentum<T>
where
    &'b T: for<'a> Add<&'a T, Output = T>,
{
    type Output = ThreeMomentum<T>;
    fn add(self, rhs: &ThreeMomentum<T>) -> Self::Output {
        ThreeMomentum {
            px: &self.px + &rhs.px,
            py: &self.py + &rhs.py,
            pz: &self.pz + &rhs.pz,
        }
    }
}

impl<T> Add<ThreeMomentum<T>> for &ThreeMomentum<T>
where
    T: for<'a> Add<&'a T, Output = T>,
{
    type Output = ThreeMomentum<T>;
    fn add(self, rhs: ThreeMomentum<T>) -> Self::Output {
        rhs + self
    }
}

impl<T> AddAssign<ThreeMomentum<T>> for ThreeMomentum<T>
where
    T: AddAssign<T>,
{
    fn add_assign(&mut self, rhs: ThreeMomentum<T>) {
        self.px += rhs.px;
        self.py += rhs.py;
        self.pz += rhs.pz;
    }
}

impl<'a, T> AddAssign<&'a ThreeMomentum<T>> for ThreeMomentum<T>
where
    T: AddAssign<&'a T>,
{
    fn add_assign(&mut self, rhs: &'a ThreeMomentum<T>) {
        self.px += &rhs.px;
        self.py += &rhs.py;
        self.pz += &rhs.pz;
    }
}

impl<T> Sub<ThreeMomentum<T>> for ThreeMomentum<T>
where
    T: Sub<T, Output = T>,
{
    type Output = ThreeMomentum<T>;
    fn sub(self, rhs: ThreeMomentum<T>) -> Self::Output {
        ThreeMomentum {
            px: self.px - rhs.px,
            py: self.py - rhs.py,
            pz: self.pz - rhs.pz,
        }
    }
}

impl<T> Sub<&ThreeMomentum<T>> for &ThreeMomentum<T>
where
    for<'a> &'a T: Sub<&'a T, Output = T>,
{
    type Output = ThreeMomentum<T>;
    fn sub(self, rhs: &ThreeMomentum<T>) -> Self::Output {
        ThreeMomentum {
            px: &self.px - &rhs.px,
            py: &self.py - &rhs.py,
            pz: &self.pz - &rhs.pz,
        }
    }
}

impl<T> SubAssign<ThreeMomentum<T>> for ThreeMomentum<T>
where
    T: SubAssign<T>,
{
    fn sub_assign(&mut self, rhs: ThreeMomentum<T>) {
        self.px -= rhs.px;
        self.py -= rhs.py;
        self.pz -= rhs.pz;
    }
}

impl<'a, T> SubAssign<&'a ThreeMomentum<T>> for ThreeMomentum<T>
where
    T: SubAssign<&'a T>,
{
    fn sub_assign(&mut self, rhs: &'a ThreeMomentum<T>) {
        self.px -= &rhs.px;
        self.py -= &rhs.py;
        self.pz -= &rhs.pz;
    }
}

impl<T> Mul<ThreeMomentum<T>> for ThreeMomentum<T>
where
    T: Mul<T, Output = T> + Add<T, Output = T>,
{
    type Output = T;
    fn mul(self, rhs: ThreeMomentum<T>) -> Self::Output {
        self.px * rhs.px + self.py * rhs.py + self.pz * rhs.pz
    }
}

impl<T> Mul<&ThreeMomentum<T>> for ThreeMomentum<T>
where
    T: for<'a> Mul<&'a T, Output = T> + Add<T, Output = T>,
{
    type Output = T;
    fn mul(self, rhs: &ThreeMomentum<T>) -> Self::Output {
        self.px * &rhs.px + self.py * &rhs.py + self.pz * &rhs.pz
    }
}

impl<T> Mul<ThreeMomentum<T>> for &ThreeMomentum<T>
where
    T: for<'b> Mul<&'b T, Output = T> + Add<T, Output = T>,
{
    type Output = T;
    fn mul(self, rhs: ThreeMomentum<T>) -> Self::Output {
        rhs * self
    }
}

impl<T> Mul<T> for ThreeMomentum<T>
where
    T: Mul<T, Output = T> + Clone,
{
    type Output = ThreeMomentum<T>;
    fn mul(self, rhs: T) -> Self::Output {
        ThreeMomentum {
            px: self.px * rhs.clone(),
            py: self.py * rhs.clone(),
            pz: self.pz * rhs,
        }
    }
}

impl<T> Mul<&T> for ThreeMomentum<T>
where
    T: for<'a> Mul<&'a T, Output = T> + Clone,
{
    type Output = ThreeMomentum<T>;
    fn mul(self, rhs: &T) -> Self::Output {
        ThreeMomentum {
            px: self.px * rhs,
            py: self.py * rhs,
            pz: self.pz * rhs,
        }
    }
}

impl<T> Mul<T> for &ThreeMomentum<T>
where
    T: Mul<T, Output = T> + Clone,
{
    type Output = ThreeMomentum<T>;
    fn mul(self, rhs: T) -> Self::Output {
        ThreeMomentum {
            px: self.px.clone() * rhs.clone(),
            py: self.py.clone() * rhs.clone(),
            pz: self.pz.clone() * rhs,
        }
    }
}

impl<T> Mul<&T> for &ThreeMomentum<T>
where
    T: for<'b> Mul<&'b T, Output = T> + Clone,
{
    type Output = ThreeMomentum<T>;
    fn mul(self, rhs: &T) -> Self::Output {
        ThreeMomentum {
            px: self.px.clone() * rhs,
            py: self.py.clone() * rhs,
            pz: self.pz.clone() * rhs,
        }
    }
}

impl<T> Neg for ThreeMomentum<T>
where
    T: Neg<Output = T>,
{
    type Output = ThreeMomentum<T>;
    fn neg(self) -> Self::Output {
        ThreeMomentum {
            px: -self.px,
            py: -self.py,
            pz: -self.pz,
        }
    }
}

impl<T> Neg for &ThreeMomentum<T>
where
    T: Neg<Output = T> + Clone,
{
    type Output = ThreeMomentum<T>;
    fn neg(self) -> Self::Output {
        ThreeMomentum {
            px: -self.px.clone(),
            py: -self.py.clone(),
            pz: -self.pz.clone(),
        }
    }
}

impl<T: Display> Display for ThreeMomentum<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "px: {}, py: {}, pz: {}", self.px, self.py, self.pz)
    }
}

impl<T: LowerExp> LowerExp for ThreeMomentum<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "px: {:e}, py: {:e}, pz: {:e}", self.px, self.py, self.pz)
    }
}

impl<T: Default> Default for ThreeMomentum<T> {
    fn default() -> Self {
        ThreeMomentum {
            px: T::default(),
            py: T::default(),
            pz: T::default(),
        }
    }
}

impl<T> ThreeMomentum<T> {
    pub(crate) fn norm_squared(&self) -> T
    where
        T: for<'a> Mul<&'a T, Output = T> + Add<T, Output = T> + Clone,
    {
        self.px.clone() * &self.px + self.py.clone() * &self.py + self.pz.clone() * &self.pz
    }

    pub(crate) fn on_shell_energy(&self, mass: Option<T>) -> Energy<T>
    where
        T: Mul<T, Output = T> + Add<T, Output = T> + Clone + std::ops::Add<Output = T> + Real,
    {
        let energy_squared = self.on_shell_energy_squared(mass);
        Energy {
            value: energy_squared.value.sqrt(),
        }
    }

    pub(crate) fn on_shell_energy_squared(&self, mass: Option<T>) -> Energy<T>
    where
        T: for<'a> Mul<&'a T, Output = T>
            + Add<T, Output = T>
            + Clone
            + std::ops::Add<Output = T>
            + Display,
    {
        let p2 = self.norm_squared();
        if let Some(mass) = mass {
            // println!("mass: {}", mass);
            Energy {
                value: p2 + mass.clone() * &mass,
            }
        } else {
            Energy { value: p2 }
        }
    }

    pub(crate) fn norm(&self) -> T
    where
        T: for<'a> Mul<&'a T, Output = T> + Add<T> + Real,
    {
        self.norm_squared().sqrt()
    }

    pub(crate) fn into_four_momentum_parametric(self, id: usize) -> FourMomentum<T, Atom> {
        let energy = Energy::new_parametric(id);
        FourMomentum {
            temporal: energy,
            spatial: self,
        }
    }

    pub(crate) fn into_on_shell_four_momentum(self, mass: Option<T>) -> FourMomentum<T, T>
    where
        T: Mul<T> + Add<T> + std::ops::Add<Output = T> + Real,
    {
        FourMomentum::new_on_shell(self, mass)
    }

    pub(crate) fn cast<U>(&self) -> ThreeMomentum<U>
    where
        T: Clone + Into<U>,
    {
        ThreeMomentum {
            px: (self.px.clone()).into(),
            py: (self.py.clone()).into(),
            pz: (self.pz.clone()).into(),
        }
    }

    #[allow(clippy::wrong_self_convention)]
    pub(crate) fn into_f64(&self) -> ThreeMomentum<f64>
    where
        T: FloatLike,
    {
        ThreeMomentum {
            px: self.px.to_f64(),
            py: self.py.to_f64(),
            pz: self.pz.to_f64(),
        }
    }
}

impl<T> From<[T; 3]> for ThreeMomentum<T> {
    fn from(data: [T; 3]) -> Self {
        let [px, py, pz] = data;
        ThreeMomentum { px, py, pz }
    }
}

impl<T> From<ThreeMomentum<T>> for [T; 3] {
    fn from(data: ThreeMomentum<T>) -> Self {
        [data.px, data.py, data.pz]
    }
}

impl<T> From<(T, T, T)> for ThreeMomentum<T> {
    fn from(data: (T, T, T)) -> Self {
        let (px, py, pz) = data;
        ThreeMomentum { px, py, pz }
    }
}

impl<T> From<ThreeMomentum<T>> for (T, T, T) {
    fn from(data: ThreeMomentum<T>) -> Self {
        (data.px, data.py, data.pz)
    }
}

#[derive(Default, Debug, PartialEq, Eq, Clone, Copy, Serialize, Deserialize, Encode, Decode)]
pub struct FourMomentum<T, U = T> {
    pub temporal: Energy<U>,
    pub spatial: ThreeMomentum<T>,
}

impl<T: FloatLike> Sub<&FourMomentum<F<T>>> for &FourMomentum<F<T>> {
    type Output = FourMomentum<F<T>>;
    fn sub(self, rhs: &FourMomentum<F<T>>) -> Self::Output {
        FourMomentum {
            temporal: Energy {
                value: &self.temporal.value - &rhs.temporal.value,
            },
            spatial: &self.spatial - &rhs.spatial,
        }
    }
}

impl<T> FourMomentum<T> {
    pub(crate) fn map_ref<U>(&self, f: &impl Fn(&T) -> U) -> FourMomentum<U> {
        FourMomentum {
            temporal: self.temporal.map_ref(f),
            spatial: self.spatial.map_ref(f),
        }
    }

    pub(crate) fn map<U>(self, f: &impl Fn(T) -> U) -> FourMomentum<U> {
        FourMomentum {
            temporal: self.temporal.map(f),
            spatial: self.spatial.map(f),
        }
    }
}

impl<T: FloatLike> ApproxEq<FourMomentum<F<T>>, F<T>> for FourMomentum<F<T>> {
    fn approx_eq(&self, other: &FourMomentum<F<T>>, threshold: &F<T>) -> bool {
        self.temporal.approx_eq(&other.temporal, threshold)
            && self.spatial.approx_eq(&other.spatial, threshold)
    }
}

impl<T: FloatLike> ApproxEq<Polarization<Complex<F<T>>>, F<T>> for FourMomentum<F<T>> {
    fn approx_eq(&self, other: &Polarization<Complex<F<T>>>, tolerance: &F<T>) -> bool {
        if other.tensor.size().unwrap() != 4 {
            false
        } else {
            self.into_iter()
                .zip(other.tensor.iter_flat())
                .all(|(a, (_, b))| a.approx_eq(b, tolerance))
        }
    }

    fn approx_eq_res(
        &self,
        other: &Polarization<Complex<F<T>>>,
        tolerance: &F<T>,
    ) -> color_eyre::Result<()> {
        if other.tensor.size().unwrap() != 4 {
            Err(eyre::eyre!("Polarization tensor has wrong size."))
        } else {
            self.into_iter()
                .zip(other.tensor.iter_flat())
                .try_for_each(|(a, (i, b))| {
                    a.approx_eq_res(b, tolerance).wrap_err(format!(
                        "Polarization tensor element {} does not match. ",
                        i
                    ))
                })
        }
    }
}

impl<T: FloatLike> FourMomentum<F<T>> {
    pub(crate) fn from_ff64(four_momentum: &FourMomentum<F<f64>>) -> Self {
        let temporal = Energy::from_ff64(four_momentum.temporal);
        let spatial = ThreeMomentum::from_ff64(four_momentum.spatial);
        FourMomentum { temporal, spatial }
    }

    pub(crate) fn rescale_spatial(&self, x: &F<T>) -> Self {
        FourMomentum {
            temporal: self.temporal.clone(),
            spatial: ThreeMomentum {
                px: &self.spatial.px * x,
                py: &self.spatial.py * x,
                pz: &self.spatial.pz * x,
            },
        }
    }

    pub(crate) fn set_energy_on_shell(&self, mass: Option<F<T>>) -> Self {
        let energy = self.spatial.on_shell_energy(mass);
        FourMomentum {
            temporal: energy,
            spatial: self.spatial.clone(),
        }
    }

    pub(crate) fn rapidity(&self) -> F<T> {
        let pt = self.spatial.pt();
        self.rapidity_with_pt2(&pt.square())
    }

    pub(crate) fn rapidity_with_pt2(&self, pt2: &F<T>) -> F<T> {
        const MAX_RAPIDITY_SENTINEL: usize = 100_000;

        let zero = self.temporal.value.zero();
        let pz = self.spatial.pz.clone();
        let abs_pz = pz.norm();
        let energy = self.temporal.value.clone();

        if *pt2 == zero.clone() && energy == abs_pz {
            let max_rapidity_here = abs_pz.from_usize(MAX_RAPIDITY_SENTINEL) + abs_pz;
            if pz >= zero {
                max_rapidity_here
            } else {
                -max_rapidity_here
            }
        } else {
            let mass2 = energy.square()
                - (self.spatial.px.square() + self.spatial.py.square() + self.spatial.pz.square());
            let effective_mass2 = if mass2 < zero { zero.clone() } else { mass2 };
            let e_plus_abs_pz = energy + abs_pz;
            let half = e_plus_abs_pz.one() / e_plus_abs_pz.from_usize(2);
            let mut rapidity =
                ((pt2.clone() + effective_mass2) / (e_plus_abs_pz.clone() * e_plus_abs_pz)).ln()
                    * half;
            if pz > zero {
                rapidity = -rapidity;
            }
            rapidity
        }
    }
}

#[derive(
    Debug, PartialEq, Eq, Clone, Copy, Serialize, Deserialize, Encode, Decode, Hash, JsonSchema,
)]
// #[trait_decode(trait= GammaLoopContext)]
#[serde(untagged)]
pub enum ExternalMomenta<T> {
    Dependent(Dep),
    Independent([T; 4]),
}

impl<T: FloatLike> Rotatable for ExternalMomenta<F<T>> {
    fn rotate(&self, rotation: &Rotation) -> Self {
        match self {
            ExternalMomenta::Dependent(_) => ExternalMomenta::Dependent(Dep::Dep),
            ExternalMomenta::Independent(data) => {
                let fm = FourMomentum::from(data.clone());
                fm.rotate(rotation).into()
            }
        }
    }
}

#[derive(
    Debug, PartialEq, Eq, Clone, Copy, Serialize, Deserialize, Encode, Decode, Hash, JsonSchema,
)]
pub enum Dep {
    #[serde(rename = "dependent")]
    Dep,
}

#[derive(Error, Debug)]
pub enum ExternalMomentaError {
    #[error("Dependent momenta cannot be converted to FourMomentum")]
    Dependent,
}

impl<T> TryFrom<ExternalMomenta<T>> for FourMomentum<T> {
    type Error = ExternalMomentaError;
    fn try_from(value: ExternalMomenta<T>) -> Result<Self, Self::Error> {
        match value {
            ExternalMomenta::Dependent(_) => Err(ExternalMomentaError::Dependent),
            ExternalMomenta::Independent(data) => Ok(FourMomentum::from(data)),
        }
    }
}

impl<T> From<FourMomentum<T>> for ExternalMomenta<T> {
    fn from(value: FourMomentum<T>) -> Self {
        ExternalMomenta::Independent(value.into())
    }
}

impl<T> From<[T; 4]> for ExternalMomenta<T> {
    fn from(data: [T; 4]) -> Self {
        ExternalMomenta::Independent(data)
    }
}

impl<T: RefZero, U: RefZero> RefZero for FourMomentum<T, U> {
    fn ref_zero(&self) -> Self {
        FourMomentum {
            temporal: self.temporal.ref_zero(),
            spatial: self.spatial.ref_zero(),
        }
    }
}

impl<T: RefZero, U: RefZero> RefZero<FourMomentum<T, U>> for &FourMomentum<T, U> {
    fn ref_zero(&self) -> FourMomentum<T, U> {
        FourMomentum {
            temporal: self.temporal.ref_zero(),
            spatial: self.spatial.ref_zero(),
        }
    }
}

#[derive(
    Debug, PartialEq, Clone, Copy, Serialize, Deserialize, Encode, Decode, Eq, PartialOrd, Ord,
)]
pub struct PolDef {
    pub pol_type: PolType,
    pub eid: EdgeIndex,
}

#[derive(
    Debug, PartialEq, Clone, Copy, Serialize, Deserialize, Encode, Decode, Eq, PartialOrd, Ord,
)]
pub enum PolType {
    U,
    V,
    UBar,
    VBar,
    Scalar,
    Epsilon,
    EpsilonBar,
}

impl PolType {
    pub(crate) fn bar(self) -> Self {
        match self {
            Self::Epsilon => Self::EpsilonBar,
            Self::Scalar => Self::Scalar,
            Self::U => Self::UBar,
            Self::UBar => Self::U,
            Self::EpsilonBar => Self::Epsilon,
            Self::VBar => Self::V,
            Self::V => Self::VBar,
        }
    }
}

impl Display for PolType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::U => write!(f, "u"),
            Self::V => write!(f, "v"),
            Self::VBar => write!(f, "vbar"),
            Self::UBar => write!(f, "ubar"),
            Self::Scalar => write!(f, ""),
            Self::Epsilon => write!(f, "ϵ"),
            Self::EpsilonBar => write!(f, "ϵbar"),
        }
    }
}

#[derive(Debug, PartialEq, Clone, bincode_trait_derive::Encode, bincode_trait_derive::Decode)]
#[trait_decode(trait = symbolica::state::HasStateMap)]
pub struct Polarization<T> {
    pub tensor: DenseTensor<T, IndexLess<LibraryRep>>,
    pol_type: PolType,
}

impl<T: FloatLike> ApproxEq<Polarization<Complex<F<T>>>, F<T>> for Polarization<Complex<F<T>>> {
    fn approx_eq(&self, other: &Polarization<Complex<F<T>>>, tolerance: &F<T>) -> bool {
        self.pol_type == other.pol_type
            && self
                .tensor
                .flat_iter()
                .zip(other.tensor.flat_iter())
                .all(|((_, a), (_, b))| a.approx_eq(b, tolerance))
    }

    fn approx_eq_res(
        &self,
        other: &Polarization<Complex<F<T>>>,
        tolerance: &F<T>,
    ) -> color_eyre::Result<()> {
        if self.pol_type != other.pol_type {
            Err(eyre::eyre!(
                "Polarization types do not match. self: {}, other: {}",
                self,
                other
            ))
        } else {
            self.tensor
                .flat_iter()
                .zip(other.tensor.flat_iter())
                .try_for_each(|((i, a), (_, b))| {
                    a.approx_eq_res(b, tolerance).wrap_err(format!(
                        "Polarization tensor element {} does not match. self: {}, other: {}",
                        i, self, other
                    ))
                })
        }
    }
}

impl<T: for<'a> std::ops::AddAssign<&'a T>> AddAssign<Polarization<T>> for Polarization<T> {
    fn add_assign(&mut self, rhs: Polarization<T>) {
        self.tensor += rhs.tensor;
    }
}

impl<T: for<'a> SubAssign<&'a T>> SubAssign<Polarization<T>> for Polarization<T> {
    fn sub_assign(&mut self, rhs: Polarization<T>) {
        self.tensor -= rhs.tensor;
    }
}

impl<T: FloatLike> Display for Polarization<F<T>> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "Pol {}: {}", self.pol_type, self.tensor)
    }
}

impl<T: FloatLike> Display for Polarization<Complex<F<T>>> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "Pol {}: {}", self.pol_type, self.tensor)
    }
}

impl<T: FloatLike> LowerExp for Polarization<Complex<F<T>>> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "Pol {}: {:+e}", self.pol_type, self.tensor)
    }
}

impl<T> Polarization<T> {
    pub(crate) fn map<U>(&self, f: impl Fn(&T) -> U) -> Polarization<U> {
        Polarization {
            tensor: self.tensor.map_data_ref(f),
            pol_type: self.pol_type,
        }
    }
}

impl<T: Clone> Polarization<T> {
    pub(crate) fn is_scalar(&self) -> bool {
        self.pol_type == PolType::Scalar
    }
    pub(crate) fn scalar(value: T) -> Self {
        let structure = IndexLess::new(vec![]);
        Polarization {
            tensor: DenseTensor {
                data: vec![value],
                structure,
            },
            pol_type: PolType::Scalar,
        }
    }

    pub(crate) fn lorentz(value: [T; 4]) -> Self {
        let structure = IndexLess::new(vec![Minkowski {}.new_rep(4).cast()]);
        Polarization {
            tensor: DenseTensor {
                data: value.to_vec(),
                structure,
            },
            pol_type: PolType::Epsilon,
        }
    }

    pub(crate) fn bispinor_u(value: [T; 4]) -> Self {
        let structure = IndexLess::new(vec![GR.bis.new_rep(4)]);

        Polarization {
            tensor: DenseTensor {
                data: value.to_vec(),
                structure,
            },
            pol_type: PolType::U,
        }
    }

    pub(crate) fn bispinor_v(value: [T; 4]) -> Self {
        let structure = IndexLess::new(vec![GR.bis.new_rep(4)]);

        Polarization {
            tensor: DenseTensor {
                data: value.to_vec(),
                structure,
            },
            pol_type: PolType::V,
        }
    }

    pub(crate) fn shadow(&self) -> DenseTensor<Atom, IndexLess<LibraryRep>> {
        self.tensor
            .structure
            .clone()
            .to_dense_labeled(|_, i| FlatCoefficent::<NoArgs> {
                index: i,
                name: Some(symbol!(self.pol_type.to_string())),
                args: None,
            })
            .unwrap()
    }
}

impl<T> Index<usize> for Polarization<T> {
    type Output = T;
    fn index(&self, index: usize) -> &Self::Output {
        &self.tensor[index.into()]
    }
}

impl<T> RefZero for Polarization<T>
where
    T: RefZero + Clone,
{
    fn ref_zero(&self) -> Self {
        Polarization {
            tensor: <DenseTensor<_, _> as RefZero>::ref_zero(&self.tensor),
            pol_type: self.pol_type,
        }
    }
}

impl<'a, T, U> Mul<&'a U> for Polarization<T>
where
    T: Clone + MulAssign<&'a U>,
{
    type Output = Polarization<T>;
    fn mul(mut self, rhs: &'a U) -> Self::Output {
        for i in self.tensor.data.iter_mut() {
            *i *= rhs;
        }
        self
    }
}

impl<'a, T> MulAssign<&'a T> for Polarization<T>
where
    T: Clone + MulAssign<&'a T>,
{
    fn mul_assign(&mut self, rhs: &'a T) {
        for i in self.tensor.data.iter_mut() {
            *i *= rhs;
        }
    }
}

impl<T, U, Out> FallibleAdd<Polarization<U>> for Polarization<T>
where
    T: FallibleAdd<U, Output = Out>,
{
    type Output = Polarization<Out>;
    fn add_fallible(&self, rhs: &Polarization<U>) -> Option<Self::Output> {
        Some(Polarization {
            tensor: self.tensor.add_fallible(&rhs.tensor)?,
            pol_type: self.pol_type,
        })
    }
}

impl<T> Polarization<T> {
    pub(crate) fn cast<U>(&self) -> Polarization<U>
    where
        T: Clone,
        U: Clone + From<T>,
    {
        Polarization {
            tensor: self.tensor.cast(),
            pol_type: self.pol_type,
        }
    }
}

impl<T> Polarization<Complex<T>> {
    pub(crate) fn complex_cast<U>(&self) -> Polarization<Complex<U>>
    where
        T: Clone,
        U: Clone + From<T>,
    {
        Polarization {
            tensor: self
                .tensor
                .map_data_ref(|d| d.map_ref(|r| r.clone().into())),
            pol_type: self.pol_type,
        }
    }
}

impl<T: FloatLike> Polarization<Complex<F<T>>> {
    pub(crate) fn higher(&self) -> Polarization<Complex<F<T::Higher>>>
    where
        T::Higher: FloatLike,
    {
        Polarization {
            tensor: self.tensor.map_data_ref(|t| t.map_ref(|r| r.higher())),
            pol_type: self.pol_type,
        }
    }

    pub(crate) fn lower(&self) -> Polarization<Complex<F<T::Lower>>>
    where
        T::Lower: FloatLike,
    {
        Polarization {
            tensor: self.tensor.map_data_ref(|t| t.map_ref(|r| r.lower())),
            pol_type: self.pol_type,
        }
    }
}

impl<T: FloatLike, U: FloatLike> From<FourMomentum<T, U>> for FourMomentum<F<T>, F<U>> {
    fn from(value: FourMomentum<T, U>) -> Self {
        FourMomentum {
            temporal: value.temporal.into(),
            spatial: value.spatial.into(),
        }
    }
}

impl<T: FloatLike, U: FloatLike> FourMomentum<F<T>, F<U>> {
    pub(crate) fn higher(&self) -> FourMomentum<F<T::Higher>, F<U::Higher>>
    where
        T::Higher: FloatLike,
        U::Higher: FloatLike,
    {
        FourMomentum {
            temporal: self.temporal.higher(),
            spatial: self.spatial.higher(),
        }
    }

    pub(crate) fn lower(&self) -> FourMomentum<F<T::Lower>, F<U::Lower>>
    where
        T::Lower: FloatLike,
        U::Lower: FloatLike,
    {
        FourMomentum {
            temporal: self.temporal.lower(),
            spatial: self.spatial.lower(),
        }
    }
}

impl<T, U> FourMomentum<T, U> {
    pub(crate) fn new(energy: Energy<U>, three_momentum: ThreeMomentum<T>) -> Self {
        FourMomentum {
            temporal: energy,
            spatial: three_momentum,
        }
    }
}

impl<T: Real> RefDefault for FourMomentum<T, T> {
    fn default(&self) -> Self {
        let zero = self.temporal.value.zero();
        FourMomentum {
            temporal: Energy::new(zero.clone()),
            spatial: ThreeMomentum::new(zero.clone(), zero.clone(), zero.clone()),
        }
    }
}

impl<T> Mul<T> for &FourMomentum<T, T>
where
    T: Mul<T, Output = T> + Clone,
{
    type Output = FourMomentum<T, T>;
    fn mul(self, rhs: T) -> Self::Output {
        FourMomentum {
            temporal: Energy {
                value: self.temporal.value.clone() * rhs.clone(),
            },
            spatial: self.spatial.clone() * rhs,
        }
    }
}

impl<T> Mul<&T> for &FourMomentum<T, T>
where
    T: for<'b> Mul<&'b T, Output = T> + Clone,
{
    type Output = FourMomentum<T, T>;
    fn mul(self, rhs: &T) -> Self::Output {
        FourMomentum {
            temporal: Energy {
                value: self.temporal.value.clone() * rhs,
            },
            spatial: self.spatial.clone() * rhs,
        }
    }
}

impl<T> FourMomentum<T, T> {
    pub(crate) fn zero(&self) -> Self
    where
        T: Real,
    {
        let zero = self.temporal.value.zero();
        FourMomentum {
            temporal: Energy::new(zero.clone()),
            spatial: ThreeMomentum::new(zero.clone(), zero.clone(), zero.clone()),
        }
    }

    pub(crate) fn square(&self) -> T
    where
        T: for<'a> Mul<&'a T, Output = T> + Add<T, Output = T> + Clone + Sub<T, Output = T>,
    {
        let temporal = self.temporal.value.clone();
        let spatial = self.spatial.norm_squared();
        temporal * &self.temporal.value - spatial
    }

    pub(crate) fn norm(&self) -> T
    where
        T: Real,
    {
        self.square().sqrt()
    }

    pub(crate) fn from_args(energy: T, px: T, py: T, pz: T) -> Self {
        let energy = Energy::new(energy);
        let three_momentum = ThreeMomentum { px, py, pz };
        FourMomentum {
            temporal: energy,
            spatial: three_momentum,
        }
    }
    pub(crate) fn new_on_shell(three_momentum: ThreeMomentum<T>, mass: Option<T>) -> Self
    where
        T: Mul<T> + Add<T> + std::ops::Add<Output = T> + Real,
    {
        let energy = three_momentum.on_shell_energy(mass);
        // println!("{}", energy);
        FourMomentum {
            temporal: energy,
            spatial: three_momentum,
        }
    }

    pub(crate) fn into_dense(self, index: AbstractIndex) -> DenseTensor<T, OrderedStructure>
    where
        T: Clone,
    {
        let structure =
            Canonicalized::from_iter([LibraryRep::new_slot(Minkowski {}.into(), 4, index)])
                .into_canonical();
        DenseTensor::from_storage_data(
            vec![
                self.temporal.value,
                self.spatial.px,
                self.spatial.py,
                self.spatial.pz,
            ],
            structure,
        )
        .unwrap()
    }

    pub(crate) fn into_dense_named(
        self,
        index: AbstractIndex,
        name: Symbol,
        num: usize,
    ) -> DenseTensor<T, NamedStructure<Symbol, usize>>
    where
        T: Clone,
    {
        let structure = Canonicalized::<OrderedStructure>::from_iter([LibraryRep::new_slot(
            Minkowski {}.into(),
            4,
            index,
        )])
        .into_canonical()
        .to_named(name, Some(num));
        DenseTensor::from_storage_data(
            vec![
                self.temporal.value,
                self.spatial.px,
                self.spatial.py,
                self.spatial.pz,
            ],
            structure,
        )
        .unwrap()
    }

    pub(crate) fn cast<U>(&self) -> FourMomentum<U, U>
    where
        T: Clone + Into<U>,
    {
        FourMomentum {
            temporal: Energy::new((self.temporal.value.clone()).into()),
            spatial: ThreeMomentum::cast(&self.spatial),
        }
    }

    pub(crate) fn boost(&self, boost_vector: &ThreeMomentum<T>) -> FourMomentum<T>
    where
        T: Real + SingleFloat + PartialOrd,
    {
        let b2 = boost_vector.norm_squared();
        let one = b2.one();
        let zero = one.zero();
        let gamma = (one.clone() - &b2).sqrt().inv();

        let bp = self.spatial.clone() * boost_vector;
        let gamma2 = if b2 > zero {
            (gamma.clone() - &one) / b2
        } else {
            zero
        };
        let factor = gamma2 * &bp + gamma.clone() * &self.temporal.value;

        FourMomentum::from_args(
            (bp + &self.temporal.value) * &gamma,
            self.spatial.px.mul_add(&factor, &boost_vector.px),
            self.spatial.py.mul_add(&factor, &boost_vector.py),
            self.spatial.pz.mul_add(&factor, &boost_vector.pz),
        )
    }
}

impl<T: FloatLike> FourMomentum<F<T>, F<T>> {
    pub(crate) fn dot(&self, p2: &FourMomentum<F<T>>) -> F<T> {
        &self.temporal.value * &p2.temporal.value
            - (&self.spatial.px * &p2.spatial.px
                + &self.spatial.py * &p2.spatial.py
                + &self.spatial.pz * &p2.spatial.pz)
    }
    /// Compute the phi-angle separation with p2.
    pub(crate) fn getdelphi(&self, p2: &FourMomentum<F<T>>) -> F<T>
    where
        T: Real,
    {
        self.spatial.getdelphi(&p2.spatial)
    }

    /// Compute the deltaR separation with momentum p2.
    #[inline]
    pub(crate) fn delta_r(&self, p2: &FourMomentum<F<T>>) -> F<T>
    where
        T: Real,
    {
        self.spatial.delta_r(&p2.spatial)
    }

    pub(crate) fn pt(&self) -> F<T>
    where
        T: Real,
    {
        self.spatial.pt()
    }

    pub(crate) fn to_f64(&self) -> FourMomentum<F<f64>, F<f64>> {
        FourMomentum {
            temporal: Energy {
                value: F(self.temporal.value.to_f64()),
            },
            spatial: self.spatial.to_f64().cast(),
        }
    }

    pub(crate) fn pol_one(&self) -> [F<T>; 4]
    where
        T: FloatLike,
    {
        // definition from helas_ref A.2

        // debug!("pol_one in: {}", self);

        let pt = self.pt();
        let p = self.spatial.norm();

        let (e1, e2, e3) = if pt.is_zero() {
            (pt.one(), pt.zero(), pt.zero())
        } else {
            (
                &self.spatial.px * &self.spatial.pz / (&pt * &p),
                &self.spatial.py * &self.spatial.pz / (&pt * &p),
                -(&pt / &p),
            )
        };

        // debug!(
        //     " (pt.zero(), e1, e2, e3) {} {} {} {}",
        //     pt.zero(),
        //     e1,
        //     e2,
        //     e3
        // );
        [pt.zero(), e1, e2, e3]

        // debug!("pol :{pol}");
    }

    pub(crate) fn pol_two(&self) -> [F<T>; 4]
    where
        T: FloatLike,
    {
        // definition from helas_ref A.2
        let pt = self.pt();
        let (e1, e2, e3) = if pt.is_zero() {
            if self.spatial.pz.positive() {
                (pt.zero(), pt.one(), pt.zero())
            } else {
                (pt.zero(), -pt.one(), pt.zero())
            }
        } else {
            (-(&self.spatial.py / &pt), &self.spatial.px / &pt, pt.zero())
        };
        [pt.zero(), e1, e2, e3]
    }

    pub(crate) fn pol_three(&self) -> Polarization<F<T>>
    where
        T: FloatLike,
    {
        // definition from helas_ref A.2
        let m = self.norm();
        let p = self.spatial.norm();
        let emp = &self.temporal.value / (&m * &p);
        let e0 = p / &m;
        let e1 = &self.spatial.px * &emp;
        let e2 = &self.spatial.py * &emp;
        let e3 = &self.spatial.pz * &emp;

        Polarization::lorentz([e0, e1, e2, e3])
    }

    pub(crate) fn eps_pol<Hel: Into<Helicity>>(&self, lambda: Hel) -> Polarization<Complex<F<T>>> {
        match lambda.into() {
            Helicity::Signed(lambda) => {
                if lambda.is_zero() {
                    self.pol_three().cast()
                } else {
                    let one = self.temporal.value.one();
                    let sqrt_2_inv: F<T> = (&one + &one).sqrt().inv();

                    let [eone0, eone1, eone2, eone3] = self.pol_one();

                    let [etwo0, etwo1, etwo2, etwo3] = self.pol_two();

                    let components = [
                        Complex {
                            re: -lambda * eone0 * &sqrt_2_inv,
                            im: -etwo0 * &sqrt_2_inv, //using opposite convention with respect to helas to align with madgraph
                        },
                        Complex {
                            re: -lambda * eone1 * &sqrt_2_inv,
                            im: -etwo1 * &sqrt_2_inv,
                        },
                        Complex {
                            re: -lambda * eone2 * &sqrt_2_inv,
                            im: -etwo2 * &sqrt_2_inv,
                        },
                        Complex {
                            re: -lambda * eone3 * &sqrt_2_inv,
                            im: -etwo3 * &sqrt_2_inv,
                        },
                    ];

                    Polarization::lorentz(components)
                }
            }
            _ => {
                let zero = self.temporal.value.zero();
                let cmpz = Complex::new_re(zero);
                Polarization::lorentz([cmpz.clone(), cmpz.clone(), cmpz.clone(), cmpz.clone()])
            }
        }
    }

    pub(crate) fn omega(&self, lambda: Sign) -> Complex<F<T>> {
        match lambda {
            Sign::Positive => (&self.temporal.value + self.spatial.norm()).complex_sqrt(),
            Sign::Negative => (&self.temporal.value - self.spatial.norm()).complex_sqrt(),
        }
    }

    pub(crate) fn u(&self, lambda: Sign) -> Polarization<Complex<F<T>>> {
        let xi = self.xi(lambda);
        Polarization::bispinor_u([
            self.omega(-lambda) * &xi[0],
            self.omega(-lambda) * &xi[1],
            self.omega(lambda) * &xi[0],
            self.omega(lambda) * &xi[1],
        ])
    }

    pub(crate) fn v(&self, lambda: Sign) -> Polarization<Complex<F<T>>> {
        let xi = self.xi(-lambda);
        Polarization::bispinor_v([
            (-lambda) * self.omega(lambda) * &xi[0],
            (-lambda) * self.omega(lambda) * &xi[1],
            lambda * self.omega(-lambda) * &xi[0],
            lambda * self.omega(-lambda) * &xi[1],
        ])
    }

    pub(crate) fn xi(&self, lambda: Sign) -> [Complex<F<T>>; 2] {
        if (self.spatial.pz < self.spatial.pz.zero()
            && self.spatial.py.is_zero()
            && self.spatial.px.is_zero())
            || (self.spatial.pz == -self.spatial.norm())
        {
            let zero: Complex<F<T>> = self.temporal.value.zero().into();
            let one = zero.one();
            // We are defining using madgraph conventions not helas, taking py =0 and the limit px =0 from below
            match lambda {
                Sign::Positive => [zero, -one],
                Sign::Negative => [one, zero],
            }
        } else {
            let prefactor: F<T> = ((F::from_f64(2.)
                * self.spatial.norm()
                * (self.spatial.norm() + &self.spatial.pz))
                .sqrt())
            .inv();
            let mut xi: [Complex<F<T>>; 2] = [
                Complex::new_re(&prefactor * (self.spatial.norm() + &self.spatial.pz)),
                Complex::new(self.spatial.px.clone(), self.spatial.py.clone()) * &prefactor,
            ]; //plus

            if matches!(lambda, Sign::Negative) {
                xi.swap(0, 1);
                xi[0].re = -xi[0].re.clone();
            }
            xi
        }
    }
}

impl<T: FloatLike> Polarization<Complex<F<T>>> {
    pub(crate) fn bar(&self) -> Self {
        let mut tensor = self.tensor.map_data_ref(Complex::conj);

        if matches!(
            self.pol_type,
            PolType::U | PolType::V | PolType::UBar | PolType::VBar
        ) {
            tensor.data.swap(0, 2);
            tensor.data.swap(1, 3);
        }
        Polarization {
            tensor,
            pol_type: self.pol_type.bar(),
        }
    }
}

impl<T> IntoIterator for Polarization<T> {
    type IntoIter = std::vec::IntoIter<Self::Item>;
    type Item = T;

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

impl<'a, T> IntoIterator for &'a Polarization<T> {
    type IntoIter = std::slice::Iter<'a, T>;
    type Item = &'a T;

    fn into_iter(self) -> Self::IntoIter {
        self.tensor.data.iter()
    }
}

#[derive(
    Debug, PartialEq, Eq, Clone, Copy, PartialOrd, Ord, Hash, bincode::Encode, bincode::Decode,
)]
pub enum Helicity {
    Signed(SignOrZero),
    Summed,
    SummedAveraged,
}

impl Display for Helicity {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Helicity::Signed(s) => write!(f, "{}", s),
            Helicity::Summed => write!(f, "summed"),
            Helicity::SummedAveraged => write!(f, "summed_averaged"),
        }
    }
}

impl Serialize for Helicity {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        match self {
            Helicity::Signed(SignOrZero::Plus) => serializer.serialize_i8(1),
            Helicity::Signed(SignOrZero::Minus) => serializer.serialize_i8(-1),
            Helicity::Signed(SignOrZero::Zero) => serializer.serialize_i8(0),
            Helicity::Summed => serializer.serialize_str("summed"),
            Helicity::SummedAveraged => serializer.serialize_str("summed_averaged"),
        }
    }
}

#[derive(Deserialize)]
#[serde(untagged)]
enum HelicitySerde {
    Numeric(i8),
    Text(String),
}

impl<'de> Deserialize<'de> for Helicity {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        let value = HelicitySerde::deserialize(deserializer)?;
        match value {
            HelicitySerde::Numeric(value) => {
                Helicity::try_from(value).map_err(serde::de::Error::custom)
            }
            HelicitySerde::Text(value) => {
                Helicity::from_str(&value).map_err(serde::de::Error::custom)
            }
        }
    }
}

impl JsonSchema for Helicity {
    fn schema_name() -> std::borrow::Cow<'static, str> {
        "Helicity".into()
    }

    fn json_schema(_generator: &mut schemars::SchemaGenerator) -> schemars::Schema {
        schemars::json_schema!({
            "description": "Helicity input. Use -1, 0, or 1 for fixed helicities; or use a string for summed modes.",
            "oneOf": [
                {
                    "type": "integer",
                    "enum": [-1, 0, 1]
                },
                {
                    "type": "string",
                    "enum": [
                        "+",
                        "-",
                        "0",
                        "plus",
                        "minus",
                        "zero",
                        "summed",
                        "summed_averaged",
                        "summed-averaged"
                    ]
                }
            ]
        })
    }
}

impl TryFrom<i8> for Helicity {
    type Error = SignError;
    fn try_from(value: i8) -> Result<Self, Self::Error> {
        Ok(Helicity::Signed(SignOrZero::try_from(value)?))
    }
}

impl From<Sign> for Helicity {
    fn from(sign: Sign) -> Self {
        Self::Signed(sign.into())
    }
}

impl From<SignOrZero> for Helicity {
    fn from(sign: SignOrZero) -> Self {
        Self::Signed(sign)
    }
}

impl TryFrom<Helicity> for SignOrZero {
    type Error = SignError;

    fn try_from(value: Helicity) -> Result<Self, Self::Error> {
        match value {
            Helicity::Signed(sign) => Ok(sign),
            Helicity::Summed | Helicity::SummedAveraged => Err(SignError::InvalidValue),
        }
    }
}

impl TryFrom<Helicity> for Sign {
    type Error = SignError;

    fn try_from(value: Helicity) -> Result<Self, Self::Error> {
        match value {
            Helicity::Signed(sign) => sign.try_into(),
            Helicity::Summed | Helicity::SummedAveraged => Err(SignError::InvalidValue),
        }
    }
}

#[allow(non_upper_case_globals)]
impl Helicity {
    pub const ZERO: Self = Self::Signed(SignOrZero::Zero);
    pub const PLUS: Self = Self::Signed(SignOrZero::Plus);
    pub const MINUS: Self = Self::Signed(SignOrZero::Minus);
    pub const Zero: Self = Self::ZERO;
    pub const Plus: Self = Self::PLUS;
    pub const Minus: Self = Self::MINUS;
}

impl FromStr for Helicity {
    type Err = String;

    fn from_str(value: &str) -> Result<Self, Self::Err> {
        match value.trim().to_ascii_lowercase().as_str() {
            "+" | "plus" => Ok(Self::PLUS),
            "-" | "minus" => Ok(Self::MINUS),
            "0" | "zero" => Ok(Self::ZERO),
            "summed" => Ok(Self::Summed),
            "summed_averaged" | "summed-averaged" => Ok(Self::SummedAveraged),
            _ => Err(format!(
                "Invalid helicity '{value}'. Expected one of -1, 0, 1, plus, minus, zero, summed, summed_averaged"
            )),
        }
    }
}

#[derive(
    Debug, PartialEq, Eq, Clone, Serialize, Deserialize, Encode, Decode, PartialOrd, Ord, Hash,
)]
pub struct Signature(pub Vec<SignOrZero>);

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

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

impl FromIterator<i8> for Signature {
    fn from_iter<I: IntoIterator<Item = i8>>(iter: I) -> Self {
        Signature(
            iter.into_iter()
                .map(|x| match x {
                    0 => SignOrZero::Zero,
                    1 => SignOrZero::Plus,
                    -1 => SignOrZero::Minus,
                    _ => panic!("Invalid value for Signature"),
                })
                .collect(),
        )
    }
}

impl FromIterator<isize> for Signature {
    fn from_iter<I: IntoIterator<Item = isize>>(iter: I) -> Self {
        Signature(
            iter.into_iter()
                .map(|x| match x {
                    0 => SignOrZero::Zero,
                    1 => SignOrZero::Plus,
                    -1 => SignOrZero::Minus,
                    _ => panic!("Invalid value for Signature"),
                })
                .collect(),
        )
    }
}

impl From<Vec<i8>> for Signature {
    fn from(value: Vec<i8>) -> Self {
        Signature::from_iter(value)
    }
}

impl IntoIterator for Signature {
    type Item = SignOrZero;
    type IntoIter = std::vec::IntoIter<Self::Item>;

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

impl<'a> IntoIterator for &'a Signature {
    type Item = SignOrZero;
    type IntoIter = std::iter::Copied<std::slice::Iter<'a, Self::Item>>;

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

impl Index<usize> for Signature {
    type Output = SignOrZero;
    fn index(&self, index: usize) -> &Self::Output {
        &self.0[index]
    }
}

impl Signature {
    pub(crate) fn validate_basis<T>(&self, basis: &[T]) -> bool {
        self.len() == basis.len()
    }

    pub(crate) fn sum(&mut self, other: &Self) {
        for (i, sign) in other.iter().enumerate() {
            match (self[i], sign) {
                (SignOrZero::Zero, SignOrZero::Zero) => self.0[i] = SignOrZero::Zero,
                (SignOrZero::Zero, SignOrZero::Plus) => self.0[i] = SignOrZero::Plus,
                (SignOrZero::Zero, SignOrZero::Minus) => self.0[i] = SignOrZero::Minus,
                (SignOrZero::Plus, SignOrZero::Zero) => self.0[i] = SignOrZero::Plus,
                (SignOrZero::Plus, SignOrZero::Plus) => panic!("cannot add two positive signs"),
                (SignOrZero::Plus, SignOrZero::Minus) => self.0[i] = SignOrZero::Zero,
                (SignOrZero::Minus, SignOrZero::Zero) => self.0[i] = SignOrZero::Minus,
                (SignOrZero::Minus, SignOrZero::Plus) => self.0[i] = SignOrZero::Zero,
                (SignOrZero::Minus, SignOrZero::Minus) => panic!("cannot add two negative signs"),
            }
        }
    }

    pub(crate) fn panic_validate_basis<T>(&self, basis: &[T]) {
        if !self.validate_basis(basis) {
            panic!(
                "Invalid basis for Signature, expected length {}, got length {}",
                self.len(),
                basis.len()
            );
        }
    }

    pub(crate) fn to_momtrop_format(&self) -> Vec<isize> {
        self.0
            .iter()
            .map(|x| match x {
                SignOrZero::Zero => 0,
                SignOrZero::Plus => 1,
                SignOrZero::Minus => -1,
            })
            .collect()
    }
    pub(crate) fn len(&self) -> usize {
        self.0.len()
    }

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

    pub(crate) fn is_empty(&self) -> bool {
        self.0.is_empty()
    }
    pub(crate) fn apply<T>(&self, basis: &[T]) -> T
    where
        T: RefZero + Clone + Neg<Output = T> + AddAssign<T>,
    {
        // self.panic_validate_basis(basis);
        let mut result = basis[0].ref_zero();
        for (&sign, t) in self.0.iter().zip(basis.iter().cloned()) {
            result += sign * t;
        }
        result
    }

    pub(crate) fn apply_iter<I, T>(&self, basis: I) -> Option<T>
    where
        I: IntoIterator,
        I::Item: RefZero<T>,
        T: Clone + SubAssign<I::Item> + AddAssign<I::Item>,
    {
        let mut basis_iter = basis.into_iter();
        let mut signature_iter = self.into_iter();

        while let (Some(sign), Some(item)) = (signature_iter.next(), basis_iter.next()) {
            if sign.is_sign() {
                // Initialize the result based on the first non-zero sign
                let mut result = item.ref_zero();
                match sign {
                    SignOrZero::Zero => {
                        panic!("unreachable");
                        // return None;
                    }
                    SignOrZero::Plus => {
                        result += item;
                    }
                    SignOrZero::Minus => {
                        result -= item;
                    }
                }

                // Continue processing the rest of the iterator
                while let (Some(sign), Some(item)) = (signature_iter.next(), basis_iter.next()) {
                    match sign {
                        SignOrZero::Zero => {}
                        SignOrZero::Plus => {
                            result += item;
                        }
                        SignOrZero::Minus => {
                            result -= item;
                        }
                    }
                }

                return Some(result);
            }
        }

        // Return None if no non-zero sign was found
        None
    }

    pub(crate) fn label_with(&self, label: &str) -> String {
        let mut result = String::new();
        let mut first = true;
        for (i, sign) in self.0.iter().enumerate() {
            if !first {
                result.push_str(&sign.to_string());
            } else {
                first = false;
            }
            if sign.is_sign() {
                result.push_str(&format!("{}_{}", label, i));
            }
        }
        result
    }
}

#[test]
fn test_signature() {
    let sig = Signature(vec![SignOrZero::Plus, SignOrZero::Minus]);
    let basis: Vec<i32> = vec![1, 2];
    assert_eq!(sig.apply(&basis), 1 - 2);
    assert_eq!(sig.apply_iter(basis.iter()), Some(-1));

    let basis: [FourMomentum<i32>; 4] = [
        FourMomentum::from_args(1, 1, 0, 0),
        FourMomentum::from_args(1, 0, 1, 0),
        FourMomentum::from_args(1, 0, 0, 1),
        FourMomentum::from_args(1, 1, 1, 1),
    ];

    let sig = Signature(vec![
        SignOrZero::Plus,
        SignOrZero::Minus,
        SignOrZero::Zero,
        SignOrZero::Plus,
    ]);

    assert_eq!(sig.apply(&basis), FourMomentum::from_args(1, 2, 0, 1));
    let sig = Signature(vec![
        SignOrZero::Zero,
        SignOrZero::Zero,
        SignOrZero::Zero,
        SignOrZero::Zero,
    ]);
    assert_eq!(sig.apply_iter(basis.iter()), None);
    let sig = Signature(vec![
        SignOrZero::Zero,
        SignOrZero::Zero,
        SignOrZero::Zero,
        SignOrZero::Minus,
    ]);
    assert_eq!(sig.apply_iter(basis.iter()), Some(-basis[3]));
}

// impl Serialize for Helicity {
//     fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
//     where
//         S: Serializer,
//     {
//         match *self {
//             Helicity::Sign(ref sign) => sign.serialize(serializer),
//             Helicity::ZERO => serializer.serialize_i32(0),
//         }
//     }
// }

// impl<'de> Deserialize<'de> for Helicity {
//     fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
//     where
//         D: Deserializer<'de>,
//     {
//         let value = i32::deserialize(deserializer)?;
//         match value {
//             1 => Ok(Helicity::Sign(Sign::Positive)),
//             -1 => Ok(Helicity::Sign(Sign::Negative)),
//             0 => Ok(Helicity::ZERO),
//             _ => Err(de::Error::custom(
//                 "Expected 1 for Positive, -1 for Negative, or 0 for Zero",
//             )),
//         }
//     }
// }

impl<T> IntoIterator for FourMomentum<T> {
    type Item = T;
    type IntoIter = std::array::IntoIter<T, 4>;

    fn into_iter(self) -> Self::IntoIter {
        let [e, px, py, pz] = [
            self.temporal.value,
            self.spatial.px,
            self.spatial.py,
            self.spatial.pz,
        ];
        [e, px, py, pz].into_iter()
    }
}

impl<'a, T> IntoIterator for &'a FourMomentum<T> {
    type Item = &'a T;
    type IntoIter = std::array::IntoIter<&'a T, 4>;

    fn into_iter(self) -> Self::IntoIter {
        let [e, px, py, pz] = [
            &self.temporal.value,
            &self.spatial.px,
            &self.spatial.py,
            &self.spatial.pz,
        ];
        [e, px, py, pz].into_iter()
    }
}

impl<T> From<[T; 4]> for FourMomentum<T, T> {
    fn from(data: [T; 4]) -> Self {
        let [t, px, py, pz] = data;
        FourMomentum {
            temporal: Energy::new(t),
            spatial: ThreeMomentum { px, py, pz },
        }
    }
}

impl<T> From<FourMomentum<T, T>> for [T; 4] {
    fn from(data: FourMomentum<T, T>) -> Self {
        [
            data.temporal.value,
            data.spatial.px,
            data.spatial.py,
            data.spatial.pz,
        ]
    }
}

impl<T> From<(T, T, T, T)> for FourMomentum<T, T> {
    fn from(data: (T, T, T, T)) -> Self {
        let (t, px, py, pz) = data;
        FourMomentum {
            temporal: Energy::new(t),
            spatial: ThreeMomentum { px, py, pz },
        }
    }
}

impl<T> From<FourMomentum<T, T>> for (T, T, T, T) {
    fn from(data: FourMomentum<T, T>) -> Self {
        (
            data.temporal.value,
            data.spatial.px,
            data.spatial.py,
            data.spatial.pz,
        )
    }
}

impl<T> FourMomentum<T, Atom> {
    pub(crate) fn into_dense_param(
        self,
        index: AbstractIndex,
    ) -> DenseTensor<MultivariatePolynomial<RationalField, T>, OrderedStructure>
    where
        T: Clone + Into<Coefficient> + Exponent,
    {
        let structure =
            Canonicalized::from_iter([LibraryRep::new_slot(Minkowski {}.into(), 4, index)])
                .into_canonical();
        let energy = self
            .temporal
            .value
            .to_polynomial(&RationalField::new(IntegerRing {}), None);

        let px: MultivariatePolynomial<RationalField, _> =
            Atom::num(self.spatial.px).to_polynomial(&RationalField::new(IntegerRing {}), None);
        let py =
            Atom::num(self.spatial.py).to_polynomial(&RationalField::new(IntegerRing {}), None);
        let pz =
            Atom::num(self.spatial.pz).to_polynomial(&RationalField::new(IntegerRing {}), None);

        DenseTensor::from_storage_data(vec![energy, px, py, pz], structure).unwrap()
    }
}

impl<T, U> Add<FourMomentum<T, U>> for FourMomentum<T, U>
where
    T: Add<T, Output = T>,
    U: Add<U, Output = U>,
{
    type Output = FourMomentum<T, U>;
    fn add(self, rhs: FourMomentum<T, U>) -> Self::Output {
        FourMomentum {
            temporal: self.temporal + rhs.temporal,
            spatial: self.spatial + rhs.spatial,
        }
    }
}

impl<T, U> AddAssign<FourMomentum<T, U>> for FourMomentum<T, U>
where
    T: AddAssign<T>,
    U: AddAssign<U>,
{
    fn add_assign(&mut self, rhs: FourMomentum<T, U>) {
        self.temporal += rhs.temporal;
        self.spatial += rhs.spatial;
    }
}

impl<'a, T, U> AddAssign<&'a FourMomentum<T, U>> for FourMomentum<T, U>
where
    T: AddAssign<&'a T>,
    U: AddAssign<&'a U>,
{
    fn add_assign(&mut self, rhs: &'a FourMomentum<T, U>) {
        self.temporal += &rhs.temporal;
        self.spatial += &rhs.spatial;
    }
}

impl<T, U> Sub<FourMomentum<T, U>> for FourMomentum<T, U>
where
    T: Sub<T, Output = T>,
    U: Sub<U, Output = U>,
{
    type Output = FourMomentum<T, U>;
    fn sub(self, rhs: FourMomentum<T, U>) -> Self::Output {
        FourMomentum {
            temporal: self.temporal - rhs.temporal,
            spatial: self.spatial - rhs.spatial,
        }
    }
}

impl<T, U> Neg for FourMomentum<T, U>
where
    T: Neg<Output = T>,
    U: Neg<Output = U>,
{
    type Output = FourMomentum<T, U>;
    fn neg(self) -> Self::Output {
        FourMomentum {
            temporal: -self.temporal,
            spatial: -self.spatial,
        }
    }
}

impl<T, U> Neg for &FourMomentum<T, U>
where
    T: Neg<Output = T> + Clone,
    U: Neg<Output = U> + Clone,
{
    type Output = FourMomentum<T, U>;
    fn neg(self) -> Self::Output {
        FourMomentum {
            temporal: -self.temporal.clone(),
            spatial: -self.spatial.clone(),
        }
    }
}

impl<T, U> SubAssign<FourMomentum<T, U>> for FourMomentum<T, U>
where
    T: SubAssign<T>,
    U: SubAssign<U>,
{
    fn sub_assign(&mut self, rhs: FourMomentum<T, U>) {
        self.temporal -= rhs.temporal;
        self.spatial -= rhs.spatial;
    }
}

impl<'a, T, U> SubAssign<&'a FourMomentum<T, U>> for FourMomentum<T, U>
where
    T: SubAssign<&'a T>,
    U: SubAssign<&'a U>,
{
    fn sub_assign(&mut self, rhs: &'a FourMomentum<T, U>) {
        self.temporal -= &rhs.temporal;
        self.spatial -= &rhs.spatial;
    }
}

impl<T: Display, U: Display> Display for FourMomentum<T, U> {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "{}, {}", self.temporal, self.spatial)
    }
}

impl<T: LowerExp, U: LowerExp> LowerExp for FourMomentum<T, U> {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "{:e}, {:e}", self.temporal, self.spatial)
    }
}

impl<T: FloatLike> From<Vector<F<T>, 3>> for ThreeMomentum<F<T>> {
    fn from(value: Vector<F<T>, 3>) -> Self {
        ThreeMomentum::new(value[0].clone(), value[1].clone(), value[2].clone())
    }
}

pub enum LorentzTransformation<T> {
    Boost(ThreeMomentum<T>),
    Rotation(Rotation),
    General {
        boost: ThreeMomentum<T>,
        rotation: Rotation,
    },
}

pub trait LorentzTransformable<T>: Rotatable {
    fn lorentz_transform(&self, transformation: &LorentzTransformation<T>) -> Self;
}

// pub struct IrriducibleLorentzRep {
//     m: usize,
//     n: usize,
// }

// impl IrriducibleLorentzRep {
//     pub(crate) fn scalar() -> Self {
//         Self { m: 0, n: 0 }
//     }

//     pub(crate) fn vector() -> Self {
//         Self { m: 1, n: 1 }
//     }

//     pub(crate) fn tensor() -> Self {
//         Self { m: 2, n: 2 }
//     }

//     pub(crate) fn left_weyl() -> Self {
//         Self { m: 1, n: 0 }
//     }

//     pub(crate) fn right_weyl() -> Self {
//         Self { m: 0, n: 1 }
//     }
// }

pub enum LorentzRep {
    // Irriducible(IrriducibleLorentzRep),
    // DirectSum(Vec<LorentzRep>),
    Scalar,
    Vector,
    Bispinor,
}

#[derive(Clone, Copy, Debug, Encode, Decode)]
pub enum RotationMethod {
    EulerAngles(f64, f64, f64),
    Pi2X,
    Pi2Y,
    Pi2Z,
    Identity,
}

impl Display for RotationMethod {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            RotationMethod::EulerAngles(alpha, beta, gamma) => {
                write!(f, "EulerAngles({}, {}, {})", alpha, beta, gamma)
            }
            RotationMethod::Pi2X => write!(f, "Pi/2 rotation around x-axis"),
            RotationMethod::Pi2Y => write!(f, "Pi/2 rotation around y-axis"),
            RotationMethod::Pi2Z => write!(f, "Pi/2 rotation around z-axis"),
            RotationMethod::Identity => write!(f, "Identity rotation"),
        }
    }
}

impl From<RotationMethod> for Rotation {
    fn from(method: RotationMethod) -> Self {
        Rotation::new(method)
    }
}

#[derive(Clone, Encode, Decode)]
#[trait_decode(trait = GammaLoopContext)]
pub struct Rotation {
    pub method: RotationMethod,
    pub lorentz_rotation: EvalTensor<ExpressionEvaluator<Complex<F<f64>>>, OrderedStructure>,
    pub bispinor_rotation: EvalTensor<ExpressionEvaluator<Complex<F<f64>>>, OrderedStructure>,
    // phantom: PhantomData<T>,
}

impl Rotation {
    pub(crate) fn is_identity(&self) -> bool {
        matches!(self.method, RotationMethod::Identity)
    }

    pub(crate) fn inverse_rotate_three<T: FloatLike>(
        &self,
        momentum: &ThreeMomentum<F<T>>,
    ) -> ThreeMomentum<F<T>> {
        match self.method {
            RotationMethod::Identity => momentum.clone(),
            RotationMethod::Pi2X => ThreeMomentum {
                px: momentum.px.clone(),
                py: momentum.pz.clone(),
                pz: -momentum.py.clone(),
            },
            RotationMethod::Pi2Y => ThreeMomentum {
                px: -momentum.pz.clone(),
                py: momentum.py.clone(),
                pz: momentum.px.clone(),
            },
            RotationMethod::Pi2Z => ThreeMomentum {
                px: momentum.py.clone(),
                py: -momentum.px.clone(),
                pz: momentum.pz.clone(),
            },
            RotationMethod::EulerAngles(alpha, beta, gamma) => {
                // Rotation settings are persisted as f64, so reconstructing them at the
                // active numeric precision is an explicit f64 boundary.
                let alpha = F::<T>::from_f64(alpha);
                let beta = F::<T>::from_f64(beta);
                let gamma = F::<T>::from_f64(gamma);
                let sin_alpha = alpha.sin();
                let cos_alpha = alpha.cos();
                let sin_beta = beta.sin();
                let cos_beta = beta.cos();
                let sin_gamma = gamma.sin();
                let cos_gamma = gamma.cos();

                let px = momentum.px.clone();
                let py = momentum.py.clone();
                let pz = momentum.pz.clone();

                ThreeMomentum {
                    px: cos_gamma.clone() * &cos_beta * &px + sin_gamma.clone() * &cos_beta * &py
                        - sin_beta.clone() * &pz,
                    py: (-(cos_alpha.clone()) * &sin_gamma
                        + sin_alpha.clone() * &sin_beta * &cos_gamma)
                        * &px
                        + (cos_alpha.clone() * &cos_gamma
                            + sin_alpha.clone() * &sin_beta * &sin_gamma)
                            * &py
                        + cos_beta.clone() * &sin_alpha * &pz,
                    pz: (sin_alpha.clone() * &sin_gamma
                        + cos_alpha.clone() * &sin_beta * &cos_gamma)
                        * &px
                        + (-sin_alpha.clone() * &cos_gamma
                            + cos_alpha.clone() * &sin_beta * &sin_gamma)
                            * &py
                        + cos_alpha * &cos_beta * &pz,
                }
            }
        }
    }

    pub(crate) fn inverse_rotate_four<T: FloatLike>(
        &self,
        momentum: &FourMomentum<F<T>>,
    ) -> FourMomentum<F<T>> {
        FourMomentum {
            temporal: momentum.temporal.clone(),
            spatial: self.inverse_rotate_three(&momentum.spatial),
        }
    }

    pub(crate) fn setting(&self) -> RotationSetting {
        match self.method {
            RotationMethod::EulerAngles(alpha, beta, gamma) => {
                RotationSetting::EulerAngles { alpha, beta, gamma }
            }
            RotationMethod::Pi2X => RotationSetting::Pi2X {},
            RotationMethod::Pi2Y => RotationSetting::Pi2Y {},
            RotationMethod::Pi2Z => RotationSetting::Pi2Z {},
            RotationMethod::Identity => RotationSetting::None {},
        }
    }
    pub(crate) fn new(method: RotationMethod) -> Self {
        let mu = Minkowski::slot(4, 1);

        let al = Minkowski::slot(4, 3);
        let mud = mu.dual();

        let shadow: NamedStructure<String, ()> = Canonicalized::<OrderedStructure>::from_iter([mu])
            .into_canonical()
            .to_named("eps".to_string(), None);
        let shadow_t: MixedTensor<_, OrderedStructure> =
            ParamOrConcrete::param(shadow.to_shell().expanded_shadow().unwrap().into())
                .cast_structure();

        let rotation: MixedTensor<f64, _> = method
            .lorentz_tensor(mud, al)
            .try_into_dense()
            .unwrap()
            .into();

        let fn_map = FunctionMap::new();

        let lorentz_eval: EvalTensor<ExpressionEvaluator<SymComplex<Rational>>, OrderedStructure> =
            shadow_t
                .contract(&rotation)
                .unwrap()
                .try_into_parametric()
                .unwrap()
                .to_evaluation_tree(
                    &fn_map,
                    &shadow_t.try_into_parametric().unwrap().tensor.data(),
                )
                .unwrap()
                .linearize(
                    &OptimizationSettings::new()
                        .cpe_iterations(Some(1))
                        .verbose(false),
                );

        let i = GR.bis.new_slot(4, 1);

        let j = GR.bis.new_slot(4, 3);

        let shadow: NamedStructure<String, ()> = Canonicalized::from_iter([i.cast::<LibraryRep>()])
            .into_canonical()
            .to_named("u".to_string(), None);
        let shadow_t: MixedTensor<_, OrderedStructure> =
            ParamOrConcrete::param(shadow.to_shell().expanded_shadow().unwrap().into())
                .cast_structure();

        let rotation: MixedTensor<f64, _> =
            ParamOrConcrete::Concrete(RealOrComplexTensor::Complex(method.bispinor_tensor(i, j)));

        let res = shadow_t
            .contract(&rotation)
            .unwrap()
            .try_into_parametric()
            .unwrap();

        let fn_map = FunctionMap::new();
        let mut params = shadow_t.try_into_parametric().unwrap().tensor.data();
        params.push(Atom::i());

        let spinor_eval: EvalTensor<ExpressionEvaluator<SymComplex<Rational>>, OrderedStructure> =
            res.to_evaluation_tree(&fn_map, &params).unwrap().linearize(
                &OptimizationSettings::new()
                    .cpe_iterations(Some(1))
                    .verbose(false),
            );

        Self {
            method,
            lorentz_rotation: lorentz_eval.map_coeff(&|f| {
                Complex::new(F::from_f64(f.re.to_f64()), F::from_f64(f.im.to_f64()))
            }),
            bispinor_rotation: spinor_eval.map_coeff(&|f| {
                Complex::new(F::from_f64(f.re.to_f64()), F::from_f64(f.im.to_f64()))
            }),
        }
    }
}

impl RotationMethod {
    pub(crate) fn generator(
        &self,
        i: AbstractIndex,
        j: AbstractIndex,
    ) -> DataTensor<f64, OrderedStructure> {
        let structure = Canonicalized::from_iter([
            LibraryRep::new_slot(Minkowski {}.into(), 4, i),
            LibraryRep::new_slot(Minkowski {}.into(), 4, j),
        ]);
        let layout = TensorDataLayout::from_canonicalized(&structure)
            .expect("rotation generators have concrete dimensions");
        let zero = 0.;
        match self {
            RotationMethod::Identity => {
                let omega = SparseTensor::empty(structure.into_canonical(), zero);
                omega.into()
            }
            RotationMethod::Pi2X => {
                let mut omega = SparseTensor::empty(structure.into_canonical(), zero);
                omega
                    .set(
                        &layout
                            .logical_expanded_to_storage_expanded(&[2, 3])
                            .unwrap(),
                        -zero.PIHALF(),
                    )
                    .unwrap();
                omega
                    .set(
                        &layout
                            .logical_expanded_to_storage_expanded(&[3, 2])
                            .unwrap(),
                        zero.PIHALF(),
                    )
                    .unwrap();
                omega.into()
            }
            RotationMethod::Pi2Y => {
                let mut omega = SparseTensor::empty(structure.into_canonical(), zero);
                omega
                    .set(
                        &layout
                            .logical_expanded_to_storage_expanded(&[1, 3])
                            .unwrap(),
                        zero.PIHALF(),
                    )
                    .unwrap();
                omega
                    .set(
                        &layout
                            .logical_expanded_to_storage_expanded(&[3, 1])
                            .unwrap(),
                        -zero.PIHALF(),
                    )
                    .unwrap();
                omega.into()
            }
            RotationMethod::Pi2Z => {
                let mut omega = SparseTensor::empty(structure.into_canonical(), zero);
                omega
                    .set(
                        &layout
                            .logical_expanded_to_storage_expanded(&[1, 2])
                            .unwrap(),
                        -zero.PIHALF(),
                    )
                    .unwrap();
                omega
                    .set(
                        &layout
                            .logical_expanded_to_storage_expanded(&[2, 1])
                            .unwrap(),
                        zero.PIHALF(),
                    )
                    .unwrap();
                omega.into()
            }
            RotationMethod::EulerAngles(alpha, beta, gamma) => DenseTensor::from_storage_data(
                layout
                    .reorder_to_storage(vec![
                        // row 0
                        zero, zero, zero, zero, // row 1
                        zero, zero, -gamma, *beta, // row 2
                        zero, *gamma, zero, -alpha, // row 3
                        zero, -beta, *alpha, zero,
                    ])
                    .unwrap(),
                structure.into_canonical(),
            )
            .unwrap()
            .into(),
        }
    }

    pub(crate) fn lorentz_tensor(
        &self,
        i: Slot<Minkowski>,
        j: Slot<Minkowski>,
    ) -> DataTensor<f64, OrderedStructure> {
        let structure = Canonicalized::from_iter([i.cast::<LibraryRep>(), j.cast()]);
        let layout = TensorDataLayout::from_canonicalized(&structure)
            .expect("Lorentz rotations have concrete dimensions");

        match self {
            RotationMethod::Identity => {
                let rot = DenseTensor::from_storage_data(
                    layout
                        .reorder_to_storage(vec![
                            1., 0., 0., 0., // row 1
                            0., -1., 0., 0., // row 2
                            0., 0., -1., 0., // row 3
                            0., 0., 0., -1.,
                        ])
                        .unwrap(),
                    structure.into_canonical(),
                )
                .unwrap();
                rot.into()
            }
            RotationMethod::Pi2X => {
                let rot = DenseTensor::from_storage_data(
                    layout
                        .reorder_to_storage(vec![
                            1., 0., 0., 0., // row 1
                            0., -1., 0., 0., // row 2
                            0., 0., 0., -1., // row 3
                            0., 0., 1., 0.,
                        ])
                        .unwrap(),
                    structure.into_canonical(),
                )
                .unwrap();
                rot.into()
            }
            RotationMethod::Pi2Y => {
                let rot = DenseTensor::from_storage_data(
                    layout
                        .reorder_to_storage(vec![
                            1., 0., 0., 0., // row 1
                            0., 0., 0., 1., // row 2
                            0., 0., -1., 0., // row 3
                            0., -1., 0., 0.,
                        ])
                        .unwrap(),
                    structure.into_canonical(),
                )
                .unwrap();
                rot.into()
            }
            RotationMethod::Pi2Z => {
                let rot = DenseTensor::from_storage_data(
                    layout
                        .reorder_to_storage(vec![
                            1., 0., 0., 0., // row 1
                            0., 0., -1., 0., // row 2
                            0., 1., 0., 0., // row 3
                            0., 0., 0., -1.,
                        ])
                        .unwrap(),
                    structure.into_canonical(),
                )
                .unwrap();
                rot.into()
            }
            RotationMethod::EulerAngles(alpha, beta, gamma) => DenseTensor::from_storage_data(
                layout
                    .reorder_to_storage(vec![
                        // row 0
                        1.,
                        0.,
                        0.,
                        0.,
                        // row 1
                        0.,
                        -gamma.cos() * beta.cos(),
                        alpha.sin() * beta.sin() * gamma.cos() - alpha.cos() * gamma.sin(),
                        alpha.sin() * gamma.sin() + alpha.cos() * beta.sin() * gamma.cos(),
                        // row 2
                        0.,
                        gamma.sin() * beta.cos(),
                        -alpha.cos() * gamma.cos() - alpha.sin() * beta.sin() * gamma.sin(),
                        -alpha.sin() * gamma.cos() + alpha.cos() * beta.sin() * gamma.sin(),
                        // row 3
                        0.,
                        -beta.sin(),
                        alpha.sin() * beta.cos(),
                        -alpha.cos() * beta.cos(),
                    ])
                    .unwrap(),
                structure.into_canonical(),
            )
            .unwrap()
            .into(),
            // Rotation::EulerAngles(alpha, beta, gamma) => DenseTensor::from_storage_data(
            //     vec![
            //         // row 0
            //         zero.one(),
            //         zero.clone(),
            //         zero.clone(),
            //         zero.clone(),
            //         // row 1
            //         zero.clone(),
            //         alpha.cos() * beta.cos(),
            //         alpha.sin() * beta.sin() * gamma.cos() - alpha.cos() * gamma.sin(),
            //         alpha.sin() * gamma.sin() + alpha.cos() * beta.sin() * gamma.cos(),
            //         // row 2
            //         zero.clone(),
            //         gamma.sin() * beta.cos(),
            //         alpha.cos() * gamma.cos() + alpha.sin() * beta.sin() * gamma.sin(),
            //         -alpha.sin() * gamma.cos() + alpha.cos() * beta.sin() * gamma.sin(),
            //         // row 3
            //         zero.clone(),
            //         -beta.sin(),
            //         alpha.sin() * beta.cos(),
            //         alpha.cos() * beta.cos(),
            //     ],
            //     structure,
            // )
            // .unwrap()
            // .into(),
        }
    }

    pub(crate) fn bispinor_tensor(
        &self,
        i: Slot<LibraryRep>,
        j: Slot<LibraryRep>,
    ) -> DataTensor<Complex<f64>, OrderedStructure> {
        let structure = Canonicalized::from_iter([i.cast::<LibraryRep>(), j.cast()]);
        let layout = TensorDataLayout::from_canonicalized(&structure)
            .expect("bispinor rotations have concrete dimensions");
        let storage_indices = |logical: &[usize]| {
            layout
                .logical_expanded_to_storage_expanded(logical)
                .expect("hard-coded bispinor coordinates are in bounds")
        };
        let zero = 0.; // F::new_zero();
        let zeroc = Complex::new_re(zero);

        match self {
            RotationMethod::Identity => {
                let mut rot = SparseTensor::empty(structure.into_canonical(), zeroc);
                rot.set(&storage_indices(&[0, 0]), zeroc.one()).unwrap();
                rot.set(&storage_indices(&[1, 1]), zeroc.one()).unwrap();
                rot.set(&storage_indices(&[2, 2]), zeroc.one()).unwrap();
                rot.set(&storage_indices(&[3, 3]), zeroc.one()).unwrap();
                rot.into()
            }
            RotationMethod::Pi2X => {
                let mut rot = SparseTensor::empty(structure.into_canonical(), zeroc);
                let sqrt2_half = zero.SQRT_2_HALF();
                let sqrt2_halfim = Complex::new_im(sqrt2_half);
                let sqrt2_halfre = Complex::new_re(sqrt2_half);

                rot.set(&storage_indices(&[0, 0]), sqrt2_halfim).unwrap();
                rot.set(&storage_indices(&[0, 1]), sqrt2_halfim).unwrap();
                rot.set(&storage_indices(&[1, 0]), sqrt2_halfim).unwrap();
                rot.set(&storage_indices(&[1, 1]), sqrt2_halfre).unwrap();
                rot.set(&storage_indices(&[2, 2]), sqrt2_halfim).unwrap();
                rot.set(&storage_indices(&[2, 3]), sqrt2_halfim).unwrap();
                rot.set(&storage_indices(&[3, 2]), sqrt2_halfim).unwrap();
                rot.set(&storage_indices(&[3, 3]), sqrt2_halfre).unwrap();
                rot.into()
            }
            RotationMethod::Pi2Y => {
                let mut rot = SparseTensor::empty(structure.into_canonical(), zeroc);
                let sqrt2_half = zero.SQRT_2_HALF();
                let sqrt2_halfim = Complex::new_im(sqrt2_half);
                let sqrt2_halfre = Complex::new_re(sqrt2_half);
                let nsqrt2_half = -sqrt2_halfre;

                rot.set(&storage_indices(&[0, 0]), sqrt2_halfim).unwrap();
                rot.set(&storage_indices(&[0, 1]), nsqrt2_half).unwrap();
                rot.set(&storage_indices(&[1, 0]), nsqrt2_half).unwrap();
                rot.set(&storage_indices(&[1, 1]), sqrt2_halfre).unwrap();
                rot.set(&storage_indices(&[2, 2]), sqrt2_halfim).unwrap();
                rot.set(&storage_indices(&[2, 3]), nsqrt2_half).unwrap();
                rot.set(&storage_indices(&[3, 2]), nsqrt2_half).unwrap();
                rot.set(&storage_indices(&[3, 3]), sqrt2_halfre).unwrap();
                rot.into()
            }
            RotationMethod::Pi2Z => {
                let mut rot = SparseTensor::empty(structure.into_canonical(), zeroc);
                let sqrt2_half = zero.SQRT_2_HALF();
                let sqrt2_halfc = Complex::new(sqrt2_half, sqrt2_half);
                let sqrt2_halfcc = sqrt2_halfc.conj();

                rot.set(&storage_indices(&[0, 0]), sqrt2_halfc).unwrap();
                rot.set(&storage_indices(&[1, 1]), sqrt2_halfcc).unwrap();
                rot.set(&storage_indices(&[2, 2]), sqrt2_halfc).unwrap();
                rot.set(&storage_indices(&[3, 3]), sqrt2_halfcc).unwrap();
                rot.into()
            }
            RotationMethod::EulerAngles(alpha, beta, gamma) => {
                let norm = alpha.square() + beta.square() + gamma.square();

                if alpha.is_zero() && beta.is_zero() {
                    let normhalf = &norm / 2.; //F::from_f64(2.);
                    let cos_phihalf = normhalf.cos();
                    let sin_phihalf = normhalf.sin();

                    let e = Complex::new(cos_phihalf, sin_phihalf);
                    let econj = e.conj();
                    return DenseTensor::from_storage_data(
                        layout
                            .reorder_to_storage(vec![
                                // row 0
                                e, zeroc, zeroc, zeroc, // row 1
                                zeroc, econj, zeroc, zeroc, // row 2
                                zeroc, zeroc, e, zeroc, // row 3
                                zeroc, zeroc, zeroc, econj,
                            ])
                            .unwrap(),
                        structure.into_canonical(),
                    )
                    .unwrap()
                    .into();
                }

                let complex_phi = Complex::new(*alpha, *beta);

                let normhalf = &norm / 2.; //F::from_f64(2.);
                let cos_phihalf = normhalf.cos();
                let sin_phihalf = normhalf.sin();

                let a_00 = Complex::new(gamma / norm * cos_phihalf, sin_phihalf);
                let a_01 =
                    Complex::new_im((norm - gamma.square() / norm) * sin_phihalf) / complex_phi;
                let a_10 = complex_phi * Complex::new_im(sin_phihalf / norm);
                let a_11 = Complex::new(cos_phihalf, -gamma / norm * sin_phihalf);

                DenseTensor::from_storage_data(
                    layout
                        .reorder_to_storage(vec![
                            // row 0
                            a_00, a_01, zeroc, zeroc, // row 1
                            a_10, a_11, zeroc, zeroc, // row 2
                            zeroc, zeroc, a_00, a_01, // row 3
                            zeroc, zeroc, a_10, a_11,
                        ])
                        .unwrap(),
                    structure.into_canonical(),
                )
                .unwrap()
                .into()
            }
        }
    }
}

pub trait Rotatable {
    fn rotate(&self, rotation: &Rotation) -> Self;
}

impl<T: FloatLike> Rotatable for ThreeMomentum<F<T>> {
    fn rotate(&self, rotation: &Rotation) -> Self {
        match rotation.method {
            RotationMethod::Identity => self.clone(),
            RotationMethod::Pi2X => ThreeMomentum::perform_pi2_rotation_x(self),
            RotationMethod::Pi2Y => ThreeMomentum::perform_pi2_rotation_y(self),
            RotationMethod::Pi2Z => ThreeMomentum::perform_pi2_rotation_z(self),
            RotationMethod::EulerAngles(alpha, beta, gamma) => {
                let mut result = self.clone();
                result.rotate_mut(&F::from_f64(alpha), &F::from_f64(beta), &F::from_f64(gamma));
                result
            }
        }
    }
}

impl<T: FloatLike> Rotatable for ThreeMomentum<HyperDual<F<T>>> {
    fn rotate(&self, rotation: &Rotation) -> Self {
        match rotation.method {
            RotationMethod::Identity => self.clone(),
            RotationMethod::Pi2X => ThreeMomentum::perform_pi2_rotation_x(self),
            RotationMethod::Pi2Y => ThreeMomentum::perform_pi2_rotation_y(self),
            RotationMethod::Pi2Z => ThreeMomentum::perform_pi2_rotation_z(self),
            RotationMethod::EulerAngles(alpha, beta, gamma) => {
                let mut result = self.clone();
                result.rotate_mut(
                    &new_constant(&self.px, &F::from_f64(alpha)),
                    &new_constant(&self.py, &F::from_f64(beta)),
                    &new_constant(&self.pz, &F::from_f64(gamma)),
                );
                result
            }
        }
    }
}

impl<T: FloatLike> ThreeMomentum<HyperDual<F<T>>> {
    fn rotate_mut(
        &mut self,
        alpha: &HyperDual<F<T>>,
        beta: &HyperDual<F<T>>,
        gamma: &HyperDual<F<T>>,
    ) {
        let sin_alpha = alpha.sin();
        let cos_alpha = alpha.cos();
        let sin_beta = beta.sin();
        let cos_beta = beta.cos();
        let sin_gamma = gamma.sin();
        let cos_gamma = gamma.cos();

        let px = self.px.clone();
        let py = self.py.clone();
        let pz = self.pz.clone();

        self.px = cos_gamma.clone() * &cos_beta * &px
            + (-(cos_alpha.clone()) * &sin_gamma + sin_alpha.clone() * &sin_beta * &cos_gamma)
                * &py
            + (sin_alpha.clone() * &sin_gamma + cos_alpha.clone() * &sin_beta * &cos_gamma) * &pz;

        self.py = sin_gamma.clone() * &cos_beta * &px
            + (cos_alpha.clone() * &cos_gamma + sin_alpha.clone() * &sin_beta * &sin_gamma) * &py
            + (-sin_alpha.clone() * &cos_gamma + cos_alpha.clone() * &sin_beta * &sin_gamma) * &pz;

        self.pz =
            -sin_beta * &px + cos_beta.clone() * &sin_alpha * &py + cos_alpha * &cos_beta * &pz;
    }
}

impl<T: FloatLike> Rotatable for FourMomentum<F<T>> {
    fn rotate(&self, rotation: &Rotation) -> Self {
        Self {
            temporal: self.temporal.clone(),
            spatial: self.spatial.rotate(rotation),
        }
    }
}

impl<T: FloatLike + EvaluationDomain> Rotatable for Polarization<Complex<F<T>>> {
    fn rotate(&self, rotation: &Rotation) -> Self {
        let rotated = match self.pol_type {
            PolType::Epsilon | PolType::EpsilonBar => rotation
                .lorentz_rotation
                .clone()
                .map_coeff(&|t| t.map(|f| F::from_ff64(f)))
                .evaluate(&self.tensor.data)
                .try_into_dense()
                .unwrap()
                .cast_structure(),
            PolType::U | PolType::UBar | PolType::V | PolType::VBar | PolType::Scalar => {
                self.tensor.clone()
            }
        };

        Polarization {
            tensor: rotated,
            pol_type: self.pol_type,
        }
    }
}

#[cfg(test)]
mod tests {

    use core::f64;

    use eyre::Context;
    use schemars::schema_for;

    use crate::utils::F;

    use super::*;
    use crate::utils::ApproxEq;

    #[test]
    fn polarization() {
        let mom = FourMomentum::from_args(F(1.), F(1.), F(0.), F(0.));

        let pol: Polarization<F<f64>> = Polarization::lorentz(mom.pol_one());
        println!("{}", pol)

        // let pol2 = pol.clone();

        // print!("{}", pol.add_fallible(&pol2).unwrap());

        // println!("pol_one: {:?}", pol);

        // let structure = pol.tensor.structure.clone();

        // println!("{}", structure.flat_index([2]).unwrap());

        // pol.tensor
        //     .iter_flat()
        //     .for_each(|(i, d)| println!("{}{}", i, d));
    }

    #[test]
    fn serialization() {
        #[derive(Debug, Serialize, Deserialize, PartialEq)]
        struct HelicityToml {
            helicities: Vec<Helicity>,
        }

        let hels = vec![
            Helicity::PLUS,
            Helicity::MINUS,
            Helicity::ZERO,
            Helicity::Summed,
            Helicity::SummedAveraged,
        ];
        let serialized = serde_json::to_string(&hels).unwrap();
        assert_eq!(serialized, r#"[1,-1,0,"summed","summed_averaged"]"#);

        let deserialized: Vec<Helicity> =
            serde_json::from_str(r#"[1,"minus","0","summed","summed-averaged"]"#).unwrap();
        assert_eq!(deserialized, hels);

        let toml_value = toml::from_str::<HelicityToml>(
            r#"
helicities = [1, "minus", "zero", "summed", "summed_averaged"]
"#,
        )
        .unwrap();
        assert_eq!(toml_value.helicities, hels);

        let schema = serde_json::to_value(schema_for!(Helicity)).unwrap();
        let allowed_strings = &schema["enum"];
        assert!(allowed_strings.is_null());
        assert!(
            schema["oneOf"][0]["enum"]
                .as_array()
                .unwrap()
                .contains(&(-1).into())
        );
        assert!(
            schema["oneOf"][1]["enum"]
                .as_array()
                .unwrap()
                .contains(&"summed_averaged".into())
        );
    }

    #[test]
    fn eps() {
        let mom = FourMomentum::from_args(
            F(156.2565490076973),
            F(-108.59017233120495),
            F(-100.2097685717689),
            F(50.81619686346704),
        );

        let pol = mom.eps_pol(SignOrZero::Plus);
        println!("{}", pol);

        let mom = FourMomentum::from_args(
            F(441.7831721921727),
            F(237.7632083614734),
            F(368.76330416225863),
            F(-51.523329523343534),
        );

        let pol = mom.eps_pol(SignOrZero::Plus);
        println!("{}", pol);

        let mom = FourMomentum::from_args(F(485.0355), F(0.), F(0.), F(485.0355));

        let pol = mom.eps_pol(SignOrZero::Plus);

        println!("{}", pol);

        println!("pt{}", mom.pt());

        let mom = FourMomentum::from_args(F(485.0355), F(-107.6044), F(-431.5174), F(193.5805));

        let pol = mom.eps_pol(SignOrZero::Plus).bar();

        // Helicity=           1           1           1           1
        //  W(*,1)=               (5.4707403520912967,0.0000000000000000)               (0.0000000000000000,0.0000000000000000)               (31.622776601683793,0.0000000000000000)               (0.0000000000000000,0.0000000000000000)
        //  W(*,2)=               (0.0000000000000000,0.0000000000000000)             (-0.70710678118654757,0.0000000000000000)              (0.0000000000000000,0.70710678118654757)               (0.0000000000000000,0.0000000000000000)
        //  W(*,3)=               (17.333409072341226,0.0000000000000000)              (6.3994499859138338,-25.663202641304650)               (2.9986797695150327,0.0000000000000000)              (1.1071048475630934,-4.4397340569457056)
        //  W(*,4)=               (0.0000000000000000,0.0000000000000000)        (6.82818793416064135E-002,0.68609707126238750)            (0.27382536157480858,-0.17108713804717862)              (0.64834964048112464,0.0000000000000000)

        println!("{}", pol);
        println!("pt{}", mom.pt());
    }

    #[test]
    fn spinors_degenerate() {
        let mom = FourMomentum::from_args(F(2.), F(0.), F(0.), F(-2.));

        assert_eq!(
            [Complex::new_zero(), Complex::new_re(F(-1.))],
            mom.xi(Sign::Positive)
        );

        assert_eq!(
            [Complex::new_re(F(1.)), Complex::new_re(F(0.))],
            mom.xi(Sign::Negative)
        );

        let u_p = mom.u(Sign::Positive);
        let u_p_target: Polarization<Complex<F<_>>> =
            Polarization::bispinor_u([F(0.), F(0.), F(0.), F(-2.)]).cast();

        u_p.approx_eq_res(&u_p_target, &F(0.001))
            .wrap_err("u+((2,0,0,-2)) does not match target: (0,0,0,-2)")
            .unwrap();

        let mut u_p_bar_target: Polarization<Complex<F<_>>> =
            Polarization::bispinor_u([F(0.), F(-2.), F(0.), F(0.)]).cast();

        u_p_bar_target.pol_type = PolType::UBar;

        u_p.bar()
            .approx_eq_res(&u_p_bar_target, &F(0.001))
            .wrap_err("u+bar((2,0,0,-2)) does not match target: (0,-2,0,0)")
            .unwrap();

        let u_m = mom.u(Sign::Negative);
        let u_m_target: Polarization<Complex<F<_>>> =
            Polarization::bispinor_u([F(2.), F(0.), F(0.), F(0.)]).cast();

        u_m.approx_eq_res(&u_m_target, &F(0.001))
            .wrap_err("u-((2,0,0,-2)) does not match target: (2,0,0,0)")
            .unwrap();

        let mut u_m_bar_target: Polarization<Complex<F<_>>> =
            Polarization::bispinor_u([F(0.), F(0.), F(2.), F(0.)]).cast();

        u_m_bar_target.pol_type = PolType::UBar;

        u_m.bar()
            .approx_eq_res(&u_m_bar_target, &F(0.001))
            .wrap_err("u-bar((2,0,0,-2)) does not match target: (0,0,2,0)")
            .unwrap();

        let v_p = mom.v(Sign::Positive);
        let v_p_target: Polarization<Complex<F<_>>> =
            Polarization::bispinor_v([F(-2.), F(0.), F(0.), F(0.)]).cast();

        v_p.approx_eq_res(&v_p_target, &F(0.001))
            .wrap_err("v+((2,0,0,-2)) does not match target: (-2,0,0,0)")
            .unwrap();

        let mut v_p_bar_target: Polarization<Complex<F<_>>> =
            Polarization::bispinor_v([F(0.), F(0.), F(-2.), F(0.)]).cast();

        v_p_bar_target.pol_type = PolType::VBar;

        v_p.bar()
            .approx_eq_res(&v_p_bar_target, &F(0.001))
            .wrap_err("v+bar((2,0,0,-2)) does not match target: (0,0,-2,0)")
            .unwrap();

        let v_m = mom.v(Sign::Negative);
        let v_m_target: Polarization<Complex<F<_>>> =
            Polarization::bispinor_v([F(0.), F(0.), F(0.), F(2.)]).cast();

        v_m.approx_eq_res(&v_m_target, &F(0.001))
            .wrap_err("v-((2,0,0,-2)) does not match target: (0,0,0,2)")
            .unwrap();

        let mut v_m_bar_target: Polarization<Complex<F<_>>> =
            Polarization::bispinor_v([F(0.), F(2.), F(0.), F(0.)]).cast();

        v_m_bar_target.pol_type = PolType::VBar;

        v_m.bar()
            .approx_eq_res(&v_m_bar_target, &F(0.001))
            .wrap_err("v-bar((2,0,0,-2)) does not match target: (0,2,0,0)")
            .unwrap();
    }

    #[test]
    fn spinors() {
        let mom = FourMomentum::from_args(F(4.), F(0.), F(4.), F(0.));

        assert!(
            Complex::approx_eq_slice(
                &[
                    Complex::new_re(F(f64::consts::FRAC_1_SQRT_2)),
                    Complex::new_im(F(f64::consts::FRAC_1_SQRT_2)),
                ],
                &mom.xi(Sign::Positive),
                &F(0.001),
            ),
            "xi+((4,0,4,0)) does not match target: (1/sqrt(2),i/sqrt(2))"
        );

        assert!(
            Complex::approx_eq_slice(
                &[
                    Complex::new_im(F(f64::consts::FRAC_1_SQRT_2)),
                    Complex::new_re(F(f64::consts::FRAC_1_SQRT_2)),
                ],
                &mom.xi(Sign::Negative),
                &F(0.001),
            ),
            "xi-((4,0,4,0)) does not match target: (i/sqrt(2),1/sqrt(2))"
        );

        let zero: Complex<_> = F(0.).into();
        let u_p = mom.u(Sign::Positive);
        let u_p_target: Polarization<Complex<F<_>>> =
            Polarization::bispinor_u([zero, zero, Complex::new_re(F(2.)), Complex::new_im(F(2.))])
                .cast();

        u_p.approx_eq_res(&u_p_target, &F(0.001))
            .wrap_err("u+((4,0,4,0)) does not match target: (0,0,2,i2)")
            .unwrap();

        let mut u_p_bar_target: Polarization<Complex<F<_>>> =
            Polarization::bispinor_u([Complex::new_re(F(2.)), Complex::new_im(-F(2.)), zero, zero])
                .cast();

        u_p_bar_target.pol_type = PolType::UBar;

        u_p.bar()
            .approx_eq_res(&u_p_bar_target, &F(0.001))
            .wrap_err("u+bar((4,0,4,0)) does not match target: (2,-2i,0,0)")
            .unwrap();

        let u_m = mom.u(Sign::Negative);
        let u_m_target: Polarization<Complex<F<_>>> =
            Polarization::bispinor_u([Complex::new_im(F(2.)), Complex::new_re(F(2.)), zero, zero])
                .cast();

        u_m.approx_eq_res(&u_m_target, &F(0.001))
            .wrap_err("u-((4,0,4,0)) does not match target: (2i,2,0,0)")
            .unwrap();

        let mut u_m_bar_target: Polarization<Complex<F<_>>> =
            Polarization::bispinor_u([zero, zero, Complex::new_im(-F(2.)), Complex::new_re(F(2.))])
                .cast();

        u_m_bar_target.pol_type = PolType::UBar;

        u_m.bar()
            .approx_eq_res(&u_m_bar_target, &F(0.001))
            .wrap_err("u-bar((4,0,4,0)) does not match target: (0,0,-2i,2)")
            .unwrap();

        let v_p = mom.v(Sign::Positive);
        let v_p_target: Polarization<Complex<F<_>>> = Polarization::bispinor_v([
            Complex::new_im(-F(2.)),
            Complex::new_re(-F(2.)),
            zero,
            zero,
        ])
        .cast();

        v_p.approx_eq_res(&v_p_target, &F(0.001))
            .wrap_err("v+((4,0,4,0)) does not match target: (-2i,-2,0,0)")
            .unwrap();

        let mut v_p_bar_target: Polarization<Complex<F<_>>> =
            Polarization::bispinor_v([zero, zero, Complex::new_im(F(2.)), Complex::new_re(F(-2.))])
                .cast();

        v_p_bar_target.pol_type = PolType::VBar;

        v_p.bar()
            .approx_eq_res(&v_p_bar_target, &F(0.001))
            .wrap_err("v+bar((4,0,4,0)) does not match target: (0,0,2i,-2)")
            .unwrap();

        let v_m = mom.v(Sign::Negative);
        let v_m_target: Polarization<Complex<F<_>>> = Polarization::bispinor_v([
            zero,
            zero,
            Complex::new_re(-F(2.)),
            Complex::new_im(-F(2.)),
        ])
        .cast();

        v_m.approx_eq_res(&v_m_target, &F(0.001))
            .wrap_err("v-((4,0,4,0)) does not match target: (0,0,-2,-2i)")
            .unwrap();

        let mut v_m_bar_target: Polarization<Complex<F<_>>> =
            Polarization::bispinor_v([Complex::new_re(-F(2.)), Complex::new_im(F(2.)), zero, zero])
                .cast();

        v_m_bar_target.pol_type = PolType::VBar;

        v_m.bar()
            .approx_eq_res(&v_m_bar_target, &F(0.001))
            .wrap_err("v-bar((4,0,4,0)) does not match target: (-2,2i,0,0)")
            .unwrap();
    }

    #[test]
    fn vectors_degenerate() {
        let mom = FourMomentum::from_args(F(2.), F(0.), F(0.), F(-2.));

        assert!(
            F::approx_eq_slice(&mom.pol_one(), &[0., 1., 0., 0.].map(F), &F(0.001)),
            "pol_one((2,0,0,-2)) does not match target: (0,1,0,0)"
        );

        assert!(
            F::approx_eq_slice(&mom.pol_two(), &[0., 0., -1., 0.].map(F), &F(0.001)),
            "pol_two((2,0,0,-2)) does not match target: (0,0,-1,0)"
        );

        let e_p = mom.eps_pol(SignOrZero::Plus);

        let zero = F(0.).into();
        let e_p_target = Polarization::lorentz([
            zero,
            Complex::new_re(F(-f64::consts::FRAC_1_SQRT_2)),
            Complex::new_im(F(f64::consts::FRAC_1_SQRT_2)),
            zero,
        ]);

        let mut e_p_bar_target = Polarization::lorentz([
            zero,
            Complex::new_re(F(-f64::consts::FRAC_1_SQRT_2)),
            Complex::new_im(F(-f64::consts::FRAC_1_SQRT_2)),
            zero,
        ]);

        e_p_bar_target.pol_type = PolType::EpsilonBar;

        e_p.approx_eq_res(&e_p_target, &F(0.0001))
            .wrap_err("e+((2,0,0,-2)) does not match target: (0,-1/sqrt(2),i/sqrt(2),0)")
            .unwrap();

        e_p.bar()
            .approx_eq_res(&e_p_bar_target, &F(0.0001))
            .wrap_err("e+bar((2,0,0,-2)) does not match target: (0,-1/sqrt(2),-i/sqrt(2),0)")
            .unwrap();

        let e_m = mom.eps_pol(SignOrZero::Minus);

        let e_m_target = Polarization::lorentz([
            zero,
            Complex::new_re(F(f64::consts::FRAC_1_SQRT_2)),
            Complex::new_im(F(f64::consts::FRAC_1_SQRT_2)),
            zero,
        ]);

        let mut e_m_bar_target = Polarization::lorentz([
            zero,
            Complex::new_re(F(f64::consts::FRAC_1_SQRT_2)),
            Complex::new_im(F(-f64::consts::FRAC_1_SQRT_2)),
            zero,
        ]);

        e_m_bar_target.pol_type = PolType::EpsilonBar;

        e_m.approx_eq_res(&e_m_target, &F(0.0001))
            .wrap_err("e-((2,0,0,-2)) does not match target: (0,1/sqrt(2),i/sqrt(2),0)")
            .unwrap();

        e_m.bar()
            .approx_eq_res(&e_m_bar_target, &F(0.0001))
            .wrap_err("e-bar((2,0,0,-2)) does not match target: (0,1/sqrt(2),-i/sqrt(2),0)")
            .unwrap();
    }

    #[test]
    fn vectors() {
        let mom = FourMomentum::from_args(F(2.), F(0.), F(0.), F(-2.));

        assert!(
            F::approx_eq_slice(&mom.pol_one(), &[0., 1., 0., 0.].map(F), &F(0.001)),
            "pol_one((2,0,0,-2)) does not match target: (0,1,0,0)"
        );

        assert!(
            F::approx_eq_slice(&mom.pol_two(), &[0., 0., -1., 0.].map(F), &F(0.001)),
            "pol_two((2,0,0,-2)) does not match target: (0,0,-1,0)"
        );

        let e_p = mom.eps_pol(SignOrZero::Plus);

        let zero = F(0.).into();
        let e_p_target = Polarization::lorentz([
            zero,
            Complex::new_re(F(-f64::consts::FRAC_1_SQRT_2)),
            Complex::new_im(F(f64::consts::FRAC_1_SQRT_2)),
            zero,
        ]);

        let mut e_p_bar_target = Polarization::lorentz([
            zero,
            Complex::new_re(F(-f64::consts::FRAC_1_SQRT_2)),
            Complex::new_im(F(-f64::consts::FRAC_1_SQRT_2)),
            zero,
        ]);

        e_p_bar_target.pol_type = PolType::EpsilonBar;

        e_p.approx_eq_res(&e_p_target, &F(0.0001))
            .wrap_err("e+((2,0,0,-2)) does not match target: (0,-1/sqrt(2),i/sqrt(2),0)")
            .unwrap();

        e_p.bar()
            .approx_eq_res(&e_p_bar_target, &F(0.0001))
            .wrap_err("e+bar((2,0,0,-2)) does not match target: (0,-1/sqrt(2),-i/sqrt(2),0)")
            .unwrap();

        let e_m = mom.eps_pol(SignOrZero::Minus);

        let e_m_target = Polarization::lorentz([
            zero,
            Complex::new_re(F(f64::consts::FRAC_1_SQRT_2)),
            Complex::new_im(F(f64::consts::FRAC_1_SQRT_2)),
            zero,
        ]);

        let mut e_m_bar_target = Polarization::lorentz([
            zero,
            Complex::new_re(F(f64::consts::FRAC_1_SQRT_2)),
            Complex::new_im(F(-f64::consts::FRAC_1_SQRT_2)),
            zero,
        ]);

        e_m_bar_target.pol_type = PolType::EpsilonBar;

        e_m.approx_eq_res(&e_m_target, &F(0.0001))
            .wrap_err("e-((2,0,0,-2)) does not match target: (0,1/sqrt(2),i/sqrt(2),0)")
            .unwrap();

        e_m.bar()
            .approx_eq_res(&e_m_bar_target, &F(0.0001))
            .wrap_err("e-bar((2,0,0,-2)) does not match target: (0,1/sqrt(2),-i/sqrt(2),0)")
            .unwrap();
    }

    #[test]
    fn rotations() {
        let mom = FourMomentum::from_args(F(2.), F(3.), F(1.), F(2.));

        let momc: [F<f64>; 4] = mom.into();
        let e = Polarization::lorentz(momc.map(Complex::new_re));
        let u = mom.u(Sign::Negative);

        let rot: Rotation = RotationMethod::EulerAngles(std::f64::consts::PI / 2., 0., 0.).into();
        let rotx = RotationMethod::Pi2X.into();

        let rotye: Rotation = RotationMethod::EulerAngles(0., std::f64::consts::PI / 2., 0.).into();
        let roty = RotationMethod::Pi2Y.into();

        let rotze: Rotation = RotationMethod::EulerAngles(0., 0., std::f64::consts::PI / 2.).into();
        let rotz = RotationMethod::Pi2Z.into();

        let rotid = RotationMethod::Identity.into();
        // println!(
        //     "{}",
        //     rot.bispinor_tensor(
        //         Bispinor::new_slot_selfless(4, 2),
        //         Bispinor::new_slot_selfless(4, 1),
        //     )
        // );

        let mom_rot = mom.rotate(&rot);
        let mom_id = mom.rotate(&rotid);
        let other_mom_rot = mom.rotate(&rotx);
        let mom_rotye = mom.rotate(&rotye);
        let mom_roty = mom.rotate(&roty);
        let mom_rotze = mom.rotate(&rotze);
        let mom_rotz = mom.rotate(&rotz);

        println!("id:\t{}", mom_id);
        mom_id
            .approx_eq_res(&mom, &F(0.001))
            .wrap_err("mom is not identical to mom transformed with identity")
            .unwrap();
        println!("orig:\t{}", mom);

        println!("rotxe:\t{}", mom_rot);
        println!("pi2x:\t{}", other_mom_rot);
        mom_rot.approx_eq_res(&other_mom_rot, &F(0.001)).unwrap();

        println!("rotye:\t{}", mom_rotye);
        println!("pi2y:\t{}", mom_roty);

        mom_rotye.approx_eq_res(&mom_roty, &F(0.001)).unwrap();
        println!("rotze:\t{}", mom_rotze);
        println!("pi2z:\t{}", mom_rotz);
        mom_rotze.approx_eq_res(&mom_rotz, &F(0.001)).unwrap();

        let e_rot = e.rotate(&rot);
        println!("rotxe:{}", e_rot);
        mom_rot.approx_eq_res(&e_rot, &F(0.001)).unwrap();

        let e_id = e.rotate(&rotid);
        mom_id.approx_eq_res(&e_id, &F(0.001)).unwrap();
        let e_rotye = e.rotate(&rotye);

        println!("rotye:{}", e_rotye);
        mom_rotye.approx_eq_res(&e_rotye, &F(0.001)).unwrap();
        let e_roty = e.rotate(&roty);
        mom_roty.approx_eq_res(&e_roty, &F(0.001)).unwrap();
        let e_rotze = e.rotate(&rotze);
        mom_rotze.approx_eq_res(&e_rotze, &F(0.001)).unwrap();
        let e_rotz = e.rotate(&rotz);
        mom_rotz.approx_eq_res(&e_rotz, &F(0.001)).unwrap();
        let u_rot = u.rotate(&rot);
        let u_id = u.rotate(&rotid);

        let other_e_rot = e.rotate(&rotx);

        let other_u_rot = u.rotate(&rotx);

        println!("id:{}", e_id);
        println!("orig:{}", e);
        println!("euler:{}", e_rot);
        println!("Pi2X:{}", other_e_rot);
        println!("rotye:{}", e_rotye);
        println!("pi2y:{}", e_roty);
        println!("rotze:{}", e_rotze);
        println!("pi2z:{}", e_rotz);

        println!("id:{}", u_id);
        println!("orig:{}", u);
        println!("euler:{}", u_rot);
        println!("Pi2X {}", other_u_rot);
    }

    // #[test]
    // fn omega() {
    //     let mu = PhysReps::new_slot(Minkowski {}.into(), 4, 0);

    //     let nu = PhysReps::new_slot(Minkowski {}.into(), 4, 1);

    //     let mud = mu.dual();

    //     let nud = nu.dual();

    //     let i = PhysReps::new_slot(GR.bis.into(), 4, 2);

    //     let j = PhysReps::new_slot(GR.bis.into(), 4, 3);

    //     let k = PhysReps::new_slot(GR.bis.into(), 4, 4);

    //     let zero = Atom::num(0);
    //     let theta_x = parse!("theta_x").unwrap();
    //     let theta_y = parse!("theta_y").unwrap();
    //     let theta_z = parse!("theta_z").unwrap();

    //     let omega = DenseTensor::from_storage_data(
    //         vec![
    //             zero.clone(),
    //             zero.clone(),
    //             zero.clone(),
    //             zero.clone(),
    //             //
    //             zero.clone(),
    //             zero.clone(),
    //             theta_z.clone(),
    //             -theta_y.clone(),
    //             //
    //             zero.clone(),
    //             -theta_z.clone(),
    //             zero.clone(),
    //             theta_x.clone(),
    //             //
    //             zero.clone(),
    //             theta_y.clone(),
    //             -theta_x.clone(),
    //             zero.clone(),
    //         ],
    //         OrderedStructure::from_iter([mu, nu]),
    //     )
    //     .unwrap();

    //     println!("{}", omega);

    //     let gammamujk: SparseTensor<Complex<f64>> =
    //         ufo::gamma_data_weyl(OrderedStructure::from_iter([mud, j, k]));

    //     let gammamukj: SparseTensor<Complex<f64>> =
    //         ufo::gamma_data_weyl(OrderedStructure::from_iter([mud, k, i]));

    //     let gammanuki: SparseTensor<Complex<f64>> =
    //         ufo::gamma_data_weyl(OrderedStructure::from_iter([nud, k, i]));

    //     let gammanujk: SparseTensor<Complex<f64>> =
    //         ufo::gamma_data_weyl(OrderedStructure::from_iter([nud, j, k]));

    //     let sigma = (gammamujk
    //         .contract(&gammanuki)
    //         .unwrap()
    //         .sub_fallible(&gammanujk.contract(&gammamukj).unwrap())
    //         .unwrap())
    //     .scalar_mul(&Complex::new_im(0.25))
    //     .unwrap();

    //     println!("{}", sigma);

    //     println!("{}", omega.contract(&sigma).unwrap());
    // }
}