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

import Foundation
import ZQGPU

/// Bayesian linear regression by Gibbs sampling (CLAUDE.md §1 MCMC
/// module). Semi-conjugate model with Stata `bayes` default priors:
///
///   y | β, σ²  ~  N(Xβ, σ²I)
///   β_j        ~  N(0, τ²),        τ² = 10 000 by default
///   σ²         ~  InvGamma(a₀, b₀), a₀ = b₀ = 0.01 by default
///
/// Full conditionals are exact (Gibbs, acceptance rate 1):
///   β | σ², y ~ N(Vₙ X'y/σ², Vₙ),  Vₙ = (I/τ² + X'X/σ²)⁻¹
///   σ² | β, y ~ InvGamma(a₀ + n/2, b₀ + ‖y − Xβ‖²/2)
///
/// Draws come from the Philox stream, so a (seed, chain) pair fully
/// determines the chain.
public struct ZQBayesCoefficient: Equatable, Sendable {
    public var name: String
    public var posteriorMean: Double
    public var posteriorSD: Double
    public var credibleLower: Double   // equal-tailed 2.5%
    public var credibleUpper: Double   // 97.5%
}

public struct ZQBayesResult: Equatable, Sendable {
    public var coefficients: [ZQBayesCoefficient]
    /// Posterior summary of σ (the residual standard deviation).
    public var sigma: ZQBayesCoefficient
    public var observationCount: Int
    public var mcmcSize: Int
    public var burnIn: Int
}

public enum ZQBayesianRegression {

    public static func fitGibbs(
        y: [Double],
        predictors: [(name: String, values: [Double])],
        includeConstant: Bool = true,
        mcmcSize: Int = 10_000,
        burnIn: Int = 2_500,
        seed: UInt64,
        coefficientPriorVariance: Double = 10_000,
        sigmaPriorShape: Double = 0.01,
        sigmaPriorRate: Double = 0.01
    ) throws -> ZQBayesResult {
        let n = y.count
        var names = predictors.map(\.name)
        if includeConstant { names.append("_cons") }
        let k = names.count
        guard n > k else {
            throw ZQStatsError("bayes: insufficient observations: n=\(n), k=\(k)")
        }
        guard mcmcSize > 10 else { throw ZQStatsError("bayes: mcmcsize too small") }

        // Design (column-major) and sufficient statistics.
        var x = [Double]()
        x.reserveCapacity(n * k)
        for column in predictors {
            guard column.values.count == n else {
                throw ZQStatsError("regressor '\(column.name)' has wrong length")
            }
            x.append(contentsOf: column.values)
        }
        if includeConstant { x.append(contentsOf: [Double](repeating: 1, count: n)) }

        var xtx = [Double](repeating: 0, count: k * k)
        var xty = [Double](repeating: 0, count: k)
        for j in 0..<k {
            for i in j..<k {
                var sum = 0.0
                for row in 0..<n { sum += x[i * n + row] * x[j * n + row] }
                xtx[j * k + i] = sum
                xtx[i * k + j] = sum
            }
            var sum = 0.0
            for row in 0..<n { sum += x[j * n + row] * y[row] }
            xty[j] = sum
        }

        // Start at OLS-ish values via a ridge solve to be safe.
        var stream = PhiloxStream(seed: seed)
        var beta = [Double](repeating: 0, count: k)
        var sigma2 = max(y.reduce(0) { $0 + $1 * $1 } / Double(n), 1e-8)

        let total = burnIn + mcmcSize
        var betaDraws = [[Double]](repeating: [], count: mcmcSize)
        var sigmaDraws = [Double](repeating: 0, count: mcmcSize)
        let priorPrecision = 1 / coefficientPriorVariance

        for iteration in 0..<total {
            // β | σ²: precision Λ = I/τ² + X'X/σ², posterior N(Λ⁻¹X'y/σ², Λ⁻¹).
            var lambda = [Double](repeating: 0, count: k * k)
            for j in 0..<k {
                for i in 0..<k {
                    lambda[j * k + i] = xtx[j * k + i] / sigma2
                }
                lambda[j * k + j] += priorPrecision
            }
            let rhs = xty.map { $0 / sigma2 }

            // Cholesky of Λ (lower L, column-major): solve for the mean and
            // draw β = mean + L⁻ᵀ z (since Λ = L Lᵀ ⇒ Var = L⁻ᵀ L⁻¹).
            let chol = try choleskyLower(lambda, k: k)
            let mean = try choleskySolve(chol, k: k, rhs: rhs)
            var z = [Double](repeating: 0, count: k)
            for j in 0..<k { z[j] = stream.nextNormal() }
            let noise = try backSolveTransposed(chol, k: k, rhs: z)
            for j in 0..<k { beta[j] = mean[j] + noise[j] }

            // σ² | β: InvGamma(a₀ + n/2, b₀ + RSS/2).
            var rss = 0.0
            for row in 0..<n {
                var fitted = 0.0
                for j in 0..<k { fitted += beta[j] * x[j * n + row] }
                rss += (y[row] - fitted) * (y[row] - fitted)
            }
            let precision = stream.nextGamma(
                shape: sigmaPriorShape + Double(n) / 2,
                rate: sigmaPriorRate + rss / 2
            )
            sigma2 = 1 / precision

            if iteration >= burnIn {
                betaDraws[iteration - burnIn] = beta
                sigmaDraws[iteration - burnIn] = sigma2.squareRoot()
            }
        }

        // Posterior summaries.
        func summarize(_ name: String, _ draws: [Double]) -> ZQBayesCoefficient {
            let m = draws.reduce(0, +) / Double(draws.count)
            let variance = draws.reduce(0) { $0 + ($1 - m) * ($1 - m) }
                / Double(draws.count - 1)
            let sorted = draws.sorted()
            func quantile(_ p: Double) -> Double {
                let position = p * Double(sorted.count - 1)
                let lower = Int(position)
                let fraction = position - Double(lower)
                if lower + 1 < sorted.count {
                    return sorted[lower] * (1 - fraction) + sorted[lower + 1] * fraction
                }
                return sorted[lower]
            }
            return ZQBayesCoefficient(
                name: name,
                posteriorMean: m,
                posteriorSD: variance.squareRoot(),
                credibleLower: quantile(0.025),
                credibleUpper: quantile(0.975)
            )
        }

        var coefficients: [ZQBayesCoefficient] = []
        for j in 0..<k {
            coefficients.append(summarize(names[j], betaDraws.map { $0[j] }))
        }
        return ZQBayesResult(
            coefficients: coefficients,
            sigma: summarize("sigma", sigmaDraws),
            observationCount: n,
            mcmcSize: mcmcSize,
            burnIn: burnIn
        )
    }

    // MARK: - Small dense Cholesky helpers (k×k, column-major)

    private static func choleskyLower(_ a: [Double], k: Int) throws -> [Double] {
        var l = [Double](repeating: 0, count: k * k)
        for j in 0..<k {
            var diagonal = a[j * k + j]
            for p in 0..<j { diagonal -= l[p * k + j] * l[p * k + j] }
            guard diagonal > 0 else {
                throw ZQStatsError("bayes: posterior precision is not positive definite")
            }
            let root = diagonal.squareRoot()
            l[j * k + j] = root
            for i in (j + 1)..<k {
                var value = a[j * k + i]
                for p in 0..<j { value -= l[p * k + i] * l[p * k + j] }
                l[j * k + i] = value / root
            }
        }
        return l
    }

    /// Solves Λx = b given the lower Cholesky factor (forward then back).
    private static func choleskySolve(
        _ l: [Double], k: Int, rhs: [Double]
    ) throws -> [Double] {
        var z = [Double](repeating: 0, count: k)
        for i in 0..<k {
            var value = rhs[i]
            for p in 0..<i { value -= l[p * k + i] * z[p] }
            z[i] = value / l[i * k + i]
        }
        return try backSolveTransposed(l, k: k, rhs: z)
    }

    /// Solves Lᵀx = b (also the covariance-square-root transform for
    /// drawing from N(0, Λ⁻¹)).
    private static func backSolveTransposed(
        _ l: [Double], k: Int, rhs: [Double]
    ) throws -> [Double] {
        var solution = [Double](repeating: 0, count: k)
        for i in stride(from: k - 1, through: 0, by: -1) {
            var value = rhs[i]
            for p in (i + 1)..<k { value -= l[i * k + p] * solution[p] }
            solution[i] = value / l[i * k + i]
        }
        return solution
    }
}