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

import Accelerate
import Foundation

/// Two-stage least squares for `ivregress 2sls`.
///
/// The projection uses the thin Q of the instrument matrix (X̂ = Q·QᵀX),
/// never forming Z'Z. Inference follows Stata's `small` convention:
/// σ² = u'u/(N−K) with t statistics; residuals come from the ORIGINAL
/// regressors (u = y − Xβ), not the projected ones.
public enum ZQIV {

    public static func fit2SLS(
        y: [Double],
        endogenous: [(name: String, values: [Double])],
        exogenous: [(name: String, values: [Double])],
        instruments: [(name: String, values: [Double])],
        includeConstant: Bool = true,
        variance: ZQVarianceEstimator = .classical,
        confidenceLevel: Double = 0.95
    ) throws -> ZQOLSResult {
        let n = y.count
        guard !endogenous.isEmpty else {
            throw ZQStatsError("ivregress: at least one endogenous regressor required")
        }
        let k = endogenous.count + exogenous.count + (includeConstant ? 1 : 0)
        let m = instruments.count + exogenous.count + (includeConstant ? 1 : 0)
        guard m >= k else {
            throw ZQStatsError(
                "ivregress: order condition fails — \(instruments.count) instruments for \(endogenous.count) endogenous regressors"
            )
        }
        guard n > k else {
            throw ZQStatsError("insufficient observations: n=\(n), k=\(k)")
        }

        func columnMajor(_ columns: [[Double]]) -> [Double] {
            var flat = [Double]()
            flat.reserveCapacity(n * columns.count)
            for column in columns { flat.append(contentsOf: column) }
            return flat
        }
        let ones = [Double](repeating: 1, count: n)

        // X = [endog, exog, 1], Z = [instruments, exog, 1].
        var xColumns = endogenous.map(\.values) + exogenous.map(\.values)
        var zColumns = instruments.map(\.values) + exogenous.map(\.values)
        if includeConstant {
            xColumns.append(ones)
            zColumns.append(ones)
        }
        for column in xColumns + zColumns where column.count != n {
            throw ZQStatsError("ivregress: variable has wrong length")
        }
        let x = columnMajor(xColumns)
        let z = columnMajor(zColumns)

        // First stage: X̂ = Q·(QᵀX) with Q the thin Q of Z.
        let projected: [Double]
        do {
            let zQR = try LinearAlgebra.QR(matrix: z, rows: n, cols: m)
            let q = try zQR.thinQ()
            var qtx = [Double](repeating: 0, count: m * k)
            cblas_dgemm(
                CblasColMajor, CblasTrans, CblasNoTrans,
                Int32(m), Int32(k), Int32(n),
                1.0, q, Int32(n), x, Int32(n), 0.0, &qtx, Int32(m)
            )
            var xhat = [Double](repeating: 0, count: n * k)
            cblas_dgemm(
                CblasColMajor, CblasNoTrans, CblasNoTrans,
                Int32(n), Int32(k), Int32(m),
                1.0, q, Int32(n), qtx, Int32(m), 0.0, &xhat, Int32(n)
            )
            projected = xhat
        } catch let error as LinearAlgebra.Failure {
            throw ZQStatsError("ivregress: instrument matrix is rank deficient — \(error)")
        }

        // Second stage.
        let secondStage: LinearAlgebra.QR
        let beta: [Double]
        let bread: [Double]
        do {
            secondStage = try LinearAlgebra.QR(matrix: projected, rows: n, cols: k)
            beta = try secondStage.solve(rhs: y)
            bread = try secondStage.crossProductInverse()
        } catch {
            throw ZQStatsError(
                "ivregress: projected design is rank deficient (weak or collinear instruments?)"
            )
        }

        // Residuals from the ORIGINAL regressors.
        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 dfResidual = n - k
        let sigma2 = rss / Double(dfResidual)
        let meanY = y.reduce(0, +) / Double(n)
        let tss = y.reduce(0) { $0 + ($1 - meanY) * ($1 - meanY) }
        let rSquared = tss > 0 ? 1 - rss / tss : .nan   // can be negative for IV
        let dfModel = k - (includeConstant ? 1 : 0)

        // Covariance: sandwich pieces built on the PROJECTED regressors.
        var clusterCount: Int?
        let vce: [Double]
        switch variance {
        case .classical:
            vce = bread.map { $0 * sigma2 }
        case .hc0, .hc1, .hc2, .hc3:
            var meat = [Double](repeating: 0, count: k * k)
            for row in 0..<n {
                let weight = residuals[row] * residuals[row]
                for j in 0..<k {
                    for i in j..<k {
                        let value = weight * projected[i * n + row] * projected[j * n + row]
                        meat[j * k + i] += value
                    }
                }
            }
            for j in 0..<k {
                for i in 0..<j { meat[j * k + i] = meat[i * k + j] }
            }
            if variance == .hc1 {
                let scale = Double(n) / Double(dfResidual)
                for index in meat.indices { meat[index] *= scale }
            }
            vce = LinearAlgebra.sandwich(bread: bread, meat: meat, k: k)
        case .cluster(let clusters):
            guard clusters.count == n else {
                throw ZQStatsError("cluster variable has wrong length")
            }
            var scores: [Int: [Double]] = [:]
            for i in 0..<n {
                var u = scores[clusters[i]] ?? [Double](repeating: 0, count: k)
                for j in 0..<k {
                    u[j] += projected[j * n + i] * residuals[i]
                }
                scores[clusters[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]
                    }
                }
            }
            let scale = Double(g) / Double(g - 1) * Double(n - 1) / Double(dfResidual)
            for index in meat.indices { meat[index] *= scale }
            vce = LinearAlgebra.sandwich(bread: bread, meat: meat, k: k)
        }

        let inferenceDF = clusterCount.map { Double($0 - 1) } ?? Double(dfResidual)
        let tCritical = ZQDistributions.studentTQuantile(
            0.5 + confidenceLevel / 2, df: inferenceDF
        )

        var names = endogenous.map(\.name) + exogenous.map(\.name)
        if includeConstant { names.append("_cons") }
        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
            ))
        }

        // Wald F on the slopes.
        var fStatistic: Double?
        var fPValue: Double?
        var fDF: (Double, Double)?
        if dfModel > 0 {
            var subVce = [Double](repeating: 0, count: dfModel * dfModel)
            var subBeta = [Double](repeating: 0, count: dfModel)
            for j in 0..<dfModel {
                subBeta[j] = beta[j]
                for i in 0..<dfModel {
                    subVce[j * dfModel + i] = 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 }
                fStatistic = wald / Double(dfModel)
                fDF = (Double(dfModel), inferenceDF)
                fPValue = ZQDistributions.fTestPValue(
                    wald / Double(dfModel), df1: Double(dfModel), df2: inferenceDF
                )
            }
        }

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