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

import Foundation
import Testing
import ZQData
import ZQStats

/// Numerical tests against R-generated fixtures (Tests/Fixtures/generate.R)
/// at 1e-10 relative tolerance (CLAUDE.md §0).
@Suite("ZQStats vs R fixtures")
struct StatsTests {
    let fixtures: Fixtures
    let y: [Double]                 // log(revenue), estimation sample
    let price: [Double]
    let region: [Double]
    let firmID: [Int]

    init() async throws {
        self.fixtures = try Fixtures()

        // Load the fixture CSV through the real data path (DuckDB).
        let store = try ZQDataStore()
        let frame = try await store.load(contentsOf: fixtures.datasetURL)
        let (revenue, revenueMissing) = try frame.requireNumeric("revenue")
        let (priceAll, priceMissing) = try frame.requireNumeric("price")
        let (regionAll, _) = try frame.requireNumeric("region")
        let (firmAll, _) = try frame.requireNumeric("firm_id")

        // Listwise deletion, mirroring generate.R.
        var y: [Double] = [], price: [Double] = []
        var region: [Double] = [], firm: [Int] = []
        for i in 0..<frame.rowCount where !priceMissing[i] && !revenueMissing[i] {
            y.append(Foundation.log(revenue[i]))
            price.append(priceAll[i])
            region.append(regionAll[i])
            firm.append(Int(firmAll[i]))
        }
        self.y = y
        self.price = price
        self.region = region
        self.firmID = firm
    }

    @Test("summarize matches R moments")
    func summarize() async throws {
        let store = try ZQDataStore()
        let frame = try await store.load(contentsOf: fixtures.datasetURL)
        let (values, missing) = try frame.requireNumeric("revenue")
        let summary = try #require(ZQSummarize.summary(
            name: "revenue", values: values, missing: missing, detail: true
        ))
        #expect(summary.observationCount == Int(fixtures["sum_revenue_n"]))
        expectClose(summary.mean, fixtures["sum_revenue_mean"], "mean")
        expectClose(summary.variance, fixtures["sum_revenue_var"], "variance")
        expectClose(summary.standardDeviation, fixtures["sum_revenue_sd"], "sd")
        expectClose(summary.minimum, fixtures["sum_revenue_min"], "min")
        expectClose(summary.maximum, fixtures["sum_revenue_max"], "max")
        let detail = try #require(summary.detail)
        expectClose(detail.skewness, fixtures["sum_revenue_skewness"], "skewness")
        expectClose(detail.kurtosis, fixtures["sum_revenue_kurtosis"], "kurtosis")
    }

    @Test("OLS point estimates and classical inference")
    func classicalOLS() throws {
        let result = try ZQOLS.fit(
            y: y, predictors: [("price", price)], variance: .classical
        )
        #expect(result.observationCount == Int(fixtures["ols_n"]))
        let priceCoef = result.coefficients[0]
        let constant = result.coefficients[1]
        expectClose(priceCoef.estimate, fixtures["ols_b_price"], "b[price]")
        expectClose(constant.estimate, fixtures["ols_b_cons"], "b[_cons]")
        expectClose(priceCoef.standardError, fixtures["ols_se_classical_price"], "se[price]")
        expectClose(constant.standardError, fixtures["ols_se_classical_cons"], "se[_cons]")
        expectClose(priceCoef.tStatistic, fixtures["ols_t_price"], "t[price]")
        expectClose(priceCoef.pValue, fixtures["ols_p_price"], "p[price]")
        expectClose(priceCoef.confidenceLower, fixtures["ols_ci_lower_price"], "ci lower")
        expectClose(priceCoef.confidenceUpper, fixtures["ols_ci_upper_price"], "ci upper")
        expectClose(result.rSquared, fixtures["ols_r2"], "R²")
        expectClose(result.adjustedRSquared, fixtures["ols_r2a"], "adj. R²")
        expectClose(result.rootMSE, fixtures["ols_rmse"], "root MSE")
        expectClose(try #require(result.fStatistic), fixtures["ols_F"], "F")
    }

    @Test("HC0–HC3 sandwich standard errors", arguments: [
        ("hc0", ZQVarianceEstimator.hc0),
        ("hc1", ZQVarianceEstimator.hc1),
        ("hc2", ZQVarianceEstimator.hc2),
        ("hc3", ZQVarianceEstimator.hc3),
    ])
    func robustSE(_ label: String, _ variance: ZQVarianceEstimator) throws {
        let result = try ZQOLS.fit(
            y: y, predictors: [("price", price)], variance: variance
        )
        expectClose(
            result.coefficients[0].standardError,
            fixtures["ols_se_\(label)_price"], "\(label) se[price]"
        )
        expectClose(
            result.coefficients[1].standardError,
            fixtures["ols_se_\(label)_cons"], "\(label) se[_cons]"
        )
    }

    @Test("cluster-robust standard errors (Stata small-sample factor)")
    func clusterSE() throws {
        let result = try ZQOLS.fit(
            y: y, predictors: [("price", price)], variance: .cluster(firmID)
        )
        #expect(result.clusterCount == Int(fixtures["ols_G"]))
        expectClose(
            result.coefficients[0].standardError,
            fixtures["ols_se_cluster_price"], "cluster se[price]"
        )
        expectClose(
            result.coefficients[1].standardError,
            fixtures["ols_se_cluster_cons"], "cluster se[_cons]"
        )
        expectClose(
            result.coefficients[0].pValue,
            fixtures["ols_p_cluster_price"], "cluster p[price] (t on G−1 df)"
        )
    }

    @Test("OLS invariance: permuting observations leaves estimates unchanged")
    func permutationInvariance() throws {
        let base = try ZQOLS.fit(y: y, predictors: [("price", price)])
        var order = Array(0..<y.count)
        // Deterministic shuffle (LCG) — no seeding dependency in tests.
        var state: UInt64 = 88_172_645_463_325_252
        for i in stride(from: order.count - 1, through: 1, by: -1) {
            state = state &* 6_364_136_223_846_793_005 &+ 1_442_695_040_888_963_407
            order.swapAt(i, Int(state % UInt64(i + 1)))
        }
        let permuted = try ZQOLS.fit(
            y: order.map { y[$0] },
            predictors: [("price", order.map { price[$0] })]
        )
        for (a, b) in zip(base.coefficients, permuted.coefficients) {
            expectClose(a.estimate, b.estimate, "permuted b[\(a.name)]")
            expectClose(a.standardError, b.standardError, "permuted se[\(a.name)]")
        }
    }

    @Test("OLS invariance: scaling a regressor rescales its coefficient")
    func scalingInvariance() throws {
        let base = try ZQOLS.fit(y: y, predictors: [("price", price)])
        let scaled = try ZQOLS.fit(
            y: y, predictors: [("price", price.map { $0 * 100 })]
        )
        expectClose(
            scaled.coefficients[0].estimate * 100,
            base.coefficients[0].estimate,
            "scaled coefficient"
        )
        expectClose(scaled.rSquared, base.rSquared, "scaled R²")
    }
}

@Suite("Distribution functions vs R")
struct DistributionTests {
    let fixtures = try! Fixtures()

    @Test("Student t CDF")
    func studentT() {
        expectClose(
            ZQDistributions.studentTCDF(2.5, df: 10),
            fixtures["dist_pt_2p5_df10"], rtol: 1e-12, "pt(2.5, 10)"
        )
        expectClose(
            ZQDistributions.studentTCDF(-1.3, df: 3),
            fixtures["dist_pt_m1p3_df3"], rtol: 1e-12, "pt(-1.3, 3)"
        )
        expectClose(
            ZQDistributions.studentTCDF(0.05, df: 57),
            fixtures["dist_pt_0p05_df57"], rtol: 1e-12, "pt(0.05, 57)"
        )
    }

    @Test("F CDF")
    func fisher() {
        expectClose(
            ZQDistributions.fCDF(3.7, df1: 2, df2: 30),
            fixtures["dist_pf_3p7_2_30"], rtol: 1e-12, "pf(3.7, 2, 30)"
        )
        expectClose(
            ZQDistributions.fCDF(0.5, df1: 5, df2: 100),
            fixtures["dist_pf_0p5_5_100"], rtol: 1e-12, "pf(0.5, 5, 100)"
        )
    }

    @Test("t quantile")
    func tQuantile() {
        expectClose(
            ZQDistributions.studentTQuantile(0.975, df: 12),
            fixtures["dist_qt_0p975_df12"], rtol: 1e-10, "qt(0.975, 12)"
        )
        expectClose(
            ZQDistributions.studentTQuantile(0.995, df: 4),
            fixtures["dist_qt_0p995_df4"], rtol: 1e-10, "qt(0.995, 4)"
        )
    }

    @Test("normal CDF")
    func normal() {
        expectClose(
            ZQDistributions.normalCDF(1.64),
            fixtures["dist_pnorm_1p64"], rtol: 1e-14, "pnorm(1.64)"
        )
    }
}