qudit-expr 0.3.4

use serde::{Deserialize, Serialize};

/// A unique identifier for a tensor index.
///
/// Tensor indices are identified in two distinct locations by this type alias:
///
/// 1) Local Tensor indices: `TensorExpression` objects generate tensor data
///    in a specific order. The ordering of their local tensor indices is identified
///    by their local index ids. These do not change unless the `TensorExpression`'s
///    internal symbolic expression changes. Therefore, it is important to track
///    it's local tensor index ordering, so we can accurately perform contractions
///    dictated by the network. See `[TensorIndex]` for more information.
///
/// 2) Global Tensor Network indices: The entire network will have a list of
///    indices and each will have its own id. See `[NetworkIndex]` for more
///    information.
pub type IndexId = usize;

/// Represents the size or dimension of an index in a tensor network.
///
/// This type alias is used to specify the number of possible values for an index.
/// From a quantum circuit perspective, Input and Output indices usually have
/// size equal to radix of the qudit captured by the leg (2 for qubits, 3 for
/// qutrits, etc); however, during contraction simplification, many edges can
/// be grouped together to simplify the number of contractions, leading to a
/// larger index size.
pub type IndexSize = usize;

/// Represents a weighted-indexed edge in a tensor network.
///
/// This provides the core index information for contraction ordering solvers.
pub type WeightedIndex = (IndexId, IndexSize);

/// IndexDirection represents a tensor's leg direction from the quantum circuit perspective.
///
/// In the context of quantum circuits and tensor networks, indices (or "legs") of a tensor
/// are categorized by their direction to signify their role in an operation.
///
/// While tensors may conceptually have multiple indices (legs), in the `qudit-expr` library,
/// `TensorExpression` objects are always generated as 1-D, 2-D, or 3-D tensors. This enum
/// links the conceptual/graphical tensor indices to the generated expression's dimensions by grouping
/// them along directions. So all `Input`, `Output`, and `Batch` legs correspond to a
/// specific dimension of the generated tensor. This tensor network library can
/// use this to predict and manipulate the shape of a generated tensor.
///
/// For example, kets capture quantum states, and are represented conceptually
/// by column vectors. In the tensor network perspective, they are represented by
/// a tensor with only `Output` indices. They have as many indices as there are
/// qudits in the state and the size of each index is the radix of the qudit.
/// For qubits, the radix is 2, so all indices would have dimension 2. However,
/// when generated in the `qudit-expr` library, they are generated as a vector.
/// The `IndexDirection` is what is used by the `qudit-tensor` library to expect
/// a vector shape, since a ket tensor will only have indices in one direction.
///
/// Similarly, gates which capture quantum operations, are represented by unitary
/// matrices. As a tensor, they have indices that point in both `Input` and `Output`
/// directions. Even though there will likely be more than 2 indices, the `qudit-tensor`
/// library expects a matrix to be generated by the underlying `TensorExpression`
/// because of the two distinct directions.
///
/// Batch indices are used for measurements/Kraus operators or for batch executions
/// of a network.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
pub enum IndexDirection {
    /// For gradients and hessians
    Derivative,

    /// Batch index of a tensor. These are typically represented in quantum circuits
    /// through interactions with classical systems through measurements and/or
    /// feed-forward operations. Additionally, they can capture batch computations
    /// of a network.
    Batch,

    /// Output leg of a tensor, corresponds to rows of an operation. In circuit
    /// diagrams, these represent the wires leaving a gate or ket.
    Output,

    /// Input leg of a tensor, corresponds to columns of an operation. In circuit
    /// diagrams, these represent the wires going into a gate or bra.
    Input,
}

impl std::fmt::Display for IndexDirection {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            IndexDirection::Input => write!(f, "Input"),
            IndexDirection::Output => write!(f, "Output"),
            IndexDirection::Batch => write!(f, "Batch"),
            IndexDirection::Derivative => write!(f, "Derivative"),
        }
    }
}

impl IndexDirection {
    pub fn is_derivative(&self) -> bool {
        matches!(self, IndexDirection::Derivative)
    }
    pub fn is_batch(&self) -> bool {
        matches!(self, IndexDirection::Batch)
    }
    pub fn is_output(&self) -> bool {
        matches!(self, IndexDirection::Output)
    }
    pub fn is_input(&self) -> bool {
        matches!(self, IndexDirection::Input)
    }
}

/// Represents an index of a tensor.
///
/// Indices in this format are associated with a direction, and as a result,
/// an expected output or expected generation (input) shape and ordering.
/// For example, a `[super::network::QuditTensorNetwork]` will store its output
/// indices in this format because after evaluation, there will
/// be an expected shape and ordering. The shape is determined by all the
/// indices directions and sizes, and the ordering will be determined by the ids.
///
/// # See Also
///
/// - `[super::tensor::QuditTensor]` - Tensor object capturing how tensors are generated.
/// - `[super::network::QuditTensorNetwork]` - Tensor network describing a desired calculation.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
pub struct TensorIndex {
    /// The direction of the tensor leg. Important for mapping to dimensions
    /// of a TensorExpression's generated data.
    direction: IndexDirection,

    /// A unique identifier for the index associated with this leg.
    index_id: IndexId,

    /// The size or dimension of the index associated with this leg. This `index_size`
    /// specifies the number of possible values that the index can take.
    index_size: IndexSize,
}

impl TensorIndex {
    /// Creates a new `TensorIndex` instance.
    ///
    /// # Arguments
    ///
    /// * `direction` - The direction of the tensor leg.
    /// * `index_id` - A unique identifier for the index.
    /// * `index_size` - The dimension or size of the index.
    ///
    /// # Returns
    ///
    /// A new `TensorIndex` instance with the specified properties.
    ///
    /// # Examples
    ///
    /// ```
    /// use qudit_expr::index::{IndexDirection, TensorIndex};
    ///
    /// let idx = TensorIndex::new(IndexDirection::Input, 0, 2);
    /// assert_eq!(idx.index_id(), 0);
    /// assert_eq!(idx.index_size(), 2);
    /// assert_eq!(idx.direction(), IndexDirection::Input);
    /// ```
    pub fn new(direction: IndexDirection, index_id: IndexId, index_size: IndexSize) -> Self {
        Self {
            direction,
            index_id,
            index_size,
        }
    }

    /// Returns the `IndexDirection` of the tensor leg.
    ///
    /// # Examples
    ///
    /// ```
    /// use qudit_expr::index::{IndexDirection, TensorIndex};
    ///
    /// let leg = TensorIndex::new(IndexDirection::Output, 1, 3);
    /// assert_eq!(leg.direction(), IndexDirection::Output);
    /// ```
    pub fn direction(&self) -> IndexDirection {
        self.direction
    }

    /// Returns the `IndexId` of the index associated with this leg.
    ///
    /// # Examples
    ///
    /// ```
    /// use qudit_expr::index::{IndexDirection, TensorIndex};
    ///
    /// let leg = TensorIndex::new(IndexDirection::Batch, 2, 4);
    /// assert_eq!(leg.index_id(), 2);
    /// ```
    pub fn index_id(&self) -> IndexId {
        self.index_id
    }

    /// Returns the `IndexSize` of the index associated with this leg.
    ///
    /// # Examples
    ///
    /// ```
    /// use qudit_expr::index::{IndexDirection, TensorIndex};
    ///
    /// let leg = TensorIndex::new(IndexDirection::Input, 3, 5);
    /// assert_eq!(leg.index_size(), 5);
    /// ```
    pub fn index_size(&self) -> IndexSize {
        self.index_size
    }
}

impl std::fmt::Display for TensorIndex {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "Leg {}, (size={}, dir={})",
            self.index_id, self.index_size, self.direction,
        )
    }
}