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

import Foundation

/// Sequential random-variate stream over the Philox counter space, for
/// samplers whose draws are inherently ordered (MCMC chains). Unlike the
/// bootstrap's counter-addressable draws, a stream consumes positions one
/// by one — rejection samplers use a variable number — but the sequence
/// is fully determined by (seed, stream id), so chains are reproducible
/// and independent chains never overlap.
///
/// Stream ids occupy bit 62 of the word-counter space plus a 2⁴⁸-word
/// block per id, disjoint from bootstrap draws (low positions) and
/// permutation keys (bit 63).
public struct PhiloxStream: Sendable {
    private let generator: Philox4x32
    private let base: UInt64
    private var position: UInt64 = 0

    public init(seed: UInt64, stream: UInt64 = 0) {
        self.generator = Philox4x32(seed: seed)
        self.base = (UInt64(1) << 62) &+ stream &* (UInt64(1) << 48)
    }

    /// Uniform in (0, 1) — endpoints excluded so inverse-CDF transforms
    /// and logs stay finite.
    public mutating func nextUniform() -> Double {
        position &+= 1
        let u = generator.uniform(at: base &+ position)
        return min(max(u, 5e-324), 1 - 2.2e-16)
    }

    /// Standard normal via Box–Muller (two uniforms per pair, the spare
    /// is cached).
    private var cachedNormal: Double?
    public mutating func nextNormal() -> Double {
        if let cached = cachedNormal {
            cachedNormal = nil
            return cached
        }
        let u1 = nextUniform()
        let u2 = nextUniform()
        let radius = (-2 * Foundation.log(u1)).squareRoot()
        let angle = 2 * Double.pi * u2
        cachedNormal = radius * Foundation.sin(angle)
        return radius * Foundation.cos(angle)
    }

    /// Gamma(shape, rate) via Marsaglia–Tsang squeeze (with the standard
    /// boost for shape < 1).
    public mutating func nextGamma(shape: Double, rate: Double) -> Double {
        precondition(shape > 0 && rate > 0)
        if shape < 1 {
            // Gamma(a) = Gamma(a+1) · U^(1/a)
            let boosted = nextGamma(shape: shape + 1, rate: rate)
            return boosted * Foundation.pow(nextUniform(), 1 / shape)
        }
        let d = shape - 1.0 / 3.0
        let c = 1 / (9 * d).squareRoot()
        while true {
            let z = nextNormal()
            let v = (1 + c * z) * (1 + c * z) * (1 + c * z)
            guard v > 0 else { continue }
            let u = nextUniform()
            if Foundation.log(u) < 0.5 * z * z + d - d * v + d * Foundation.log(v) {
                return d * v / rate
            }
        }
    }
}