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

import Foundation

/// Variance estimator selection for linear regression.
public enum ZQVarianceEstimator: Equatable, Sendable {
    /// Classical homoskedastic σ²(X'X)⁻¹.
    case classical
    /// Heteroskedasticity-consistent sandwich estimators. `robust` in
    /// Stata parlance maps to `.hc1`.
    case hc0, hc1, hc2, hc3
    /// Cluster-robust: `cluster` holds one group label per observation.
    case cluster([Int])
}

/// One row of a regression coefficient table.
public struct ZQCoefficient: Equatable, Sendable {
    public var name: String
    public var estimate: Double
    public var standardError: Double
    public var tStatistic: Double
    public var pValue: Double
    public var confidenceLower: Double
    public var confidenceUpper: Double

    public init(
        name: String, estimate: Double, standardError: Double,
        tStatistic: Double, pValue: Double,
        confidenceLower: Double, confidenceUpper: Double
    ) {
        self.name = name
        self.estimate = estimate
        self.standardError = standardError
        self.tStatistic = tStatistic
        self.pValue = pValue
        self.confidenceLower = confidenceLower
        self.confidenceUpper = confidenceUpper
    }
}

/// Full OLS fit result.
public struct ZQOLSResult: Equatable, Sendable {
    public var coefficients: [ZQCoefficient]
    public var observationCount: Int
    public var degreesOfFreedomResidual: Int
    /// Degrees of freedom used for t statistics and CIs: n−k classically,
    /// G−1 under clustering (Stata convention).
    public var inferenceDF: Double
    public var rSquared: Double
    public var adjustedRSquared: Double
    public var rootMSE: Double
    public var fStatistic: Double?
    public var fPValue: Double?
    public var fDF: (Double, Double)? {
        get { fDFStorage.map { ($0.first, $0.second) } }
        set { fDFStorage = newValue.map { FDF($0.0, $0.1) } }
    }
    public var clusterCount: Int?
    public var residuals: [Double]
    /// Full covariance matrix of the coefficients, column-major k×k,
    /// aligned with `coefficients` — needed by delta-method
    /// post-estimation (margins).
    public var vce: [Double]

    private struct FDF: Equatable, Sendable {
        let first: Double
        let second: Double
        init(_ first: Double, _ second: Double) {
            self.first = first
            self.second = second
        }
        var asTuple: (Double, Double) { (first, second) }
    }

    private var fDFStorage: FDF?

    public init(
        coefficients: [ZQCoefficient],
        observationCount: Int,
        degreesOfFreedomResidual: Int,
        inferenceDF: Double,
        rSquared: Double,
        adjustedRSquared: Double,
        rootMSE: Double,
        fStatistic: Double?,
        fPValue: Double?,
        fDF: (Double, Double)?,
        clusterCount: Int?,
        residuals: [Double],
        vce: [Double] = []
    ) {
        self.coefficients = coefficients
        self.observationCount = observationCount
        self.degreesOfFreedomResidual = degreesOfFreedomResidual
        self.inferenceDF = inferenceDF
        self.rSquared = rSquared
        self.adjustedRSquared = adjustedRSquared
        self.rootMSE = rootMSE
        self.fStatistic = fStatistic
        self.fPValue = fPValue
        self.fDFStorage = fDF.map { FDF($0.0, $0.1) }
        self.clusterCount = clusterCount
        self.residuals = residuals
        self.vce = vce
    }
}

public struct ZQStatsError: Error, Equatable, Sendable, CustomStringConvertible {
    public var message: String
    public init(_ message: String) { self.message = message }
    public var description: String { message }
}

/// Ordinary least squares on fully observed (already listwise-deleted)
/// data. Solved by LAPACK QR (dgeqrf/dormqr) — X'X is never formed for
/// the coefficient path (CLAUDE.md §5).
public enum ZQOLS {

    /// - Parameters:
    ///   - y: response, length n
    ///   - predictors: named regressor columns (length n each), intercept
    ///     excluded — it is appended automatically as `_cons` unless
    ///     `includeConstant` is false.
    public static func fit(
        y: [Double],
        predictors: [(name: String, values: [Double])],
        includeConstant: Bool = true,
        variance: ZQVarianceEstimator = .classical,
        confidenceLevel: Double = 0.95
    ) throws -> ZQOLSResult {
        let n = y.count
        var names = predictors.map(\.name)
        if includeConstant { names.append("_cons") }
        let k = names.count
        guard n > k else {
            throw ZQStatsError("insufficient observations: n=\(n), k=\(k)")
        }
        for column in predictors where column.values.count != n {
            throw ZQStatsError("regressor '\(column.name)' has wrong length")
        }

        // Design matrix, column-major.
        var x = [Double]()
        x.reserveCapacity(n * k)
        for column in predictors { x.append(contentsOf: column.values) }
        if includeConstant { x.append(contentsOf: [Double](repeating: 1, count: n)) }

        let qr: LinearAlgebra.QR
        let beta: [Double]
        let xtxInverse: [Double]
        do {
            qr = try LinearAlgebra.QR(matrix: x, rows: n, cols: k)
            beta = try qr.solve(rhs: y)
            xtxInverse = try qr.crossProductInverse()
        } catch {
            throw ZQStatsError(
                "design matrix is rank deficient or QR failed: \(error)"
            )
        }

        // Residuals and fit statistics.
        let fitted = LinearAlgebra.multiply(matrix: x, rows: n, cols: k, vector: beta)
        var residuals = [Double](repeating: 0, count: n)
        var rss = 0.0
        for i in 0..<n {
            residuals[i] = y[i] - fitted[i]
            rss += residuals[i] * residuals[i]
        }
        let meanY = y.reduce(0, +) / Double(n)
        let tss = y.reduce(0) { $0 + ($1 - meanY) * ($1 - meanY) }
        let dfResidual = n - k
        let sigma2 = rss / Double(dfResidual)
        let rSquared = tss > 0 ? 1 - rss / tss : .nan
        let dfModel = k - (includeConstant ? 1 : 0)
        let adjustedRSquared =
            1 - (1 - rSquared) * Double(n - (includeConstant ? 1 : 0)) / Double(dfResidual)

        // Covariance matrix.
        var clusterCount: Int? = nil
        let vce: [Double]
        switch variance {
        case .classical:
            vce = xtxInverse.map { $0 * sigma2 }

        case .hc0, .hc1, .hc2, .hc3:
            let leverage: [Double]?
            if variance == .hc2 || variance == .hc3 {
                let q = try qr.thinQ()
                var h = [Double](repeating: 0, count: n)
                for j in 0..<k {
                    for i in 0..<n {
                        let value = q[j * n + i]
                        h[i] += value * value
                    }
                }
                leverage = h
            } else {
                leverage = nil
            }
            var weights = [Double](repeating: 0, count: n)
            for i in 0..<n {
                let e2 = residuals[i] * residuals[i]
                switch variance {
                case .hc0: weights[i] = e2
                case .hc1: weights[i] = e2 * Double(n) / Double(dfResidual)
                case .hc2: weights[i] = e2 / (1 - leverage![i])
                case .hc3: weights[i] = e2 / ((1 - leverage![i]) * (1 - leverage![i]))
                default: break
                }
            }
            let meat = weightedCrossProduct(x: x, n: n, k: k, weights: weights)
            vce = LinearAlgebra.sandwich(bread: xtxInverse, meat: meat, k: k)

        case .cluster(let groups):
            guard groups.count == n else {
                throw ZQStatsError("cluster variable has wrong length")
            }
            var scores: [Int: [Double]] = [:]
            for i in 0..<n {
                var u = scores[groups[i]] ?? [Double](repeating: 0, count: k)
                for j in 0..<k {
                    u[j] += x[j * n + i] * residuals[i]
                }
                scores[groups[i]] = u
            }
            let g = scores.count
            guard g > 1 else {
                throw ZQStatsError("cluster variable must define at least 2 groups")
            }
            clusterCount = g
            var meat = [Double](repeating: 0, count: k * k)
            for u in scores.values {
                for j in 0..<k {
                    for i in 0..<k {
                        meat[j * k + i] += u[i] * u[j]
                    }
                }
            }
            // Stata regress small-sample factor: G/(G−1) · (n−1)/(n−k).
            let scale = Double(g) / Double(g - 1) * Double(n - 1) / Double(dfResidual)
            for index in meat.indices { meat[index] *= scale }
            vce = LinearAlgebra.sandwich(bread: xtxInverse, meat: meat, k: k)
        }

        // Inference: t statistics against n−k df, or G−1 under clustering.
        let inferenceDF: Double
        if let g = clusterCount {
            inferenceDF = Double(g - 1)
        } else {
            inferenceDF = Double(dfResidual)
        }
        let tCritical = ZQDistributions.studentTQuantile(
            0.5 + confidenceLevel / 2, df: inferenceDF
        )

        var coefficients: [ZQCoefficient] = []
        for j in 0..<k {
            let se = vce[j * k + j].squareRoot()
            let t = beta[j] / se
            coefficients.append(ZQCoefficient(
                name: names[j],
                estimate: beta[j],
                standardError: se,
                tStatistic: t,
                pValue: ZQDistributions.tTestPValue(t, df: inferenceDF),
                confidenceLower: beta[j] - tCritical * se,
                confidenceUpper: beta[j] + tCritical * se
            ))
        }

        // Overall Wald F test on all non-constant coefficients.
        var fStatistic: Double? = nil
        var fPValue: Double? = nil
        var fDF: (Double, Double)? = nil
        if dfModel > 0 {
            let restricted = Array(0..<dfModel)
            var subVce = [Double](repeating: 0, count: dfModel * dfModel)
            var subBeta = [Double](repeating: 0, count: dfModel)
            for (jj, j) in restricted.enumerated() {
                subBeta[jj] = beta[j]
                for (ii, i) in restricted.enumerated() {
                    subVce[jj * dfModel + ii] = vce[j * k + i]
                }
            }
            if let solved = try? LinearAlgebra.solveSymmetric(subVce, k: dfModel, rhs: subBeta) {
                let wald = zip(subBeta, solved).reduce(0) { $0 + $1.0 * $1.1 }
                let f = wald / Double(dfModel)
                fStatistic = f
                fDF = (Double(dfModel), inferenceDF)
                fPValue = ZQDistributions.fTestPValue(
                    f, df1: Double(dfModel), df2: inferenceDF
                )
            }
        }

        return ZQOLSResult(
            coefficients: coefficients,
            observationCount: n,
            degreesOfFreedomResidual: dfResidual,
            inferenceDF: inferenceDF,
            rSquared: rSquared,
            adjustedRSquared: adjustedRSquared,
            rootMSE: sigma2.squareRoot(),
            fStatistic: fStatistic,
            fPValue: fPValue,
            fDF: fDF,
            clusterCount: clusterCount,
            residuals: residuals,
            vce: vce
        )
    }

    /// meat = Σ_i w_i · x_i x_iᵀ for column-major X.
    private static func weightedCrossProduct(
        x: [Double], n: Int, k: Int, weights: [Double]
    ) -> [Double] {
        var meat = [Double](repeating: 0, count: k * k)
        for j in 0..<k {
            for i in j..<k {
                var sum = 0.0
                for row in 0..<n {
                    sum += weights[row] * x[i * n + row] * x[j * n + row]
                }
                meat[j * k + i] = sum
                meat[i * k + j] = sum
            }
        }
        return meat
    }
}