//
//  Philox.swift
//  Metrika
//
//  Author:  Simon-Pierre Boucher
//  Contact: contact@spboucher.ai
//  Copyright © 2026 Simon-Pierre Boucher. All rights reserved.
//

/// Philox4x32-10 counter-based random number generator (Salmon et al.,
/// SC'11). Every 128-bit counter maps to four independent 32-bit outputs,
/// so draws are addressable by index — the property that makes CPU and
/// GPU streams identical for the same seed (CLAUDE.md §5: reproducibility
/// is a release blocker). This pure-Swift implementation is the reference
/// the MLX kernel path must match bit-for-bit.
public struct Philox4x32: Sendable {
    private static let multiplier0: UInt32 = 0xD251_1F53
    private static let multiplier1: UInt32 = 0xCD9E_8D57
    private static let weyl0: UInt32 = 0x9E37_79B9
    private static let weyl1: UInt32 = 0xBB67_AE85
    private static let rounds = 10

    public let key: (UInt32, UInt32)

    /// `set seed` maps the 64-bit seed directly onto the Philox key.
    public init(seed: UInt64) {
        self.key = (UInt32(truncatingIfNeeded: seed),
                    UInt32(truncatingIfNeeded: seed >> 32))
    }

    /// One Philox block: 128-bit counter → four 32-bit words.
    public func block(counter: (UInt32, UInt32, UInt32, UInt32))
        -> (UInt32, UInt32, UInt32, UInt32) {
        var c = counter
        var k = key
        for _ in 0..<Self.rounds {
            let product0 = UInt64(Self.multiplier0) * UInt64(c.0)
            let product1 = UInt64(Self.multiplier1) * UInt64(c.2)
            let high0 = UInt32(truncatingIfNeeded: product0 >> 32)
            let low0 = UInt32(truncatingIfNeeded: product0)
            let high1 = UInt32(truncatingIfNeeded: product1 >> 32)
            let low1 = UInt32(truncatingIfNeeded: product1)
            c = (high1 ^ c.1 ^ k.0, low1, high0 ^ c.3 ^ k.1, low0)
            k.0 = k.0 &+ Self.weyl0
            k.1 = k.1 &+ Self.weyl1
        }
        return c
    }

    /// The i-th 32-bit word of the stream (block = i/4, lane = i%4).
    public func word(at index: UInt64) -> UInt32 {
        let blockIndex = index / 4
        let lane = Int(index % 4)
        let output = block(counter: (
            UInt32(truncatingIfNeeded: blockIndex),
            UInt32(truncatingIfNeeded: blockIndex >> 32),
            0, 0
        ))
        switch lane {
        case 0: return output.0
        case 1: return output.1
        case 2: return output.2
        default: return output.3
        }
    }

    /// Uniform Double in [0, 1) from the i-th draw, using 53 bits built
    /// from two consecutive 32-bit words.
    public func uniform(at index: UInt64) -> Double {
        let high = UInt64(word(at: index &* 2))
        let low = UInt64(word(at: index &* 2 &+ 1))
        let bits53 = ((high << 32) | low) >> 11
        return Double(bits53) * (1.0 / 9_007_199_254_740_992.0)  // 2^-53
    }

    /// Uniform integer in 0..<bound from the i-th draw (used for
    /// bootstrap resampling indices). Bounded by 64-bit multiply-shift;
    /// the tiny modulo bias (< 2⁻³² for realistic n) is acceptable and,
    /// critically, is computed identically on CPU and GPU.
    public func integer(at index: UInt64, bound: Int) -> Int {
        precondition(bound > 0)
        let word = UInt64(self.word(at: index))
        return Int((word &* UInt64(bound)) >> 32)
    }
}

/// Bootstrap resampling indices: replicate r draws its n indices from
/// dedicated counter positions, so any subset of replicates can be
/// regenerated independently — on either backend.
public enum ZQResampling {
    public static func pairsBootstrapIndices(
        replicate: Int, sampleSize: Int, generator: Philox4x32
    ) -> [Int] {
        let base = UInt64(replicate) &* UInt64(sampleSize)
        var indices = [Int](repeating: 0, count: sampleSize)
        for i in 0..<sampleSize {
            indices[i] = generator.integer(at: base &+ UInt64(i), bound: sampleSize)
        }
        return indices
    }

    /// A random permutation of 0..<n for one replicate, defined as the
    /// argsort of n 64-bit Philox keys. Unlike Fisher–Yates this is
    /// order-free — each key is counter-addressable — so a future GPU
    /// argsort produces the identical permutation. Key collisions occur
    /// with probability ~n²·2⁻⁶⁴ and would only perturb tie order.
    ///
    /// Counter space: permutations use a dedicated stream offset (bit 63)
    /// so they never collide with bootstrap draws under the same seed.
    public static func permutationIndices(
        replicate: Int, sampleSize: Int, generator: Philox4x32
    ) -> [Int] {
        let streamOffset = UInt64(1) << 63
        let base = UInt64(replicate) &* UInt64(sampleSize)
        var keys = [UInt64](repeating: 0, count: sampleSize)
        for i in 0..<sampleSize {
            // Word positions: offset applied after the ×2 expansion so it
            // survives (offset·2 would wrap to 0 mod 2⁶⁴).
            let wordIndex = streamOffset &+ (base &+ UInt64(i)) &* 2
            let high = UInt64(generator.word(at: wordIndex))
            let low = UInt64(generator.word(at: wordIndex &+ 1))
            keys[i] = (high << 32) | low
        }
        return (0..<sampleSize).sorted { keys[$0] < keys[$1] }
    }
}