/// Randomized Hadamard Transform (RHDH) for pseudo-random rotation.
///
/// Applies: scale * WHT(D ⊙ v), where D is a random ±1 diagonal and WHT is
/// the Walsh-Hadamard Transform. This is O(d log d) to apply and stores only
/// d bits (one sign per dimension), vs O(d²) for a full rotation matrix.
///
/// After applying to a unit-norm vector, each output coordinate is
/// approximately N(0, 1/padded_dim). Multiply by sqrt(padded_dim) to reach N(0,1).
use rand::SeedableRng;
use rand::RngCore;
use rand_chacha::ChaCha20Rng;

pub struct Rhdh {
    pub dim:        usize,
    pub padded_dim: usize,
    /// Random ±1 diagonal, length = padded_dim.
    signs:          Vec<i8>,
}

impl Rhdh {
    pub fn new(dim: usize, seed: u64) -> Self {
        let padded_dim = dim.next_power_of_two();
        let mut rng    = ChaCha20Rng::seed_from_u64(seed);
        let signs: Vec<i8> = (0..padded_dim)
            .map(|_| if rng.next_u32() & 1 == 0 { 1i8 } else { -1i8 })
            .collect();
        Self { dim, padded_dim, signs }
    }

    /// Apply RHDH to `v` (length == self.dim).
    ///
    /// Returns a vec of length **self.padded_dim** (NOT self.dim).
    ///
    /// WHY padded_dim: truncating the WHT output to dim elements would corrupt
    /// inner products (the truncated output matrix is not orthogonal). Using the
    /// full padded output preserves orthogonality exactly. The storage overhead is
    /// at most 2× (bounded by next_power_of_two), and is zero when dim is already
    /// a power of two (e.g. 1024, 4096).
    pub fn apply(&self, v: &[f32]) -> Vec<f32> {
        debug_assert_eq!(v.len(), self.dim);

        // 1. Pad input with zeros and apply sign flip.
        let mut buf = vec![0.0f32; self.padded_dim];
        for (i, &x) in v.iter().enumerate() {
            buf[i] = x * self.signs[i] as f32;
        }

        // 2. In-place Walsh-Hadamard Transform.
        fwht(&mut buf);

        // 3. Scale to preserve norm: ||WHT(v_padded)|| = sqrt(padded_dim) × ||v||.
        let scale = 1.0 / (self.padded_dim as f32).sqrt();
        buf.iter_mut().for_each(|x| *x *= scale);

        // Return ALL padded_dim elements — do NOT truncate.
        buf
    }

    /// Serialize signs to bytes (1 bit per sign, packed).
    pub fn to_bytes(&self) -> Vec<u8> {
        self.signs
            .chunks(8)
            .map(|chunk| {
                chunk.iter().enumerate().fold(0u8, |acc, (i, &s)| {
                    acc | (((s > 0) as u8) << i)
                })
            })
            .collect()
    }

    /// Deserialize from bytes (must match dim).
    pub fn from_bytes(bytes: &[u8], dim: usize) -> Self {
        let padded_dim = dim.next_power_of_two();
        let mut signs = Vec::with_capacity(padded_dim);
        for byte in bytes {
            for bit in 0..8 {
                signs.push(if (byte >> bit) & 1 == 1 { 1i8 } else { -1i8 });
            }
        }
        signs.truncate(padded_dim);
        Self { dim, padded_dim, signs }
    }
}

/// Fast Walsh-Hadamard Transform (in-place). len must be a power of 2.
fn fwht(a: &mut [f32]) {
    let n = a.len();
    debug_assert!(n.is_power_of_two());
    let mut h = 1usize;
    while h < n {
        for i in (0..n).step_by(h * 2) {
            for j in i..i + h {
                let x = a[j];
                let y = a[j + h];
                a[j]     = x + y;
                a[j + h] = x - y;
            }
        }
        h *= 2;
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use approx::assert_relative_eq;

    fn l2_norm(v: &[f32]) -> f32 {
        v.iter().map(|x| x * x).sum::<f32>().sqrt()
    }

    #[test]
    fn preserves_norm() {
        let dim = 64;
        let v: Vec<f32> = (0..dim).map(|i| (i as f32 + 1.0).recip()).collect();
        let norm_before = l2_norm(&v);
        let rhdh = Rhdh::new(dim, 42);
        let out = rhdh.apply(&v);
        let norm_after = l2_norm(&out);
        assert_relative_eq!(norm_before, norm_after, epsilon = 1e-4);
    }

    #[test]
    fn deterministic_with_same_seed() {
        let dim = 128;
        let v: Vec<f32> = (0..dim).map(|i| i as f32).collect();
        let out1 = Rhdh::new(dim, 0xDEAD).apply(&v);
        let out2 = Rhdh::new(dim, 0xDEAD).apply(&v);
        assert_eq!(out1, out2);
    }

    #[test]
    fn different_seeds_differ() {
        let dim = 64;
        let v: Vec<f32> = vec![1.0; dim];
        let out1 = Rhdh::new(dim, 1).apply(&v);
        let out2 = Rhdh::new(dim, 2).apply(&v);
        assert_ne!(out1, out2);
    }

    #[test]
    fn roundtrip_serialization() {
        let dim = 64;
        let rhdh = Rhdh::new(dim, 77);
        let bytes = rhdh.to_bytes();
        let restored = Rhdh::from_bytes(&bytes, dim);
        let v: Vec<f32> = (0..dim).map(|i| i as f32).collect();
        assert_eq!(rhdh.apply(&v), restored.apply(&v));
    }
}
