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

import Foundation
import MLX

/// GPU-batched pairs bootstrap (CLAUDE.md §5).
///
/// Resample indices come from Philox4x32-10 evaluated as vectorized MLX
/// ops with the same key schedule as the CPU reference, so GPU and CPU
/// streams are bit-identical for a given seed — the reproducibility
/// release blocker. The heavy O(reps·n·k²) work (gather + cross-products)
/// runs on the GPU in float32; the tiny O(reps·k³) normal-equation solves
/// run on the CPU in float64 (§5 hybrid).
///
/// Documented tolerance vs the CPU float64 path: coefficient draws agree
/// to ~1e-5 relative for well-conditioned designs (float32 gather sums).
public enum ZQGPUBootstrap {

    /// Whether the Metal-backed MLX device is usable in this process.
    ///
    /// SwiftPM's command-line build cannot compile Metal shaders, so under
    /// `swift build`/`swift test` the mlx metallib does not exist and any
    /// MLX op would abort. Mirror the C++ loader's search (colocated,
    /// main-bundle Cmlx resource bundle, all loaded bundles) and report
    /// GPU availability only when a metallib is actually present —
    /// xcodebuild-built apps and tests have it, CLI builds fall back to
    /// the CPU path via the planner.
    public static let isAvailable: Bool = {
        #if arch(arm64)
        var roots: [URL] = []
        if let url = Bundle.main.resourceURL { roots.append(url) }
        roots.append(Bundle.main.bundleURL)
        if let executable = Bundle.main.executableURL {
            roots.append(executable.deletingLastPathComponent())
        }
        for bundle in Bundle.allBundles {
            if let url = bundle.resourceURL { roots.append(url) }
        }
        let candidates = roots.flatMap { root in
            [
                root.appendingPathComponent("mlx.metallib"),
                root.appendingPathComponent("default.metallib"),
                root.appendingPathComponent(
                    "mlx-swift_Cmlx.bundle/Contents/Resources/default.metallib"
                ),
            ]
        }
        return candidates.contains { FileManager.default.fileExists(atPath: $0.path) }
        #else
        return false
        #endif
    }()

    /// Runs `replicates` pairs-bootstrap OLS fits and returns the
    /// coefficient draws, one row per replicate (constant last when
    /// `includeConstant`). Replicates whose resampled design is singular
    /// come back as NaN rows for the caller to drop.
    public static func pairsBootstrapOLS(
        y: [Double],
        predictors: [(name: String, values: [Double])],
        includeConstant: Bool = true,
        replicates: Int,
        seed: UInt64,
        memoryBudgetBytes: Int = 512 << 20
    ) -> [[Double]] {
        let n = y.count
        let p = predictors.count
        let k = p + (includeConstant ? 1 : 0)
        precondition(n > k, "insufficient observations")

        // Conditioning: the GPU stage forms X'X in float32, which squares
        // the condition number — raw data with an intercept loses 2–3
        // digits. Standardizing each predictor by its full-sample mean and
        // scale is an exact reparametrization (the shift/scale constants
        // are fixed across replicates), so the original coefficients are
        // recovered exactly after the float64 solve:
        //   β_j = β̃_j / s_j,   b₀ = α + ȳ − Σ c_j β_j.
        var shifts = [Double](repeating: 0, count: p)
        var scales = [Double](repeating: 1, count: p)
        var yShift = 0.0
        if includeConstant {
            yShift = y.reduce(0, +) / Double(n)
            for (j, column) in predictors.enumerated() {
                let mean = column.values.reduce(0, +) / Double(n)
                let variance = column.values.reduce(0) {
                    $0 + ($1 - mean) * ($1 - mean)
                } / Double(n)
                shifts[j] = mean
                scales[j] = variance.squareRoot() > 1e-30 ? variance.squareRoot() : 1
            }
        }

        // Standardized design matrix on the GPU, float32, row-major (n, k).
        var design = [Float]()
        design.reserveCapacity(n * k)
        for row in 0..<n {
            for (j, column) in predictors.enumerated() {
                design.append(Float((column.values[row] - shifts[j]) / scales[j]))
            }
            if includeConstant { design.append(1) }
        }
        let x = MLXArray(design, [n, k])
        let yGPU = MLXArray(y.map { Float($0 - yShift) }, [n, 1])

        // §5 chunking: keep the gathered design under the memory budget.
        let bytesPerReplicate = n * (k + 1) * 4 * 2   // gathered X & y + transposes
        let chunkSize = max(1, min(replicates, memoryBudgetBytes / bytesPerReplicate))

        var draws = [[Double]](repeating: [], count: replicates)
        var start = 0
        while start < replicates {
            let count = min(chunkSize, replicates - start)
            let (xtx, xty) = crossProducts(
                x: x, y: yGPU, n: n, k: k,
                firstReplicate: start, replicateCount: count, seed: seed
            )
            solveChunk(
                xtx: xtx, xty: xty, k: k,
                into: &draws, offset: start, count: count
            )
            start += count
        }

        // Undo the standardization.
        if includeConstant {
            for r in 0..<replicates where !draws[r].isEmpty {
                var recovered = draws[r]
                var interceptShift = yShift
                for j in 0..<p {
                    recovered[j] = draws[r][j] / scales[j]
                    interceptShift -= shifts[j] * recovered[j]
                }
                recovered[p] = draws[r][p] + interceptShift
                draws[r] = recovered
            }
        }
        return draws
    }

    /// One replicate's resample indices computed on the GPU — exposed so
    /// the cross-backend parity test can compare against the CPU
    /// reference (`ZQResampling.pairsBootstrapIndices`) bit-for-bit.
    public static func gpuIndices(
        replicate: Int, sampleSize: Int, seed: UInt64
    ) -> [Int] {
        let base = UInt64(replicate) * UInt64(sampleSize)
        let positions = MLXArray(0..<Int32(sampleSize)).asType(.uint64)
            + MLXArray(base, dtype: .uint64)
        let indices = philoxBoundedIndices(
            positions: positions, bound: sampleSize, seed: seed
        )
        eval(indices)
        return indices.asArray(Int32.self).map(Int.init)
    }

    /// Diagnostic hook: one replicate's raw GPU cross-products with the
    /// same standardization as the production path.
    public static func debugCrossProducts(
        y: [Double], predictors: [(name: String, values: [Double])],
        replicate: Int, seed: UInt64
    ) -> (xtx: [Float], xty: [Float]) {
        let n = y.count
        let p = predictors.count
        let k = p + 1
        let meanY = y.reduce(0, +) / Double(n)
        var design = [Float]()
        for row in 0..<n {
            for column in predictors {
                let mean = column.values.reduce(0, +) / Double(n)
                let variance = column.values.reduce(0) {
                    $0 + ($1 - mean) * ($1 - mean)
                } / Double(n)
                design.append(Float((column.values[row] - mean) / variance.squareRoot()))
            }
            design.append(1)
        }
        let x = MLXArray(design, [n, k])
        let yGPU = MLXArray(y.map { Float($0 - meanY) }, [n, 1])
        return crossProducts(
            x: x, y: yGPU, n: n, k: k,
            firstReplicate: replicate, replicateCount: 1, seed: seed
        )
    }

    /// Diagnostic hook: cross-products for replicate 0 computed inside a
    /// batch of `batchSize`, to expose batch-dependent kernel numerics.
    public static func debugBatchedCrossProducts(
        y: [Double], predictors: [(name: String, values: [Double])],
        batchSize: Int, seed: UInt64
    ) -> (xtx: [Float], xty: [Float]) {
        let n = y.count
        let p = predictors.count
        let k = p + 1
        let meanY = y.reduce(0, +) / Double(n)
        var design = [Float]()
        for row in 0..<n {
            for column in predictors {
                let mean = column.values.reduce(0, +) / Double(n)
                let variance = column.values.reduce(0) {
                    $0 + ($1 - mean) * ($1 - mean)
                } / Double(n)
                design.append(Float((column.values[row] - mean) / variance.squareRoot()))
            }
            design.append(1)
        }
        let x = MLXArray(design, [n, k])
        let yGPU = MLXArray(y.map { Float($0 - meanY) }, [n, 1])
        let (xtx, xty) = crossProducts(
            x: x, y: yGPU, n: n, k: k,
            firstReplicate: 0, replicateCount: batchSize, seed: seed
        )
        return (Array(xtx.prefix(k * k)), Array(xty.prefix(k)))
    }

    // MARK: - GPU stage

    /// Gathers the resampled design for a chunk and returns float32
    /// cross-products, pulled back to the CPU: X'X (count·k·k) and X'y
    /// (count·k).
    private static func crossProducts(
        x: MLXArray, y: MLXArray, n: Int, k: Int,
        firstReplicate: Int, replicateCount: Int, seed: UInt64
    ) -> (xtx: [Float], xty: [Float]) {
        // Draw positions: replicate r uses counters r·n ..< (r+1)·n,
        // matching ZQResampling.pairsBootstrapIndices.
        let base = UInt64(firstReplicate) * UInt64(n)
        let total = replicateCount * n
        let positions = MLXArray(0..<Int32(total)).asType(.uint64)
            + MLXArray(base, dtype: .uint64)

        let indices = philoxBoundedIndices(positions: positions, bound: n, seed: seed)
            .reshaped([replicateCount, n])

        // Gather rows: (count, n, k) and (count, n, 1).
        let xg = x.take(indices.flattened(), axis: 0)
            .reshaped([replicateCount, n, k])
        let yg = y.take(indices.flattened(), axis: 0)
            .reshaped([replicateCount, n, 1])

        // Batched cross-products. X'X is computed one column at a time
        // through (count,k,n)@(count,n,1) products: MLX's batched GEMM
        // mis-accumulates small k×k outputs at batch ≥ 2 (~6e-4 relative,
        // verified against float64 sums, independent of strides/fusion),
        // while the single-column path is float32-exact. k is single-digit
        // so the extra dispatches are negligible next to the gather.
        let xt = xg.transposed(0, 2, 1)               // (count, k, n)
        let xColumns = split(xg, parts: k, axis: 2)   // k × (count, n, 1)
        let xtxColumns = xColumns.map { matmul(xt, $0) }
        let xtx = concatenated(xtxColumns, axis: 2)   // (count, k, k)
        let xty = matmul(xt, yg)                      // (count, k, 1)
        eval(xtx, xty)

        return (xtx.asArray(Float.self), xty.asArray(Float.self))
    }

    /// Philox4x32-10 as vectorized MLX ops. Given absolute draw positions,
    /// returns bounded int32 indices in 0..<bound — bit-identical to
    /// `Philox4x32.integer(at:bound:)` on the CPU.
    static func philoxBoundedIndices(
        positions: MLXArray, bound: Int, seed: UInt64
    ) -> MLXArray {
        let mask32 = MLXArray(UInt64(0xFFFF_FFFF), dtype: .uint64)
        let block = positions >> 2
        let lane = (positions & MLXArray(UInt64(3), dtype: .uint64)).asType(.uint32)

        var c0 = (block & mask32)
        var c1 = (block >> 32)
        var c2 = MLXArray.zeros(positions.shape, dtype: .uint64)
        var c3 = MLXArray.zeros(positions.shape, dtype: .uint64)

        // Key schedule precomputed on the CPU — identical to the
        // reference: key_r = (seedLo + r·W0, seedHi + r·W1) mod 2³².
        let seedLo = UInt32(truncatingIfNeeded: seed)
        let seedHi = UInt32(truncatingIfNeeded: seed >> 32)
        let m0 = MLXArray(UInt64(0xD251_1F53), dtype: .uint64)
        let m1 = MLXArray(UInt64(0xCD9E_8D57), dtype: .uint64)

        for round in 0..<10 {
            let k0 = seedLo &+ UInt32(round) &* 0x9E37_79B9
            let k1 = seedHi &+ UInt32(round) &* 0xBB67_AE85
            let product0 = m0 * c0                    // c0 < 2³² so exact
            let product1 = m1 * c2
            let high0 = product0 >> 32
            let low0 = product0 & mask32
            let high1 = product1 >> 32
            let low1 = product1 & mask32
            let newC0 = (high1 ^ c1 ^ MLXArray(UInt64(k0), dtype: .uint64)) & mask32
            let newC2 = (high0 ^ c3 ^ MLXArray(UInt64(k1), dtype: .uint64)) & mask32
            c0 = newC0
            c1 = low1
            c2 = newC2
            c3 = low0
        }

        // Lane select, then bounded multiply-shift like the reference.
        let word = which(
            lane .== MLXArray(UInt32(0), dtype: .uint32), c0,
            which(
                lane .== MLXArray(UInt32(1), dtype: .uint32), c1,
                which(lane .== MLXArray(UInt32(2), dtype: .uint32), c2, c3)
            )
        )
        let bounded = (word * MLXArray(UInt64(bound), dtype: .uint64)) >> 32
        return bounded.asType(.int32)
    }

    // MARK: - CPU stage

    /// Solves the chunk's normal equations in float64 via Cholesky.
    /// Singular systems yield NaN rows.
    private static func solveChunk(
        xtx: [Float], xty: [Float], k: Int,
        into draws: inout [[Double]], offset: Int, count: Int
    ) {
        for r in 0..<count {
            var a = [Double](repeating: 0, count: k * k)
            var b = [Double](repeating: 0, count: k)
            // MLX buffers are row-major; LAPACK wants column-major. X'X is
            // symmetric so the transpose is free; X'y is a vector.
            for i in 0..<k {
                b[i] = Double(xty[r * k + i])
                for j in 0..<k {
                    a[j * k + i] = Double(xtx[r * k * k + i * k + j])
                }
            }
            if let solution = try? CholeskySolver.solve(a, k: k, rhs: b) {
                draws[offset + r] = solution
            } else {
                draws[offset + r] = [Double](repeating: .nan, count: k)
            }
        }
    }
}