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

import Foundation

/// Panel fixed-effects (within) estimator for `xtreg, fe`.
///
/// The within transformation subtracts group means and adds back grand
/// means (Stata's convention, so `_cons` is reported), then solves by the
/// usual QR path. Degrees of freedom absorb the G group effects:
/// df = N − K − G. Cluster-robust VCE clusters on the panel variable with
/// the G/(G−1) factor and t statistics on G−1 df.
public struct ZQFEResult: Equatable, Sendable {
    public var coefficients: [ZQCoefficient]
    public var observationCount: Int
    public var groupCount: Int
    public var degreesOfFreedomResidual: Int
    public var inferenceDF: Double
    /// Within R²: fit of the demeaned regression.
    public var rSquaredWithin: Double
    public var rootMSE: Double
    public var fStatistic: Double?
    public var fPValue: Double?
    public var clustered: Bool

    public init(
        coefficients: [ZQCoefficient], observationCount: Int, groupCount: Int,
        degreesOfFreedomResidual: Int, inferenceDF: Double,
        rSquaredWithin: Double, rootMSE: Double,
        fStatistic: Double?, fPValue: Double?, clustered: Bool
    ) {
        self.coefficients = coefficients
        self.observationCount = observationCount
        self.groupCount = groupCount
        self.degreesOfFreedomResidual = degreesOfFreedomResidual
        self.inferenceDF = inferenceDF
        self.rSquaredWithin = rSquaredWithin
        self.rootMSE = rootMSE
        self.fStatistic = fStatistic
        self.fPValue = fPValue
        self.clustered = clustered
    }
}

public enum ZQFixedEffects {

    /// - Parameters:
    ///   - groups: dense panel codes (0..<G), one per observation.
    ///   - clustered: cluster the VCE on the panel variable.
    public static func fitWithin(
        y: [Double],
        predictors: [(name: String, values: [Double])],
        groups: [Int],
        clustered: Bool = false,
        confidenceLevel: Double = 0.95
    ) throws -> ZQFEResult {
        let n = y.count
        let slopes = predictors.count
        guard slopes > 0 else { throw ZQStatsError("xtreg: regressors required") }
        guard groups.count == n else {
            throw ZQStatsError("panel variable has wrong length")
        }
        let groupCount = Set(groups).count
        let k = slopes + 1   // slopes + reported constant
        let dfResidual = n - slopes - groupCount
        guard dfResidual > 0 else {
            throw ZQStatsError("insufficient observations: N=\(n), K=\(slopes), G=\(groupCount)")
        }

        // Group means and grand means.
        func withinTransform(_ values: [Double]) -> (transformed: [Double], demeaned: [Double]) {
            var sums = [Double](repeating: 0, count: groupCount)
            var counts = [Double](repeating: 0, count: groupCount)
            for i in 0..<n {
                sums[groups[i]] += values[i]
                counts[groups[i]] += 1
            }
            let grand = values.reduce(0, +) / Double(n)
            var transformed = [Double](repeating: 0, count: n)
            var demeaned = [Double](repeating: 0, count: n)
            for i in 0..<n {
                let groupMean = sums[groups[i]] / counts[groups[i]]
                demeaned[i] = values[i] - groupMean
                transformed[i] = demeaned[i] + grand
            }
            return (transformed, demeaned)
        }

        let (yTransformed, yDemeaned) = withinTransform(y)
        var design = [Double]()
        design.reserveCapacity(n * k)
        var demeanedColumns: [[Double]] = []
        for column in predictors {
            guard column.values.count == n else {
                throw ZQStatsError("regressor '\(column.name)' has wrong length")
            }
            let (transformed, demeaned) = withinTransform(column.values)
            design.append(contentsOf: transformed)
            demeanedColumns.append(demeaned)
        }
        design.append(contentsOf: [Double](repeating: 1, count: n))

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

        // Within residuals and fit.
        let fitted = LinearAlgebra.multiply(matrix: design, rows: n, cols: k, vector: beta)
        var residuals = [Double](repeating: 0, count: n)
        var rss = 0.0
        for i in 0..<n {
            residuals[i] = yTransformed[i] - fitted[i]
            rss += residuals[i] * residuals[i]
        }
        let tssWithin = yDemeaned.reduce(0) { $0 + $1 * $1 }
        let sigma2 = rss / Double(dfResidual)
        let rSquaredWithin = tssWithin > 0 ? 1 - rss / tssWithin : .nan

        // Covariance.
        let vce: [Double]
        let inferenceDF: Double
        if clustered {
            var scores: [Int: [Double]] = [:]
            for i in 0..<n {
                var u = scores[groups[i]] ?? [Double](repeating: 0, count: k)
                for j in 0..<slopes {
                    u[j] += demeanedColumns[j][i] * residuals[i]
                }
                u[slopes] += residuals[i]   // constant column
                scores[groups[i]] = u
            }
            guard groupCount > 1 else {
                throw ZQStatsError("cluster-robust VCE needs at least 2 groups")
            }
            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(groupCount) / Double(groupCount - 1)
            for index in meat.indices { meat[index] *= scale }
            vce = LinearAlgebra.sandwich(bread: xtxInverse, meat: meat, k: k)
            inferenceDF = Double(groupCount - 1)
        } else {
            vce = xtxInverse.map { $0 * sigma2 }
            inferenceDF = Double(dfResidual)
        }

        let tCritical = ZQDistributions.studentTQuantile(
            0.5 + confidenceLevel / 2, df: inferenceDF
        )
        var names = predictors.map(\.name)
        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 subVce = [Double](repeating: 0, count: slopes * slopes)
        var subBeta = [Double](repeating: 0, count: slopes)
        for j in 0..<slopes {
            subBeta[j] = beta[j]
            for i in 0..<slopes {
                subVce[j * slopes + i] = vce[j * k + i]
            }
        }
        if let solved = try? LinearAlgebra.solveSymmetric(subVce, k: slopes, rhs: subBeta) {
            let wald = zip(subBeta, solved).reduce(0) { $0 + $1.0 * $1.1 }
            fStatistic = wald / Double(slopes)
            fPValue = ZQDistributions.fTestPValue(
                wald / Double(slopes), df1: Double(slopes), df2: inferenceDF
            )
        }

        return ZQFEResult(
            coefficients: coefficients,
            observationCount: n,
            groupCount: groupCount,
            degreesOfFreedomResidual: dfResidual,
            inferenceDF: inferenceDF,
            rSquaredWithin: rSquaredWithin,
            rootMSE: sigma2.squareRoot(),
            fStatistic: fStatistic,
            fPValue: fPValue,
            clustered: clustered
        )
    }
}