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

import Accelerate

/// Thin LAPACK/BLAS wrappers used by the CPU estimators. All matrices are
/// stored column-major (LAPACK convention) as flat [Double] buffers.
enum LinearAlgebra {
    struct Failure: Error, CustomStringConvertible {
        let routine: String
        let info: Int
        var description: String { "\(routine) failed with info=\(info)" }
    }

    /// QR decomposition of an n×k matrix (n ≥ k) via `dgeqrf`. CLAUDE.md §5:
    /// never form X'X — coefficient solves and (X'X)⁻¹ all go through R.
    struct QR {
        /// Factored matrix as returned by dgeqrf: R in the upper triangle,
        /// Householder vectors below the diagonal.
        private(set) var factored: [Double]
        private(set) var tau: [Double]
        let rows: Int
        let cols: Int

        init(matrix: [Double], rows: Int, cols: Int) throws {
            precondition(matrix.count == rows * cols, "matrix size mismatch")
            precondition(rows >= cols, "QR requires rows ≥ cols")
            var a = matrix
            var m = __CLPK_integer(rows)
            var n = __CLPK_integer(cols)
            var lda = m
            var tau = [Double](repeating: 0, count: cols)
            var info: __CLPK_integer = 0

            // Workspace query, then factorization.
            var lwork: __CLPK_integer = -1
            var workQuery = [Double](repeating: 0, count: 1)
            dgeqrf_(&m, &n, &a, &lda, &tau, &workQuery, &lwork, &info)
            lwork = __CLPK_integer(workQuery[0])
            var work = [Double](repeating: 0, count: max(1, Int(lwork)))
            dgeqrf_(&m, &n, &a, &lda, &tau, &work, &lwork, &info)
            guard info == 0 else { throw Failure(routine: "dgeqrf", info: Int(info)) }

            self.factored = a
            self.tau = tau
            self.rows = rows
            self.cols = cols
        }

        /// Solves min ‖Ax − b‖ using the stored factorization:
        /// x = R⁻¹ (Q'b)[0..<k].
        func solve(rhs: [Double]) throws -> [Double] {
            precondition(rhs.count == rows, "rhs size mismatch")
            var qtb = rhs
            var side: Int8 = Int8(UInt8(ascii: "L"))
            var trans: Int8 = Int8(UInt8(ascii: "T"))
            var m = __CLPK_integer(rows)
            var nrhs: __CLPK_integer = 1
            var k = __CLPK_integer(cols)
            var lda = m
            var ldc = m
            var info: __CLPK_integer = 0
            var a = factored
            var tauCopy = tau

            var lwork: __CLPK_integer = -1
            var workQuery = [Double](repeating: 0, count: 1)
            dormqr_(&side, &trans, &m, &nrhs, &k, &a, &lda, &tauCopy, &qtb, &ldc,
                    &workQuery, &lwork, &info)
            lwork = __CLPK_integer(workQuery[0])
            var work = [Double](repeating: 0, count: max(1, Int(lwork)))
            dormqr_(&side, &trans, &m, &nrhs, &k, &a, &lda, &tauCopy, &qtb, &ldc,
                    &work, &lwork, &info)
            guard info == 0 else { throw Failure(routine: "dormqr", info: Int(info)) }

            // Back-substitute R x = (Q'b)[0..<k].
            var uplo: Int8 = Int8(UInt8(ascii: "U"))
            var transN: Int8 = Int8(UInt8(ascii: "N"))
            var diag: Int8 = Int8(UInt8(ascii: "N"))
            var n = __CLPK_integer(cols)
            var solution = Array(qtb[0..<cols])
            var ldb = n
            dtrtrs_(&uplo, &transN, &diag, &n, &nrhs, &a, &lda, &solution, &ldb, &info)
            guard info == 0 else { throw Failure(routine: "dtrtrs", info: Int(info)) }
            return solution
        }

        /// (X'X)⁻¹ = R⁻¹ R⁻ᵀ, computed from the k×k triangular factor.
        func crossProductInverse() throws -> [Double] {
            // Extract the k×k upper triangle of R (column-major).
            var r = [Double](repeating: 0, count: cols * cols)
            for j in 0..<cols {
                for i in 0...j {
                    r[j * cols + i] = factored[j * rows + i]
                }
            }
            var uplo: Int8 = Int8(UInt8(ascii: "U"))
            var diag: Int8 = Int8(UInt8(ascii: "N"))
            var n = __CLPK_integer(cols)
            var lda = n
            var info: __CLPK_integer = 0
            dtrtri_(&uplo, &diag, &n, &r, &lda, &info)
            guard info == 0 else { throw Failure(routine: "dtrtri", info: Int(info)) }

            // XtXinv = Rinv · Rinvᵀ (symmetric k×k).
            var result = [Double](repeating: 0, count: cols * cols)
            for j in 0..<cols {
                for i in 0..<cols {
                    var sum = 0.0
                    // Rinv is upper triangular: Rinv[i, l] nonzero for l ≥ i.
                    for l in max(i, j)..<cols {
                        sum += r[l * cols + i] * r[l * cols + j]
                    }
                    result[j * cols + i] = sum
                }
            }
            return result
        }

        /// Explicit thin Q (n×k, column-major) via `dorgqr`, used for
        /// leverage values h_i = ‖q_i‖².
        func thinQ() throws -> [Double] {
            var a = factored
            var m = __CLPK_integer(rows)
            var n = __CLPK_integer(cols)
            var k = n
            var lda = m
            var info: __CLPK_integer = 0
            var tauCopy = tau

            var lwork: __CLPK_integer = -1
            var workQuery = [Double](repeating: 0, count: 1)
            dorgqr_(&m, &n, &k, &a, &lda, &tauCopy, &workQuery, &lwork, &info)
            lwork = __CLPK_integer(workQuery[0])
            var work = [Double](repeating: 0, count: max(1, Int(lwork)))
            dorgqr_(&m, &n, &k, &a, &lda, &tauCopy, &work, &lwork, &info)
            guard info == 0 else { throw Failure(routine: "dorgqr", info: Int(info)) }
            return a
        }
    }

    /// y ← A·x for column-major A (rows×cols).
    static func multiply(
        matrix a: [Double], rows: Int, cols: Int, vector x: [Double]
    ) -> [Double] {
        precondition(x.count == cols)
        var result = [Double](repeating: 0, count: rows)
        cblas_dgemv(
            CblasColMajor, CblasNoTrans,
            Int32(rows), Int32(cols),
            1.0, a, Int32(rows), x, 1, 0.0, &result, 1
        )
        return result
    }

    /// Symmetric sandwich product B · M · B for k×k column-major matrices.
    static func sandwich(bread: [Double], meat: [Double], k: Int) -> [Double] {
        precondition(bread.count == k * k && meat.count == k * k)
        var bm = [Double](repeating: 0, count: k * k)
        cblas_dgemm(
            CblasColMajor, CblasNoTrans, CblasNoTrans,
            Int32(k), Int32(k), Int32(k),
            1.0, bread, Int32(k), meat, Int32(k),
            0.0, &bm, Int32(k)
        )
        var result = [Double](repeating: 0, count: k * k)
        cblas_dgemm(
            CblasColMajor, CblasNoTrans, CblasNoTrans,
            Int32(k), Int32(k), Int32(k),
            1.0, bm, Int32(k), bread, Int32(k),
            0.0, &result, Int32(k)
        )
        return result
    }

    /// Solves the symmetric positive-definite system A·x = b (k×k) via
    /// Cholesky — used for Wald tests on coefficient subsets.
    static func solveSymmetric(
        _ a: [Double], k: Int, rhs: [Double]
    ) throws -> [Double] {
        var factor = a
        var solution = rhs
        var uplo: Int8 = Int8(UInt8(ascii: "U"))
        var n = __CLPK_integer(k)
        var nrhs: __CLPK_integer = 1
        var lda = n
        var ldb = n
        var info: __CLPK_integer = 0
        dposv_(&uplo, &n, &nrhs, &factor, &lda, &solution, &ldb, &info)
        guard info == 0 else { throw Failure(routine: "dposv", info: Int(info)) }
        return solution
    }
}