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

import Foundation
import Testing
import ZQEngine
import ZQStats

/// Elastic net vs glmnet fixtures. Both sides are iterative coordinate
/// descent (glmnet thresh 1e-15, ours 1e-14), so agreement is asserted at
/// 1e-6 relative — the documented tolerance for penalized solvers.
@Suite("Elastic net", .serialized)
struct ElasticNetTests {
    let fixtures: Fixtures
    let session: ZQSession

    init() async throws {
        self.fixtures = try Fixtures()
        self.session = try ZQSession(discoverUserCommands: false)
        _ = try await session.execute("use \(fixtures.datasetURL.path)")
        _ = try await session.execute("gen log_rev = ln(revenue)")
    }

    @Test("lasso matches glmnet and reproduces its selection")
    func lasso() async throws {
        let result = try await session.execute(
            "lasso log_rev price z1 z2 orders, lambda(0.05)"
        )
        expectClose(
            try #require(result.scalars["b_price"]),
            fixtures["lasso_b_price"], rtol: 1e-6, "b[price]"
        )
        expectClose(
            try #require(result.scalars["b_z1"]),
            fixtures["lasso_b_z1"], rtol: 1e-4, "b[z1] (tiny, near the threshold)"
        )
        expectClose(
            try #require(result.scalars["b__cons"]),
            fixtures["lasso_b_cons"], rtol: 1e-6, "intercept"
        )
        // glmnet zeroed z2 and orders — the selection must agree exactly.
        #expect(result.scalars["b_z2"] == 0)
        #expect(result.scalars["b_orders"] == 0)
        #expect(result.scalars["k_nonzero"] == 2)
    }

    @Test("elastic net (alpha 0.4) matches glmnet")
    func elasticNet() async throws {
        let result = try await session.execute(
            "elasticnet log_rev price z1 z2 orders, alpha(0.4) lambda(0.02)"
        )
        expectClose(
            try #require(result.scalars["b_price"]),
            fixtures["enet_b_price"], rtol: 1e-6, "b[price]"
        )
        expectClose(
            try #require(result.scalars["b_z1"]),
            fixtures["enet_b_z1"], rtol: 1e-5, "b[z1]"
        )
        expectClose(
            try #require(result.scalars["b__cons"]),
            fixtures["enet_b_cons"], rtol: 1e-6, "intercept"
        )
        #expect(result.scalars["b_orders"] == 0)
    }

    @Test("lambda 0 reproduces OLS")
    func lambdaZeroIsOLS() async throws {
        let ols = try await session.execute("reg log_rev price orders")
        let penalized = try await session.execute(
            "elasticnet log_rev price orders, alpha(1) lambda(0)"
        )
        expectClose(
            try #require(penalized.scalars["b_price"]),
            try #require(ols.scalars["b_price"]),
            rtol: 1e-8, "lambda 0 == OLS"
        )
    }

    @Test("lambda at lambda_max zeroes everything")
    func lambdaMaxZeroes() throws {
        let n = 100
        let x = (0..<n).map { Double($0) }
        let y = x.map { 2 + 3 * $0 }
        let lambdaMax = ZQElasticNet.lambdaMax(
            y: y, predictors: [("x", x)], alpha: 1
        )
        let result = try ZQElasticNet.fit(
            y: y, predictors: [("x", x)], alpha: 1, lambda: lambdaMax * 1.0000001
        )
        #expect(result.nonzeroCount == 0)
        expectClose(result.intercept, y.reduce(0, +) / Double(n), "null intercept")
    }

    @Test("predict works after a penalized fit; margins refuses")
    func postEstimation() async throws {
        _ = try await session.execute("lasso log_rev price z1 z2 orders, lambda(0.05)")
        _ = try await session.execute("predict lhat")
        let summary = try await session.execute("summarize lhat")
        #expect(summary.scalars["N"] == 57)
        await #expect(throws: ZQEngineError.self) {
            _ = try await session.execute("margins, dydx(price)")   // no VCE
        }
        _ = try await session.execute("drop lhat")
    }

    @Test("missing lambda reports lambda_max as a hint")
    func lambdaHint() async throws {
        do {
            _ = try await session.execute("lasso log_rev price")
            Issue.record("expected an error")
        } catch let error as ZQEngineError {
            #expect(error.message.contains("lambda_max"))
        }
    }
}