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

import Foundation

/// Maximum-likelihood binary/count models via iteratively reweighted
/// least squares (Fisher scoring), CLAUDE.md §5. Each IRLS step solves a
/// weighted least-squares problem through the same LAPACK QR path as OLS,
/// so X'WX is never formed explicitly.
public enum ZQGLMFamily: String, Equatable, Sendable {
    case logit, probit, poisson
}

public struct ZQGLMResult: Equatable, Sendable {
    public var family: ZQGLMFamily
    public var coefficients: [ZQCoefficient]
    public var observationCount: Int
    public var logLikelihood: Double
    public var nullLogLikelihood: Double
    /// McFadden pseudo-R²: 1 − ll/ll₀ (Stata's reported definition).
    public var pseudoRSquared: Double
    /// LR χ² against the intercept-only model (classical VCE only; under
    /// robust/cluster VCE Stata reports a Wald χ², which we mirror).
    public var chiSquared: Double?
    public var chiSquaredDF: Int
    public var chiSquaredPValue: Double?
    public var iterations: Int
    public var clusterCount: Int?
    /// Linear predictor and fitted mean at the optimum.
    public var fittedMeans: [Double]
    /// Full covariance matrix, column-major k×k, aligned with
    /// `coefficients`.
    public var vce: [Double]

    public init(
        family: ZQGLMFamily,
        coefficients: [ZQCoefficient],
        observationCount: Int,
        logLikelihood: Double,
        nullLogLikelihood: Double,
        pseudoRSquared: Double,
        chiSquared: Double?,
        chiSquaredDF: Int,
        chiSquaredPValue: Double?,
        iterations: Int,
        clusterCount: Int?,
        fittedMeans: [Double],
        vce: [Double] = []
    ) {
        self.family = family
        self.coefficients = coefficients
        self.observationCount = observationCount
        self.logLikelihood = logLikelihood
        self.nullLogLikelihood = nullLogLikelihood
        self.pseudoRSquared = pseudoRSquared
        self.chiSquared = chiSquared
        self.chiSquaredDF = chiSquaredDF
        self.chiSquaredPValue = chiSquaredPValue
        self.iterations = iterations
        self.clusterCount = clusterCount
        self.fittedMeans = fittedMeans
        self.vce = vce
    }
}

public enum ZQGLM {

    /// Fits by Fisher scoring to |Δll| < 1e-13 (tighter than R's default
    /// so both land on the same optimum to ≥1e-10).
    ///
    /// Variance estimators: `.classical` is the inverse expected
    /// information; `.hc0`/`.hc1` (and Stata's `robust`) are the score
    /// sandwich without df correction; `.cluster` applies Stata's ML
    /// factor G/(G−1). `.hc2`/`.hc3` are not defined for ML estimators.
    public static func fit(
        y: [Double],
        predictors: [(name: String, values: [Double])],
        family: ZQGLMFamily,
        includeConstant: Bool = true,
        variance: ZQVarianceEstimator = .classical,
        confidenceLevel: Double = 0.95,
        maximumIterations: Int = 100
    ) throws -> ZQGLMResult {
        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)")
        }
        if case .hc2 = variance {
            throw ZQStatsError("hc2 is not defined for maximum-likelihood estimators")
        }
        if case .hc3 = variance {
            throw ZQStatsError("hc3 is not defined for maximum-likelihood estimators")
        }

        switch family {
        case .logit, .probit:
            for value in y where value != 0 && value != 1 {
                throw ZQStatsError("\(family.rawValue): outcome must be 0/1")
            }
        case .poisson:
            for value in y where value < 0 || value != value.rounded() {
                throw ZQStatsError("poisson: outcome must be a nonnegative count")
            }
        }

        // Design matrix, column-major.
        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)) }

        // IRLS.
        var beta = [Double](repeating: 0, count: k)
        var eta = [Double](repeating: 0, count: n)
        var mu = startingMeans(y: y, family: family)
        if family == .poisson || !includeConstant {
            eta = mu.map { link(family: family, mean: $0) }
        } else {
            eta = mu.map { link(family: family, mean: $0) }
        }
        var logLikelihood = self.logLikelihood(y: y, mu: mu, family: family)
        var iterations = 0
        var finalQR: LinearAlgebra.QR?
        var weights = [Double](repeating: 0, count: n)

        for iteration in 1...maximumIterations {
            iterations = iteration

            // Working response and weights.
            var z = [Double](repeating: 0, count: n)
            for i in 0..<n {
                let derivative = meanDerivative(family: family, eta: eta[i], mean: mu[i])
                let varMu = varianceFunction(family: family, mean: mu[i])
                weights[i] = derivative * derivative / varMu
                z[i] = eta[i] + (y[i] - mu[i]) / derivative
            }

            // Weighted least squares via QR on √W·X, √W·z.
            var xw = [Double](repeating: 0, count: n * k)
            var zw = [Double](repeating: 0, count: n)
            for i in 0..<n {
                let root = weights[i].squareRoot()
                zw[i] = z[i] * root
                for j in 0..<k {
                    xw[j * n + i] = x[j * n + i] * root
                }
            }
            let qr: LinearAlgebra.QR
            do {
                qr = try LinearAlgebra.QR(matrix: xw, rows: n, cols: k)
                beta = try qr.solve(rhs: zw)
            } catch {
                throw ZQStatsError("\(family.rawValue): weighted solve failed — \(error)")
            }
            finalQR = qr

            // Update linear predictor and means.
            eta = LinearAlgebra.multiply(matrix: x, rows: n, cols: k, vector: beta)
            mu = eta.map { inverseLink(family: family, eta: $0) }
            let newLogLikelihood = self.logLikelihood(y: y, mu: mu, family: family)

            if newLogLikelihood.isNaN || newLogLikelihood.isInfinite {
                throw ZQStatsError(
                    "\(family.rawValue) failed to converge (perfect prediction?)"
                )
            }
            let done = abs(newLogLikelihood - logLikelihood)
                < 1e-13 * (abs(newLogLikelihood) + 0.1)
            logLikelihood = newLogLikelihood
            if done { break }
            if iteration == maximumIterations {
                throw ZQStatsError(
                    "\(family.rawValue) did not converge in \(maximumIterations) iterations"
                )
            }
        }

        // Bread: (X'WX)⁻¹ with the weights re-evaluated AT the converged β.
        // The last IRLS factorization carries weights from the previous
        // iterate (off by ~√tolerance), which would contaminate standard
        // errors at ~1e-7 relative — visible against R fixtures.
        guard finalQR != nil else {
            throw ZQStatsError("\(family.rawValue): no iterations performed")
        }
        let bread: [Double]
        do {
            var xw = [Double](repeating: 0, count: n * k)
            for i in 0..<n {
                let derivative = meanDerivative(family: family, eta: eta[i], mean: mu[i])
                let varMu = varianceFunction(family: family, mean: mu[i])
                let root = (derivative * derivative / varMu).squareRoot()
                for j in 0..<k {
                    xw[j * n + i] = x[j * n + i] * root
                }
            }
            let informationQR = try LinearAlgebra.QR(matrix: xw, rows: n, cols: k)
            bread = try informationQR.crossProductInverse()
        } catch {
            throw ZQStatsError("\(family.rawValue): information matrix is singular")
        }

        // Covariance.
        var clusterCount: Int? = nil
        let vce: [Double]
        switch variance {
        case .classical:
            vce = bread

        case .hc0, .hc1:
            // Score sandwich: u_i = x_i (y−μ) μ′/V(μ). No df correction —
            // Stata's `robust` for ML.
            let meat = scoreMeat(
                x: x, y: y, mu: mu, eta: eta, n: n, k: k, family: family
            )
            vce = LinearAlgebra.sandwich(bread: bread, 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 {
                let scale = scoreScale(family: family, y: y[i], mu: mu[i], eta: eta[i])
                var u = scores[groups[i]] ?? [Double](repeating: 0, count: k)
                for j in 0..<k {
                    u[j] += x[j * n + i] * scale
                }
                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 ML cluster factor: G/(G−1).
            let scale = Double(g) / Double(g - 1)
            for index in meat.indices { meat[index] *= scale }
            vce = LinearAlgebra.sandwich(bread: bread, meat: meat, k: k)

        case .hc2, .hc3:
            fatalError("unreachable — rejected above")
        }

        // Inference: z statistics (normal), Stata ML convention.
        let zCritical = ZQDistributions.normalQuantile(0.5 + confidenceLevel / 2)
        var coefficients: [ZQCoefficient] = []
        for j in 0..<k {
            let se = vce[j * k + j].squareRoot()
            let zStat = beta[j] / se
            coefficients.append(ZQCoefficient(
                name: names[j],
                estimate: beta[j],
                standardError: se,
                tStatistic: zStat,
                pValue: 2 * (1 - ZQDistributions.normalCDF(abs(zStat))),
                confidenceLower: beta[j] - zCritical * se,
                confidenceUpper: beta[j] + zCritical * se
            ))
        }

        // Null (intercept-only) log likelihood — closed form: constant mean.
        let meanY = y.reduce(0, +) / Double(n)
        let nullMu = [Double](repeating: meanY, count: n)
        let nullLogLikelihood = self.logLikelihood(y: y, mu: nullMu, family: family)

        let dfModel = k - (includeConstant ? 1 : 0)
        var chiSquared: Double? = nil
        var chiSquaredPValue: Double? = nil
        if dfModel > 0, includeConstant {
            switch variance {
            case .classical:
                let lr = 2 * (logLikelihood - nullLogLikelihood)
                chiSquared = lr
                chiSquaredPValue = ZQDistributions.chiSquarePValue(lr, df: Double(dfModel))
            default:
                // Wald χ² on the non-constant coefficients under the
                // robust/cluster VCE.
                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 }
                    chiSquared = wald
                    chiSquaredPValue = ZQDistributions.chiSquarePValue(
                        wald, df: Double(dfModel)
                    )
                }
            }
        }

        return ZQGLMResult(
            family: family,
            coefficients: coefficients,
            observationCount: n,
            logLikelihood: logLikelihood,
            nullLogLikelihood: nullLogLikelihood,
            pseudoRSquared: 1 - logLikelihood / nullLogLikelihood,
            chiSquared: chiSquared,
            chiSquaredDF: dfModel,
            chiSquaredPValue: chiSquaredPValue,
            iterations: iterations,
            clusterCount: clusterCount,
            fittedMeans: mu,
            vce: vce
        )
    }

    // MARK: - Family functions

    private static func startingMeans(y: [Double], family: ZQGLMFamily) -> [Double] {
        switch family {
        case .logit, .probit:
            return y.map { ($0 + 0.5) / 2 }
        case .poisson:
            return y.map { max($0, 0.1) }
        }
    }

    private static func link(family: ZQGLMFamily, mean: Double) -> Double {
        switch family {
        case .logit: return Foundation.log(mean / (1 - mean))
        case .probit: return ZQDistributions.normalQuantile(mean)
        case .poisson: return Foundation.log(mean)
        }
    }

    private static func inverseLink(family: ZQGLMFamily, eta: Double) -> Double {
        switch family {
        case .logit:
            let clamped = min(max(eta, -30), 30)
            return 1 / (1 + Foundation.exp(-clamped))
        case .probit:
            let p = ZQDistributions.normalCDF(eta)
            return min(max(p, 1e-12), 1 - 1e-12)
        case .poisson:
            return Foundation.exp(min(eta, 300))
        }
    }

    /// dμ/dη at the current point.
    private static func meanDerivative(
        family: ZQGLMFamily, eta: Double, mean: Double
    ) -> Double {
        switch family {
        case .logit:
            return max(mean * (1 - mean), 1e-12)
        case .probit:
            return max(normalDensity(eta), 1e-12)
        case .poisson:
            return max(mean, 1e-12)
        }
    }

    /// Var(Y | μ) up to the (unit) dispersion.
    private static func varianceFunction(family: ZQGLMFamily, mean: Double) -> Double {
        switch family {
        case .logit, .probit:
            return max(mean * (1 - mean), 1e-12)
        case .poisson:
            return max(mean, 1e-12)
        }
    }

    /// Per-observation score scale: (y−μ)·μ′/V(μ), so the score is that
    /// times x_i. For canonical links this collapses to y−μ.
    private static func scoreScale(
        family: ZQGLMFamily, y: Double, mu: Double, eta: Double
    ) -> Double {
        switch family {
        case .logit, .poisson:
            return y - mu
        case .probit:
            let density = normalDensity(eta)
            return (y - mu) * density / max(mu * (1 - mu), 1e-12)
        }
    }

    private static func scoreMeat(
        x: [Double], y: [Double], mu: [Double], eta: [Double],
        n: Int, k: Int, family: ZQGLMFamily
    ) -> [Double] {
        var meat = [Double](repeating: 0, count: k * k)
        for row in 0..<n {
            let scale = scoreScale(family: family, y: y[row], mu: mu[row], eta: eta[row])
            for j in 0..<k {
                let xj = x[j * n + row] * scale
                for i in j..<k {
                    meat[j * k + i] += xj * x[i * n + row] * scale
                }
            }
        }
        for j in 0..<k {
            for i in 0..<j {
                meat[j * k + i] = meat[i * k + j]
            }
        }
        return meat
    }

    private static func logLikelihood(
        y: [Double], mu: [Double], family: ZQGLMFamily
    ) -> Double {
        var total = 0.0
        switch family {
        case .logit, .probit:
            for i in 0..<y.count {
                let p = min(max(mu[i], 1e-12), 1 - 1e-12)
                total += y[i] == 1 ? Foundation.log(p) : Foundation.log(1 - p)
            }
        case .poisson:
            for i in 0..<y.count {
                let m = max(mu[i], 1e-300)
                total += y[i] * Foundation.log(m) - m - ZQDistributions.logGamma(y[i] + 1)
            }
        }
        return total
    }

    private static func normalDensity(_ z: Double) -> Double {
        Foundation.exp(-0.5 * z * z) / (2 * Double.pi).squareRoot()
    }
}