topos 0.13.1

use std::fmt;
use std::ops::{Add, Div, Mul, Neg, Sub};
use std::ptr;

use static_assertions::assert_impl_all;

use crate::{Element, MapOperation, Shape, Tensor};

use crate::op::Op;

use super::{Symbol, Tape};

// Entry-time contract: proxies stay thread-safe and `Copy`; the anchor
// rationale is documented in `network.rs`.
assert_impl_all!(Value<'static, f64>: Send, Sync, Copy);

/// A lightweight, `Copy` handle to a value recorded on a `Tape`.
///
/// It is an index into the tape's columns rather than a pointer, so handles
/// are cheap to copy and carry no ownership of tape memory.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub(crate) struct ValueId(pub(crate) usize);

impl ValueId {
    /// Returns the position of the value on its tape.
    pub(crate) fn index(self) -> usize {
        self.0
    }
}

/// A `Copy` proxy to a value recorded on a [`Tape`]: the operand of
/// recording.
///
/// A value stores its node position together with a borrow of the tape, so
/// it cannot outlive the construction phase — [`Tape::into_network`]
/// consumes the tape, and the borrow checker rejects a proxy that would
/// cross the seal; take [`Value::symbol`] first. Arithmetic and tensor
/// operations append computed nodes to the tape without consuming their
/// operands. Payload literals can be mixed directly into expressions, in
/// either operand order; every literal occurrence records a new leaf.
///
/// Operations validate tape identity and shape compatibility when they are
/// recorded, so invalid expressions panic before a forward run begins.
///
/// The methods in this file are opcode mnemonics: each records exactly one
/// computed node, one per `Op` variant (payload literals additionally
/// record a leaf, which is data injection rather than computation). Methods
/// that expand to several computed nodes are composites and live in the
/// composition tier of `composite.rs`.
///
/// [`Value::shape`] returns the shape inferred when the node was recorded.
/// [`Value::payload`] clones the stored payload of a leaf, parameter, or
/// input; computed values are read from a [`Run`](crate::Run), live
/// parameter payloads from [`Parameters`](crate::Parameters) — both by
/// [`Symbol`].
pub struct Value<'tape, E> {
    tape: &'tape Tape<E>,
    id: ValueId,
}

/// # Names and reads
///
/// The proxy's identity: its detached [`Symbol`], its recording
/// tape, its inferred shape, and its stored payload when it is a
/// source.
impl<'tape, E: Element> Value<'tape, E> {
    /// Binds a proxy to the node `id` recorded on `tape`.
    pub(crate) fn bind(tape: &'tape Tape<E>, id: ValueId) -> Self {
        Self { tape, id }
    }

    /// Returns the handle of the node this proxy points to.
    pub(crate) fn id(&self) -> ValueId {
        self.id
    }

    /// Returns the tape this proxy records onto: the recording phase
    /// *is* this operand, which is what lets a module express itself
    /// from its input alone.
    pub fn tape(&self) -> &'tape Tape<E> {
        self.tape
    }

    /// Returns the detached name of this value: the currency of every
    /// phase after recording, and the documented bridge across
    /// [`Tape::into_network`].
    pub fn symbol(&self) -> Symbol {
        Symbol {
            origin: self.tape.origin(),
            id: self.id,
        }
    }
}

/// The conversion form of [`Value::symbol`], for positions where a
/// list must be homogeneous in `Symbol`: `[loss.into(), stored]`.
impl<E: Element> From<Value<'_, E>> for Symbol {
    fn from(value: Value<'_, E>) -> Symbol {
        value.symbol()
    }
}

impl<'tape, E: Element> Value<'tape, E> {
    /// Returns a clone of the `Op` that produced this value.
    #[cfg(test)]
    pub(crate) fn op(&self) -> Op<Tensor<E>> {
        self.tape.with_node(self.id, |op| op.clone())
    }

    /// Returns the operand links of this value's node.
    #[cfg(test)]
    pub(crate) fn operands(&self) -> Vec<ValueId> {
        self.tape.operands_of(self.id).as_slice().to_vec()
    }

    /// Returns the shape of this value, inferred when it was recorded.
    pub fn shape(&self) -> Shape {
        self.tape.shape(self.id)
    }

    /// Returns a clone of this node's stored payload, or `None` for a computed
    /// value.
    ///
    /// Leaves return their recorded payload, parameters their record-site
    /// initial, and inputs their recorded default. Live parameter payloads
    /// are read from [`Parameters::of`](crate::Parameters::of), run results
    /// from [`Run::of`](crate::Run::of).
    pub fn payload(&self) -> Option<Tensor<E>> {
        self.tape.payload_of(self.id)
    }

    /// Records a computed node produced by `op` over the positional
    /// `operands` on the same tape and returns a proxy to it.
    fn apply(&self, op: Op<Tensor<E>>, operands: &[ValueId]) -> Self {
        let id = self.tape.record_node(op, operands);
        Self::bind(self.tape, id)
    }

    /// Records `data` as a fresh leaf on the same tape and returns a
    /// proxy to it.
    ///
    /// It backs the payload-literal operator sugar: every literal
    /// appearance records its own leaf.
    pub(crate) fn literal(&self, data: Tensor<E>) -> Self {
        Self::bind(self.tape, self.tape.record_node(Op::leaf(data), &[]))
    }

    /// Panics if `other` belongs to a different tape.
    ///
    /// The one runtime check the proxy keeps: coexisting tapes cannot
    /// be told apart by lifetimes alone, so mixing their proxies in
    /// one operator panics at the recording expression.
    fn assert_same_tape(&self, other: &Self) {
        assert!(
            ptr::eq(self.tape, other.tape),
            "values belong to different tapes"
        );
    }
}

/// # Elementary maps
///
/// One recorded node per call: the transcendentals and the order
/// pair. Arithmetic records through the standard operators.
impl<'tape, E: Element> Value<'tape, E> {
    /// Records the hyperbolic tangent of this value on the same tape
    /// and returns a proxy to it.
    pub fn tanh(self) -> Self {
        self.apply(Op::map(MapOperation::Tanh), &[self.id])
    }

    /// Records the exponential of this value on the same tape and
    /// returns a proxy to it.
    pub fn exp(self) -> Self {
        self.apply(Op::map(MapOperation::Exp), &[self.id])
    }

    /// Records the natural logarithm of this value on the same tape
    /// and returns a proxy to it.
    pub fn ln(self) -> Self {
        self.apply(Op::map(MapOperation::Ln), &[self.id])
    }

    /// Records the square root of this value on the same tape and
    /// returns a proxy to it.
    pub fn sqrt(self) -> Self {
        self.apply(Op::map(MapOperation::Sqrt), &[self.id])
    }

    /// Records the sine of this value on the same tape and returns a
    /// proxy to it.
    pub fn sin(self) -> Self {
        self.apply(Op::map(MapOperation::Sin), &[self.id])
    }

    /// Records the cosine of this value on the same tape and returns
    /// a proxy to it.
    pub fn cos(self) -> Self {
        self.apply(Op::map(MapOperation::Cos), &[self.id])
    }

    /// Records the natural logarithm of one plus this value on the
    /// same tape and returns a proxy to it.
    ///
    /// It is a distinct opcode, not sugar for `(one + x).ln()`: the
    /// composed form rounds `1 + x` first and destroys every
    /// significant digit of an `x` near zero, while this one stays
    /// accurate there. The two spellings are different specs with
    /// different bits, and both remain valid.
    pub fn log1p(self) -> Self {
        self.apply(Op::map(MapOperation::Log1p), &[self.id])
    }

    /// Records `e` raised to this value, minus one, on the same tape
    /// and returns a proxy to it.
    ///
    /// Like [`log1p`](Self::log1p), it is a distinct opcode: the
    /// composed `x.exp() - one` cancels catastrophically near zero,
    /// and this one does not.
    pub fn expm1(self) -> Self {
        self.apply(Op::map(MapOperation::Expm1), &[self.id])
    }

    /// Records the error function of this value on the same tape and
    /// returns a proxy to it.
    ///
    /// The computation delegates to the pure-Rust `libm` crate, like
    /// every transcendental; its derivative rule speaks
    /// [`erf_derivative`](Self::erf_derivative), the closed pair that
    /// keeps the constant `2/sqrt(pi)` inside per-element kernels
    /// rather than in any recorded graph.
    pub fn erf(self) -> Self {
        self.apply(Op::map(MapOperation::Erf), &[self.id])
    }

    /// Records the derivative of the error function of this value —
    /// the scaled Gaussian `(2/sqrt(pi)) * e^(-x^2)` — on the same
    /// tape and returns a proxy to it: what `differentiate` emits
    /// where the engine's rules call
    /// [`Elementary::erf_derivative`](crate::Elementary::erf_derivative).
    /// Its own derivative is `-2x` times itself, so the pair closes
    /// under differentiation.
    pub fn erf_derivative(self) -> Self {
        self.apply(Op::map(MapOperation::ErfDerivative), &[self.id])
    }

    /// Records this value raised elementwise to the power of `exponent`
    /// on the same tape and returns a proxy to it.
    ///
    /// The exponent-side gradient involves the logarithm of this value,
    /// so it is a number only where this value is positive.
    ///
    /// # Panics
    /// Panics if the operands belong to different tapes or their
    /// shapes differ.
    pub fn powf(self, exponent: Self) -> Self {
        self.assert_same_tape(&exponent);
        self.apply(Op::powf(), &[self.id, exponent.id])
    }

    /// Records the elementwise maximum of this value and `rhs` on the
    /// same network and returns a proxy to it; on a tie the gradient goes
    /// to this value, not `rhs`.
    ///
    /// # Panics
    /// Panics if the operands belong to different tapes or their
    /// shapes differ.
    pub fn maximum(self, rhs: Self) -> Self {
        self.assert_same_tape(&rhs);
        self.apply(Op::maximum(), &[self.id, rhs.id])
    }

    /// Records the elementwise 0/1 indicator of `self >= threshold` on
    /// the same network and returns a proxy to it: the Heaviside step,
    /// ties answering one.
    ///
    /// It is the derivative mask of the `maximum` family as a recorded
    /// node — what `differentiate` emits where the engine's rules call
    /// [`Elementary::step`](crate::Elementary::step) — and it carries no
    /// gradient of its own: the function is locally constant almost
    /// everywhere, so both operands are data, not differentiable
    /// dependencies.
    ///
    /// # Panics
    /// Panics if the values belong to different tapes or their
    /// shapes differ.
    pub fn step(self, threshold: Self) -> Self {
        self.assert_same_tape(&threshold);
        self.apply(Op::step(), &[self.id, threshold.id])
    }
}

/// # Tensor operations, views, windows, and index
///
/// One recorded node per call: products and reductions, the
/// explicit broadcasts, the view movers, the sliding-window pair,
/// the gather/scatter pair, and the two fused log-domain nodes.
/// Multi-node formulas are composites and live in `composite.rs`.
impl<'tape, E: Element> Value<'tape, E> {
    /// Records the matrix product of this value and `rhs` on the same
    /// tape and returns a proxy to it.
    ///
    /// The trailing two axes contract as the plain product. Any leading
    /// axes are a batch prefix, required identical on both operands:
    /// there is no broadcast batching.
    ///
    /// # Panics
    /// Panics if the operands belong to different tapes, either operand
    /// is below rank 2, their ranks or batch axes differ, or their inner
    /// dimensions differ.
    pub fn matmul(self, rhs: Self) -> Self {
        self.assert_same_tape(&rhs);
        self.apply(Op::matmul(), &[self.id, rhs.id])
    }

    /// Records the sum of every value in this payload on the same tape
    /// and returns a proxy to it.
    pub fn sum(self) -> Self {
        self.apply(Op::sum(), &[self.id])
    }

    /// Records the sum of this value along `axis` on the same tape
    /// and returns a proxy to it.
    ///
    /// # Panics
    /// Panics if `axis` is out of rank.
    pub fn sum_along(self, axis: usize) -> Self {
        self.apply(Op::sum_along(axis), &[self.id])
    }

    /// Records the explicit broadcast of this single-value payload
    /// across `shape` on the same tape and returns a proxy to it.
    ///
    /// This is the narrowest expansion opcode: the operand must hold
    /// exactly one element, and the target shape is a recorded
    /// parameter, never an alignment rule. To read the shape off
    /// another value, use [`broadcast_like`](Self::broadcast_like); for
    /// a source of any broadcastable shape, use the composite
    /// [`broadcast_to`](Self::broadcast_to), which applies the
    /// right-aligned NumPy rule over this opcode and
    /// [`broadcast_along`](Self::broadcast_along).
    ///
    /// # Panics
    /// Panics if this value's shape does not contain exactly one
    /// element.
    pub fn broadcast(self, shape: impl Into<Shape>) -> Self {
        self.apply(Op::broadcast(shape.into()), &[self.id])
    }

    /// Records the explicit repetition of this value along a new axis
    /// of `extent` inserted at `axis` on the same tape and returns a
    /// proxy to it.
    ///
    /// This opcode widens exactly one named axis and never infers an
    /// alignment. To read the extent off a reference value, use
    /// [`broadcast_along_like`](Self::broadcast_along_like); to widen
    /// several axes at once, or to expand under the right-aligned
    /// NumPy rule, use the composite
    /// [`broadcast_to`](Self::broadcast_to).
    ///
    /// # Panics
    /// Panics if `axis` exceeds this value's rank or `extent` is zero.
    pub fn broadcast_along(self, axis: usize, extent: usize) -> Self {
        self.apply(Op::broadcast_along(axis, extent), &[self.id])
    }

    /// Records a reshape of this value to `shape` on the same tape and
    /// returns a proxy to it; the elements keep their logical row-major
    /// order.
    ///
    /// # Panics
    /// Panics if `shape`'s volume differs from this value's.
    pub fn reshape(self, shape: impl Into<Shape>) -> Self {
        self.apply(Op::reshape(shape.into()), &[self.id])
    }

    /// Records a permutation of this value's axes by `order` on the same
    /// network and returns a proxy to it; axis `i` of the result takes
    /// axis `order[i]` of this value.
    ///
    /// # Panics
    /// Panics if `order` is not a permutation of `0..rank`.
    pub fn permute(self, order: impl IntoIterator<Item = usize>) -> Self {
        self.apply(Op::permute(order), &[self.id])
    }

    /// Records the window of `len` elements from `start` along `axis` on
    /// the same network and returns a proxy to it; the forward is an O(1)
    /// view and the gradient scatters back into the unselected positions
    /// as zeros.
    ///
    /// # Panics
    /// Panics if `axis` is out of rank, `len` is zero (tensors cannot be
    /// empty), or `start + len` overflows or exceeds the axis extent.
    pub fn narrow(self, axis: usize, start: usize, len: usize) -> Self {
        self.apply(Op::narrow(axis, start, len), &[self.id])
    }

    /// Records this value placed at `start ..` along `axis` inside zeros
    /// whose `axis` has extent `full_extent`, on the same tape, and
    /// returns a proxy to it: the adjoint of [`Value::narrow`], with
    /// `narrow` as its own gradient rule.
    ///
    /// # Panics
    /// Panics if `axis` is out of rank or the window overflows or
    /// exceeds `full_extent`.
    pub fn pad(self, axis: usize, start: usize, full_extent: usize) -> Self {
        self.apply(Op::pad(axis, start, full_extent), &[self.id])
    }

    /// Records the sliding windows of this value along `axis` on the
    /// same network and returns a proxy to it: the axis becomes a
    /// `(count, size)` pair where window `w` starts at `w * step` and
    /// takes every `dilation`-th element. The forward is a strided view;
    /// the gradient folds every window contribution back onto its
    /// source position, so overlapping windows accumulate.
    ///
    /// # Panics
    /// Panics if `axis` is out of rank, `size`, `step`, or `dilation` is
    /// zero, or the dilated window span `dilation * (size - 1) + 1`
    /// overflows or exceeds the axis extent.
    pub fn unfold(self, axis: usize, size: usize, step: usize, dilation: usize) -> Self {
        self.apply(Op::unfold(axis, size, step, dilation), &[self.id])
    }

    /// Records the `(count, size)` window pair at `axis`, `axis + 1`
    /// folded back onto an axis of `extent` on the same tape and
    /// returns a proxy to it: [`unfold`](Value::unfold)'s adjoint, each
    /// source position summing the window elements read from it,
    /// accumulated output-centrically so the result is deterministic
    /// under any evaluation strategy.
    ///
    /// # Panics
    /// Panics if the operand has no `(count, size)` pair at `axis`, a
    /// parameter is zero, the dilated window span exceeds `extent`, or
    /// the pair is not what unfolding an `extent` axis by these
    /// parameters produces.
    pub fn fold(
        self,
        axis: usize,
        size: usize,
        step: usize,
        dilation: usize,
        extent: usize,
    ) -> Self {
        self.apply(Op::fold(axis, size, step, dilation, extent), &[self.id])
    }

    /// Records the row gather of this value (the table) by `selection`, a
    /// one-hot `[count, vocab]` whose vocabulary matches the table's first
    /// axis: `output[i]` is the table row `selection` names for position
    /// `i`. The gradient scatter-adds into the table only; the selection is
    /// data and receives no gradient.
    ///
    /// It is the embedding lookup: feed `selection` per run, so one graph
    /// serves any batch of indices.
    ///
    /// # Panics
    /// Panics if the values belong to different tapes, `selection` is not
    /// rank 2, or its vocabulary does not match this value's first axis.
    pub fn gather(self, selection: Self) -> Self {
        self.assert_same_tape(&selection);
        self.apply(Op::gather(), &[self.id, selection.id])
    }

    /// Records the rows of this value scatter-added into one row per
    /// entry of `selection`'s vocabulary by its one-hot indices on the
    /// same tape and returns a proxy to it: [`gather`](Value::gather)'s
    /// adjoint, accumulating rows selected more than once. The
    /// selection is data and receives no gradient.
    ///
    /// # Panics
    /// Panics if the values belong to different tapes, this value is
    /// rank 0, or `selection` is not rank 2 with one row per leading
    /// entry of this value.
    pub fn scatter(self, selection: Self) -> Self {
        self.assert_same_tape(&selection);
        self.apply(Op::scatter(), &[self.id, selection.id])
    }

    /// Records the log-softmax of this value along `axis` on the same
    /// network and returns a proxy to it: the logarithm of the softmax
    /// probabilities, computed stably in one fused node.
    ///
    /// Exponentiating the result recovers the probabilities themselves; the
    /// fused form exists because the stable computation shifts by the axis
    /// maximum, which no composition of recorded operations can express.
    ///
    /// # Panics
    /// Panics if `axis` is out of rank.
    pub fn log_softmax(self, axis: usize) -> Self {
        self.apply(Op::log_softmax(axis), &[self.id])
    }

    /// Records the log-sum-exp of this value along `axis` on the same
    /// network and returns a proxy to it: the softmax family's normalizer
    /// and a smooth maximum; like `sum_along`, the reduced axis is
    /// removed.
    ///
    /// It is a fused node for the same reason as
    /// [`log_softmax`](Value::log_softmax): the stable form shifts by the
    /// axis maximum, so the result is finite for every finite operand —
    /// where the former composition over `log_softmax` returned `inf`
    /// once finite logits differed by more than the representable range.
    /// The gradient is the softmax.
    ///
    /// # Panics
    /// Panics if `axis` is out of rank.
    pub fn logsumexp(self, axis: usize) -> Self {
        self.apply(Op::log_sum_exp(axis), &[self.id])
    }
}

// Manual implementations avoid the `Data: Clone`/`Data: Copy` bounds a
// derive would add: the proxy copies a borrow and an index, never `Data`.
impl<E> Clone for Value<'_, E> {
    fn clone(&self) -> Self {
        *self
    }
}

impl<E> Copy for Value<'_, E> {}

/// It prints only the node position to avoid dumping the whole network.
impl<E> fmt::Debug for Value<'_, E> {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("Value")
            .field("id", &self.id)
            .finish()
    }
}

impl<'tape, E: Element> Add for Value<'tape, E> {
    type Output = Value<'tape, E>;

    fn add(self, rhs: Self) -> Self::Output {
        self.assert_same_tape(&rhs);
        self.apply(Op::add(), &[self.id, rhs.id])
    }
}

impl<'tape, E: Element> Sub for Value<'tape, E> {
    type Output = Value<'tape, E>;

    fn sub(self, rhs: Self) -> Self::Output {
        self.assert_same_tape(&rhs);
        self.apply(Op::sub(), &[self.id, rhs.id])
    }
}

impl<'tape, E: Element> Mul for Value<'tape, E> {
    type Output = Value<'tape, E>;

    fn mul(self, rhs: Self) -> Self::Output {
        self.assert_same_tape(&rhs);
        self.apply(Op::mul(), &[self.id, rhs.id])
    }
}

impl<'tape, E: Element> Div for Value<'tape, E> {
    type Output = Value<'tape, E>;

    fn div(self, rhs: Self) -> Self::Output {
        self.assert_same_tape(&rhs);
        self.apply(Op::div(), &[self.id, rhs.id])
    }
}

impl<'tape, E: Element> Neg for Value<'tape, E> {
    type Output = Value<'tape, E>;

    fn neg(self) -> Self::Output {
        self.apply(Op::neg(), &[self.id])
    }
}

impl<'tape, E: Element> Add<Tensor<E>> for Value<'tape, E> {
    type Output = Value<'tape, E>;

    fn add(self, rhs: Tensor<E>) -> Self::Output {
        let literal = self.literal(rhs);
        self + literal
    }
}

impl<'tape, E: Element> Sub<Tensor<E>> for Value<'tape, E> {
    type Output = Value<'tape, E>;

    fn sub(self, rhs: Tensor<E>) -> Self::Output {
        let literal = self.literal(rhs);
        self - literal
    }
}

impl<'tape, E: Element> Mul<Tensor<E>> for Value<'tape, E> {
    type Output = Value<'tape, E>;

    fn mul(self, rhs: Tensor<E>) -> Self::Output {
        let literal = self.literal(rhs);
        self * literal
    }
}

impl<'tape, E: Element> Div<Tensor<E>> for Value<'tape, E> {
    type Output = Value<'tape, E>;

    fn div(self, rhs: Tensor<E>) -> Self::Output {
        let literal = self.literal(rhs);
        self / literal
    }
}

// Element literals record a rank-0 leaf, so scalar-looking expressions
// (`w * 2.0` on a rank-0 graph) keep their spelling; a ranked operand
// still panics at the recording expression, because a rank-0 literal
// never broadcasts implicitly.
impl<'tape, E: Element> Add<E> for Value<'tape, E> {
    type Output = Value<'tape, E>;

    fn add(self, rhs: E) -> Self::Output {
        self + Tensor::from(rhs)
    }
}

impl<'tape, E: Element> Sub<E> for Value<'tape, E> {
    type Output = Value<'tape, E>;

    fn sub(self, rhs: E) -> Self::Output {
        self - Tensor::from(rhs)
    }
}

impl<'tape, E: Element> Mul<E> for Value<'tape, E> {
    type Output = Value<'tape, E>;

    fn mul(self, rhs: E) -> Self::Output {
        self * Tensor::from(rhs)
    }
}

impl<'tape, E: Element> Div<E> for Value<'tape, E> {
    type Output = Value<'tape, E>;

    fn div(self, rhs: E) -> Self::Output {
        self / Tensor::from(rhs)
    }
}

#[cfg(test)]
#[path = "tests/value_tests.rs"]
mod tests;