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

import Foundation
import Testing
import ZQGPU
import ZQParser
import ZQPlanner
import ZQStats

/// GPU bootstrap cross-checks (CLAUDE.md §9): same seed ⇒ identical
/// resample indices across backends; estimates within the documented
/// float32 tolerance.
@Suite("GPU bootstrap", .enabled(if: ZQGPUBootstrap.isAvailable))
struct GPUBootstrapTests {

    @Test("MLX Philox indices are bit-identical to the CPU reference")
    func philoxParity() {
        let seed: UInt64 = 42
        let n = 137
        let generator = Philox4x32(seed: seed)

        for replicate in [0, 1, 7, 1000] {
            let cpu = ZQResampling.pairsBootstrapIndices(
                replicate: replicate, sampleSize: n, generator: generator
            )
            let gpu = ZQGPUBootstrap.gpuIndices(replicate: replicate, sampleSize: n, seed: seed)
            #expect(cpu == gpu, "replicate \(replicate)")
        }
    }

    @Test("GPU coefficient draws match CPU within float32 tolerance")
    func drawParity() throws {
        // Small synthetic problem, well conditioned.
        let n = 200
        let generator = Philox4x32(seed: 7)
        let x = (0..<n).map { i in 5 + 10 * generator.uniform(at: UInt64(i)) }
        let y = (0..<n).map { i in
            2 + 0.5 * x[i] + (generator.uniform(at: UInt64(n + i)) - 0.5)
        }
        let reps = 50

        let gpuDraws = ZQGPUBootstrap.pairsBootstrapOLS(
            y: y, predictors: [("x", x)], replicates: reps, seed: 42
        )

        let bootstrapGenerator = Philox4x32(seed: 42)
        for replicate in 0..<reps {
            let indices = ZQResampling.pairsBootstrapIndices(
                replicate: replicate, sampleSize: n, generator: bootstrapGenerator
            )
            let fit = try ZQOLS.fit(
                y: indices.map { y[$0] },
                predictors: [("x", indices.map { x[$0] })]
            )
            for (j, coefficient) in fit.coefficients.enumerated() {
                expectClose(
                    gpuDraws[replicate][j], coefficient.estimate,
                    rtol: 1e-4, "replicate \(replicate) b[\(coefficient.name)]"
                )
            }
        }
    }

    @Test("GPU bootstrap is deterministic for a fixed seed")
    func determinism() {
        let n = 100
        let generator = Philox4x32(seed: 3)
        let x = (0..<n).map { i in generator.uniform(at: UInt64(i)) }
        let y = (0..<n).map { i in x[i] + generator.uniform(at: UInt64(n + i)) }

        let first = ZQGPUBootstrap.pairsBootstrapOLS(
            y: y, predictors: [("x", x)], replicates: 20, seed: 99
        )
        let second = ZQGPUBootstrap.pairsBootstrapOLS(
            y: y, predictors: [("x", x)], replicates: 20, seed: 99
        )
        #expect(first == second)
    }

    @Test("chunking does not change results")
    func chunkingInvariance() {
        let n = 80
        let generator = Philox4x32(seed: 5)
        let x = (0..<n).map { i in generator.uniform(at: UInt64(i)) }
        let y = (0..<n).map { i in 2 * x[i] + generator.uniform(at: UInt64(n + i)) }

        let oneChunk = ZQGPUBootstrap.pairsBootstrapOLS(
            y: y, predictors: [("x", x)], replicates: 30, seed: 11
        )
        // Budget so small every chunk holds a single replicate.
        let manyChunks = ZQGPUBootstrap.pairsBootstrapOLS(
            y: y, predictors: [("x", x)], replicates: 30, seed: 11,
            memoryBudgetBytes: 1
        )
        #expect(oneChunk == manyChunks)
    }

    @Test("planner sends large-reps bootstrap to the GPU")
    func plannerDispatch() throws {
        let planner = ZQPlanner(gpuAvailable: true)
        let parser = ZQCommandParser()

        let large = try #require(try parser.parse("bootstrap, reps(10000): reg y x"))
        #expect(planner.plan(large, rowCount: 1000).backend == .gpu)

        let small = try #require(try parser.parse("bootstrap, reps(100): reg y x"))
        #expect(planner.plan(small, rowCount: 1000).backend == .cpu)

        let noGPU = ZQPlanner(gpuAvailable: false)
        #expect(noGPU.plan(large, rowCount: 1000).backend == .cpu)
    }
}