topos 0.13.1

use std::cmp::Ordering;
use std::fmt::{self, Debug, Display};
use std::ops::{Add, Div, Mul, Neg, Sub};

use crate::backend;

use super::gemm::{self, GemmTask};
use super::{Differentiable, Element, Elementary};

/// A brain-float 16 payload: the top half of an `f32`, one sign bit,
/// eight exponent bits, and seven stored mantissa bits.
///
/// Every operation converts to `f32`, computes there, and rounds the
/// result back to the nearest bf16 (ties to even) — the standard bf16
/// semantic, deterministic on every platform. Integers are exact up
/// to 256; above that the significand steps by two, then four, and so
/// on, which bounds the [`Differentiable::from_count`] contract.
///
/// Matrix multiplication follows this per-op semantic on the composed
/// path; the accumulation contract is pinned separately at the
/// [`Elementary::gemm`] seam.
#[derive(Clone, Copy)]
pub struct Bf16(u16);

impl Bf16 {
    /// The additive identity.
    pub const ZERO: Self = Self(0x0000);

    /// The multiplicative identity.
    pub const ONE: Self = Self(0x3F80);

    /// Returns the nearest bf16 to `value`, rounding ties to even.
    ///
    /// The carry from the rounding bias propagates into the exponent,
    /// so values beyond the finite range round to infinity, exactly
    /// as IEEE round-to-nearest requires. A NaN stays a NaN: the top
    /// half is kept and a mantissa bit is forced on so the payload
    /// cannot truncate to an infinity.
    pub fn from_f32(value: f32) -> Self {
        let bits = value.to_bits();
        if value.is_nan() {
            return Self(((bits >> 16) | 0x0040) as u16);
        }
        let rounding_bias = 0x7FFF + ((bits >> 16) & 1);
        Self(((bits + rounding_bias) >> 16) as u16)
    }

    /// Returns the exact `f32` this bf16 denotes: the bits shifted
    /// into the top half of a single. Every bf16 value is exactly
    /// representable, so this conversion never rounds.
    pub fn to_f32(self) -> f32 {
        f32::from_bits((self.0 as u32) << 16)
    }

    /// Returns the raw bit pattern.
    pub fn to_bits(self) -> u16 {
        self.0
    }

    /// Returns the bf16 with the raw bit pattern `bits`.
    pub fn from_bits(bits: u16) -> Self {
        Self(bits)
    }
}

impl From<f32> for Bf16 {
    fn from(value: f32) -> Self {
        Self::from_f32(value)
    }
}

impl From<Bf16> for f32 {
    fn from(value: Bf16) -> f32 {
        value.to_f32()
    }
}

/// Value equality with float semantics: positive and negative zero
/// are equal and a NaN equals nothing, matching `f32` rather than
/// the bit pattern.
impl PartialEq for Bf16 {
    fn eq(&self, other: &Self) -> bool {
        self.to_f32() == other.to_f32()
    }
}

/// Value ordering with float semantics through the exact `f32`
/// expansion, matching [`PartialEq`]: a NaN orders against nothing.
impl PartialOrd for Bf16 {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        self.to_f32().partial_cmp(&other.to_f32())
    }
}

impl Default for Bf16 {
    /// The additive identity, matching the IEEE floats' default.
    fn default() -> Self {
        Self::ZERO
    }
}

impl From<f64> for Bf16 {
    /// Rounds to the nearest bf16 through `f32`. The double rounding
    /// is exact: `f32` carries more than twice bf16's significand
    /// bits plus two, so rounding through it agrees with rounding
    /// the double directly to the nearest bf16.
    fn from(value: f64) -> Self {
        Self::from_f32(value as f32)
    }
}

impl From<Bf16> for f64 {
    /// Widens exactly, through the exact `f32` expansion.
    fn from(value: Bf16) -> f64 {
        f64::from(value.to_f32())
    }
}

impl Debug for Bf16 {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        Debug::fmt(&self.to_f32(), formatter)
    }
}

impl Display for Bf16 {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        Display::fmt(&self.to_f32(), formatter)
    }
}

impl Add for Bf16 {
    type Output = Self;

    fn add(self, rhs: Self) -> Self {
        Self::from_f32(self.to_f32() + rhs.to_f32())
    }
}

impl Sub for Bf16 {
    type Output = Self;

    fn sub(self, rhs: Self) -> Self {
        Self::from_f32(self.to_f32() - rhs.to_f32())
    }
}

impl Mul for Bf16 {
    type Output = Self;

    fn mul(self, rhs: Self) -> Self {
        Self::from_f32(self.to_f32() * rhs.to_f32())
    }
}

impl Div for Bf16 {
    type Output = Self;

    fn div(self, rhs: Self) -> Self {
        Self::from_f32(self.to_f32() / rhs.to_f32())
    }
}

impl Neg for Bf16 {
    type Output = Self;

    /// Negation flips the sign bit directly: exact for every value,
    /// including zeros, infinities, and NaN, with no round trip.
    fn neg(self) -> Self {
        Self(self.0 ^ 0x8000)
    }
}

impl Differentiable for Bf16 {
    /// Sums of bf16 terms accumulate in `f32` and round once — the
    /// convention bf16 hardware and every mixed-precision recipe
    /// follow. Promotion is exact (bf16 is a prefix of the single
    /// format); only the one final demotion rounds.
    type Accumulator = f32;

    fn promote(&self) -> f32 {
        self.to_f32()
    }

    fn demote(accumulated: f32) -> Self {
        Self::from_f32(accumulated)
    }

    fn zero() -> Self {
        Self::ZERO
    }

    fn one() -> Self {
        Self::ONE
    }

    /// Counts are exact up to 256, the last integer bf16 represents
    /// exactly; larger counts round to nearest even.
    fn from_count(count: usize) -> Self {
        Self::from_f32(count as f32)
    }

    fn is_count(&self, count: usize) -> bool {
        *self == Self::from_f32(count as f32)
    }
}

impl Elementary for Bf16 {
    fn exp(&self) -> Self {
        Self::from_f32(Elementary::exp(&self.to_f32()))
    }

    fn ln(&self) -> Self {
        Self::from_f32(Elementary::ln(&self.to_f32()))
    }

    fn sqrt(&self) -> Self {
        Self::from_f32(self.to_f32().sqrt())
    }

    fn tanh(&self) -> Self {
        Self::from_f32(Elementary::tanh(&self.to_f32()))
    }

    fn sin(&self) -> Self {
        Self::from_f32(Elementary::sin(&self.to_f32()))
    }

    fn cos(&self) -> Self {
        Self::from_f32(Elementary::cos(&self.to_f32()))
    }

    fn log1p(&self) -> Self {
        Self::from_f32(Elementary::log1p(&self.to_f32()))
    }

    fn expm1(&self) -> Self {
        Self::from_f32(Elementary::expm1(&self.to_f32()))
    }

    fn erf(&self) -> Self {
        Self::from_f32(Elementary::erf(&self.to_f32()))
    }

    fn erf_derivative(&self) -> Self {
        Self::from_f32(Elementary::erf_derivative(&self.to_f32()))
    }

    fn powf(&self, exponent: Self) -> Self {
        Self::from_f32(Elementary::powf(&self.to_f32(), exponent.to_f32()))
    }

    /// The result is one of the operands, both exactly representable,
    /// so the round trip through `f32` never rounds.
    fn maximum(&self, other: &Self) -> Self {
        Self::from_f32(self.to_f32().max(other.to_f32()))
    }

    fn step(&self, threshold: &Self) -> Self {
        if self.to_f32() >= threshold.to_f32() {
            Self::ONE
        } else {
            Self::ZERO
        }
    }

    /// Computes the product with `f32` accumulation, rounded to bf16
    /// once per output element — the convention bf16 hardware and
    /// every mixed-precision recipe follow, and the documented matmul
    /// contract of this payload. Per-op bf16 accumulation was
    /// rejected: with eight significand bits, terms stop landing once
    /// the total reaches 256 times their size.
    ///
    /// The operands expand to `f32` exactly (bf16 is a prefix of the
    /// single format), the task is offered to the accelerated `f32`
    /// backend chain, and the composed `f32` kernel answers when no
    /// backend does, so the default build stays deterministic. The
    /// emitted StableHLO states the same semantic: a `dot_general`
    /// with an `f32` result type and an explicit `convert` back.
    fn gemm(task: &GemmTask<'_, Self>) -> Option<Vec<Self>> {
        let a: Vec<f32> = task.a().iter().map(|element| element.to_f32()).collect();
        let b: Vec<f32> = task.b().iter().map(|element| element.to_f32()).collect();
        let expanded = GemmTask::new(
            &a,
            task.a_strides(),
            &b,
            task.b_strides(),
            task.m(),
            task.k(),
            task.n(),
        );
        let product = backend::offered(&expanded).unwrap_or_else(|| gemm::multiply(&expanded));
        Some(product.into_iter().map(Self::from_f32).collect())
    }
}

impl Element for Bf16 {}

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