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

import Foundation
import Testing
import ZQGPU

/// Numerical regression tests for the GPU cross-product stage. These
/// pin down an MLX batched-GEMM issue: small k×k outputs mis-accumulate
/// (~6e-4 relative) for batch sizes ≥ 2, which the column-wise X'X
/// computation in `crossProducts` works around. If these start failing,
/// the workaround regressed or the kernel changed.
@Suite("GPU numerics regression", .enabled(if: ZQGPUBootstrap.isAvailable))
struct GPUDiagnosticTests {
    let n = 200
    let x: [Double]
    let y: [Double]

    init() {
        let generator = Philox4x32(seed: 7)
        let count = 200
        let xs = (0..<count).map { i in 5 + 10 * generator.uniform(at: UInt64(i)) }
        x = xs
        y = (0..<count).map { i in
            2 + 0.5 * xs[i] + (generator.uniform(at: UInt64(count + i)) - 0.5)
        }
    }

    @Test("batched cross-products match float64 sums of identical data")
    func crossProductPrecision() {
        // Exact float64 sums of the float32-quantized data the GPU sees.
        let indices = ZQResampling.pairsBootstrapIndices(
            replicate: 0, sampleSize: n, generator: Philox4x32(seed: 42)
        )
        let meanX = x.reduce(0, +) / Double(n)
        let varX = x.reduce(0) { $0 + ($1 - meanX) * ($1 - meanX) } / Double(n)
        let meanY = y.reduce(0, +) / Double(n)
        let xs = indices.map { Float((x[$0] - meanX) / varX.squareRoot()) }
        let ys = indices.map { Float(y[$0] - meanY) }
        var sxx = 0.0, sxy = 0.0
        for i in 0..<n {
            sxx += Double(xs[i]) * Double(xs[i])
            sxy += Double(xs[i]) * Double(ys[i])
        }

        for batch in [1, 2, 16] {
            let sums = ZQGPUBootstrap.debugBatchedCrossProducts(
                y: y, predictors: [("x", x)], batchSize: batch, seed: 42
            )
            expectClose(Double(sums.xtx[0]), sxx, rtol: 1e-5, "sxx batch=\(batch)")
            expectClose(Double(sums.xty[0]), sxy, rtol: 1e-5, "sxy batch=\(batch)")
        }
    }

    @Test("replicate draws are independent of batch size")
    func batchSizeInvariance() {
        let reference = ZQGPUBootstrap.pairsBootstrapOLS(
            y: y, predictors: [("x", x)], replicates: 1, seed: 42
        )[0]
        for batch in [2, 10, 50] {
            let draws = ZQGPUBootstrap.pairsBootstrapOLS(
                y: y, predictors: [("x", x)], replicates: batch, seed: 42
            )
            #expect(draws[0] == reference, "batch=\(batch)")
        }
    }
}