//
//  GLMTests.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

/// GLM estimators vs R glm() fixtures (epsilon 1e-12) at 1e-10 relative
/// tolerance.
@Suite("ZQGLM vs R fixtures")
struct GLMTests {
    let fixtures: Fixtures
    let purchase: [Double]
    let orders: [Double]
    let price: [Double]
    let firmID: [Int]

    init() async throws {
        self.fixtures = try Fixtures()
        let store = try ZQDataStore()
        let frame = try await store.load(contentsOf: fixtures.datasetURL)
        let (purchaseAll, _) = try frame.requireNumeric("purchase")
        let (ordersAll, _) = try frame.requireNumeric("orders")
        let (priceAll, priceMissing) = try frame.requireNumeric("price")
        let (firmAll, _) = try frame.requireNumeric("firm_id")

        var purchase: [Double] = [], orders: [Double] = []
        var price: [Double] = [], firm: [Int] = []
        for i in 0..<frame.rowCount where !priceMissing[i] {
            purchase.append(purchaseAll[i])
            orders.append(ordersAll[i])
            price.append(priceAll[i])
            firm.append(Int(firmAll[i]))
        }
        self.purchase = purchase
        self.orders = orders
        self.price = price
        self.firmID = firm
    }

    @Test("logit matches R glm(binomial)")
    func logit() throws {
        let result = try ZQGLM.fit(
            y: purchase, predictors: [("price", price)], family: .logit
        )
        expectClose(result.coefficients[0].estimate, fixtures["logit_b_price"], "b[price]")
        expectClose(result.coefficients[1].estimate, fixtures["logit_b_cons"], "b[_cons]")
        expectClose(
            result.coefficients[0].standardError, fixtures["logit_se_price"], "se[price]"
        )
        expectClose(
            result.coefficients[1].standardError, fixtures["logit_se_cons"], "se[_cons]"
        )
        expectClose(result.logLikelihood, fixtures["logit_ll"], "log likelihood")
        expectClose(result.nullLogLikelihood, fixtures["logit_ll0"], "null log likelihood")
        expectClose(try #require(result.chiSquared), fixtures["logit_chi2"], "LR chi2")
        expectClose(
            try #require(result.chiSquaredPValue),
            fixtures["logit_chi2p"], rtol: 1e-9, "Prob > chi2"
        )
    }

    @Test("logit robust and cluster standard errors")
    func logitRobust() throws {
        let robust = try ZQGLM.fit(
            y: purchase, predictors: [("price", price)], family: .logit, variance: .hc0
        )
        expectClose(
            robust.coefficients[0].standardError,
            fixtures["logit_se_hc0_price"], "robust se[price]"
        )
        let clustered = try ZQGLM.fit(
            y: purchase, predictors: [("price", price)], family: .logit,
            variance: .cluster(firmID)
        )
        expectClose(
            clustered.coefficients[0].standardError,
            fixtures["logit_se_cluster_price"], "cluster se[price]"
        )
    }

    @Test("probit matches R glm(binomial(probit))")
    func probit() throws {
        let result = try ZQGLM.fit(
            y: purchase, predictors: [("price", price)], family: .probit
        )
        expectClose(result.coefficients[0].estimate, fixtures["probit_b_price"], "b[price]")
        expectClose(result.coefficients[1].estimate, fixtures["probit_b_cons"], "b[_cons]")
        expectClose(
            result.coefficients[0].standardError, fixtures["probit_se_price"], "se[price]"
        )
        expectClose(result.logLikelihood, fixtures["probit_ll"], "log likelihood")
    }

    @Test("poisson matches R glm(poisson)")
    func poisson() throws {
        let result = try ZQGLM.fit(
            y: orders, predictors: [("price", price)], family: .poisson
        )
        expectClose(result.coefficients[0].estimate, fixtures["pois_b_price"], "b[price]")
        expectClose(result.coefficients[1].estimate, fixtures["pois_b_cons"], "b[_cons]")
        expectClose(
            result.coefficients[0].standardError, fixtures["pois_se_price"], "se[price]"
        )
        expectClose(result.logLikelihood, fixtures["pois_ll"], "log likelihood")

        let robust = try ZQGLM.fit(
            y: orders, predictors: [("price", price)], family: .poisson, variance: .hc0
        )
        expectClose(
            robust.coefficients[0].standardError,
            fixtures["pois_se_hc0_price"], "robust se[price]"
        )
    }

    @Test("logit rejects non-binary outcomes")
    func binaryValidation() {
        #expect(throws: ZQStatsError.self) {
            _ = try ZQGLM.fit(
                y: orders, predictors: [("price", price)], family: .logit
            )
        }
    }

    @Test("correlate matches R cor()")
    func correlate() async throws {
        let store = try ZQDataStore()
        let frame = try await store.load(contentsOf: fixtures.datasetURL)
        let (revenueAll, _) = try frame.requireNumeric("revenue")
        let (priceAll, priceMissing) = try frame.requireNumeric("price")
        var revenue: [Double] = [], priceComplete: [Double] = []
        for i in 0..<frame.rowCount where !priceMissing[i] {
            revenue.append(revenueAll[i])
            priceComplete.append(priceAll[i])
        }
        let matrix = ZQCorrelate.matrix(columns: [revenue, priceComplete])
        expectClose(matrix[0][1], fixtures["corr_rev_price"], "cor(revenue, price)")
        #expect(matrix[0][0] == 1 && matrix[1][1] == 1)
    }

    @Test("chi-square CDF and normal quantile vs R")
    func chiSquare() {
        expectClose(
            ZQDistributions.chiSquareCDF(3.8, df: 1),
            fixtures["dist_pchisq_3p8_1"], rtol: 1e-12, "pchisq(3.8, 1)"
        )
        expectClose(
            ZQDistributions.chiSquareCDF(25, df: 4),
            fixtures["dist_pchisq_25_4"], rtol: 1e-12, "pchisq(25, 4)"
        )
        expectClose(
            ZQDistributions.normalQuantile(0.975),
            fixtures["dist_qnorm_0p975"], rtol: 1e-12, "qnorm(0.975)"
        )
    }
}