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

import Foundation
import Testing
import ZQEngine
import ZQGPU
import ZQStats

@Suite("Bayesian regression", .serialized)
struct BayesTests {

    @Test("Philox stream normals and gammas match their theoretical moments")
    func variateGenerators() {
        var stream = PhiloxStream(seed: 42)
        let n = 200_000

        var normalSum = 0.0, normalSquares = 0.0
        for _ in 0..<n {
            let z = stream.nextNormal()
            normalSum += z
            normalSquares += z * z
        }
        #expect(abs(normalSum / Double(n)) < 0.01)
        #expect(abs(normalSquares / Double(n) - 1) < 0.02)

        // Gamma(3, rate 2): mean 1.5, variance 0.75.
        var gammaSum = 0.0, gammaSquares = 0.0
        for _ in 0..<n {
            let g = stream.nextGamma(shape: 3, rate: 2)
            gammaSum += g
            gammaSquares += g * g
        }
        let gammaMean = gammaSum / Double(n)
        let gammaVariance = gammaSquares / Double(n) - gammaMean * gammaMean
        #expect(abs(gammaMean - 1.5) < 0.01, "gamma mean \(gammaMean)")
        #expect(abs(gammaVariance - 0.75) < 0.02, "gamma variance \(gammaVariance)")

        // Shape < 1 boost path: Gamma(0.5, rate 1): mean 0.5.
        var smallSum = 0.0
        for _ in 0..<n { smallSum += stream.nextGamma(shape: 0.5, rate: 1) }
        #expect(abs(smallSum / Double(n) - 0.5) < 0.01)
    }

    @Test("flat priors recover OLS: posterior mean ≈ b̂, sd ≈ SE")
    func flatPriorAgreement() throws {
        // Synthetic data with a known DGP.
        var stream = PhiloxStream(seed: 7)
        let n = 300
        let x = (0..<n).map { _ in 5 + 10 * stream.nextUniform() }
        let y = x.map { 2 + 0.5 * $0 + 0.8 * stream.nextNormal() }

        let ols = try ZQOLS.fit(y: y, predictors: [("x", x)])
        let bayes = try ZQBayesianRegression.fitGibbs(
            y: y, predictors: [("x", x)],
            mcmcSize: 40_000, burnIn: 4_000, seed: 42,
            coefficientPriorVariance: 1e8
        )

        for (posterior, frequentist) in zip(bayes.coefficients, ols.coefficients) {
            // Monte-Carlo error with 40k (autocorrelated) draws: assert
            // within 5% of a posterior SD.
            let toleranceMean = 0.05 * frequentist.standardError
            #expect(
                abs(posterior.posteriorMean - frequentist.estimate) < toleranceMean,
                "mean[\(posterior.name)]: \(posterior.posteriorMean) vs \(frequentist.estimate)"
            )
            let sdRatio = posterior.posteriorSD / frequentist.standardError
            #expect(
                sdRatio > 0.9 && sdRatio < 1.15,
                "sd[\(posterior.name)] ratio \(sdRatio)"
            )
            // The 95% credible interval brackets the OLS estimate.
            #expect(posterior.credibleLower < frequentist.estimate)
            #expect(posterior.credibleUpper > frequentist.estimate)
        }
        // σ posterior around the DGP value 0.8.
        #expect(abs(bayes.sigma.posteriorMean - 0.8) < 0.1)
    }

    @Test("a tight prior shrinks coefficients toward zero")
    func priorShrinkage() throws {
        var stream = PhiloxStream(seed: 9)
        let n = 60
        let x = (0..<n).map { _ in stream.nextUniform() * 10 }
        let y = x.map { 3 * $0 + stream.nextNormal() }

        let flat = try ZQBayesianRegression.fitGibbs(
            y: y, predictors: [("x", x)],
            mcmcSize: 5_000, burnIn: 1_000, seed: 42,
            coefficientPriorVariance: 1e6
        )
        let tight = try ZQBayesianRegression.fitGibbs(
            y: y, predictors: [("x", x)],
            mcmcSize: 5_000, burnIn: 1_000, seed: 42,
            coefficientPriorVariance: 0.01
        )
        #expect(
            abs(tight.coefficients[0].posteriorMean)
                < abs(flat.coefficients[0].posteriorMean)
        )
        #expect(abs(flat.coefficients[0].posteriorMean - 3) < 0.2)
    }

    @Test("chains are reproducible for a fixed seed")
    func determinism() throws {
        var stream = PhiloxStream(seed: 3)
        let x = (0..<50).map { _ in stream.nextUniform() }
        let y = x.map { $0 + 0.1 * stream.nextNormal() }

        let a = try ZQBayesianRegression.fitGibbs(
            y: y, predictors: [("x", x)], mcmcSize: 1_000, burnIn: 100, seed: 42
        )
        let b = try ZQBayesianRegression.fitGibbs(
            y: y, predictors: [("x", x)], mcmcSize: 1_000, burnIn: 100, seed: 42
        )
        let c = try ZQBayesianRegression.fitGibbs(
            y: y, predictors: [("x", x)], mcmcSize: 1_000, burnIn: 100, seed: 43
        )
        #expect(a == b)
        #expect(a != c)
    }

    @Test("bayes prefix through the console")
    func consoleCommand() async throws {
        let fixtures = try Fixtures()
        let session = try ZQSession(discoverUserCommands: false)
        _ = try await session.execute("use \(fixtures.datasetURL.path)")
        _ = try await session.execute("gen log_rev = ln(revenue)")

        let ols = try await session.execute("reg log_rev price")
        let bayes = try await session.execute(
            "bayes, mcmcsize(20000) burnin(2000) seed(42) normalprior(100000000): reg log_rev price"
        )
        let posterior = try #require(bayes.scalars["b_price"])
        let frequentist = try #require(ols.scalars["b_price"])
        let se = try #require(ols.scalars["se_price"])
        #expect(abs(posterior - frequentist) < 0.05 * se)
        #expect(bayes.text.contains("Bayesian linear regression"))
        #expect(bayes.scalars["b_sigma"] != nil)
    }
}