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

import Foundation
import Testing
import ZQEngine

@Suite("margins", .serialized)
struct MarginsTests {
    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("OLS marginal effect is the coefficient with its SE")
    func olsMargins() async throws {
        let fit = try await session.execute("reg log_rev price")
        let margins = try await session.execute("margins, dydx(price)")
        #expect(margins.scalars["dydx_price"] == fit.scalars["b_price"])
        #expect(margins.scalars["se_price"] == fit.scalars["se_price"])
    }

    @Test("logit AME and delta-method SE match R")
    func logitMargins() async throws {
        _ = try await session.execute("logit purchase price")
        let margins = try await session.execute("margins, dydx(price)")
        expectClose(
            try #require(margins.scalars["dydx_price"]),
            fixtures["margins_logit_price"], "logit AME"
        )
        expectClose(
            try #require(margins.scalars["se_price"]),
            fixtures["margins_logit_se"], "logit delta SE"
        )
    }

    @Test("poisson AME and delta-method SE match R")
    func poissonMargins() async throws {
        _ = try await session.execute("poisson orders price")
        let margins = try await session.execute("margins, dydx(price)")
        expectClose(
            try #require(margins.scalars["dydx_price"]),
            fixtures["margins_pois_price"], "poisson AME"
        )
        expectClose(
            try #require(margins.scalars["se_price"]),
            fixtures["margins_pois_se"], "poisson delta SE"
        )
    }

    @Test("dydx over factor terms is rejected for now")
    func factorRejected() async throws {
        _ = try await session.execute("reg log_rev price i.region")
        await #expect(throws: ZQEngineError.self) {
            _ = try await session.execute("margins, dydx(region)")
        }
        // Continuous term of the same model still works.
        let margins = try await session.execute("margins, dydx(price)")
        #expect(margins.scalars["dydx_price"] != nil)
    }

    @Test("margins requires estimation results and a dydx() option")
    func validation() async throws {
        let fresh = try ZQSession(discoverUserCommands: false)
        _ = try await fresh.execute("use \(fixtures.datasetURL.path)")
        await #expect(throws: ZQEngineError.self) {
            _ = try await fresh.execute("margins, dydx(price)")
        }
        _ = try await session.execute("reg log_rev price")
        await #expect(throws: ZQEngineError.self) {
            _ = try await session.execute("margins")
        }
        await #expect(throws: ZQEngineError.self) {
            _ = try await session.execute("margins, dydx(nonexistent)")
        }
    }
}