//
//  PanelIVTests.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 ZQParser

/// xtreg FE and ivregress 2SLS against R fixtures, end-to-end through the
/// session (which exercises the parser's IV-group grammar too).
@Suite("Panel FE and IV", .serialized)
struct PanelIVTests {
    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)")
        _ = try await session.execute("xtset firm_id")
    }

    @Test("ivregress parser grammar")
    func parserIVGroup() throws {
        let parser = ZQCommandParser()
        let command = try #require(
            try parser.parse("ivregress 2sls y x1 (p q = z1 z2 z3), robust")
        )
        #expect(command.verb == "ivregress")
        #expect(command.subverb == "2sls")
        #expect(command.varlist == [.simple("y"), .simple("x1")])
        #expect(command.ivSpec == ZQIVSpec(
            endogenous: ["p", "q"], instruments: ["z1", "z2", "z3"]
        ))
        #expect(command.hasOption("robust"))
    }

    @Test("xtreg fe matches the R within estimator")
    func fixedEffects() async throws {
        let result = try await session.execute("xtreg log_rev price, fe")
        #expect(result.scalars["N"] == fixtures["fe_N"])
        #expect(result.scalars["N_g"] == fixtures["fe_G"])
        expectClose(try #require(result.scalars["b_price"]), fixtures["fe_b_price"], "b[price]")
        expectClose(try #require(result.scalars["b__cons"]), fixtures["fe_b_cons"], "b[_cons]")
        expectClose(
            try #require(result.scalars["se_price"]),
            fixtures["fe_se_price"], "se[price]"
        )
        expectClose(try #require(result.scalars["r2_w"]), fixtures["fe_r2_within"], "within R²")
        expectClose(try #require(result.scalars["rmse"]), fixtures["fe_rmse"], "root MSE")
        #expect(result.text.contains("Fixed-effects (within) regression"))
    }

    @Test("xtreg fe with panel-clustered SE matches R")
    func fixedEffectsClustered() async throws {
        let result = try await session.execute("xtreg log_rev price, fe cluster(firm_id)")
        expectClose(
            try #require(result.scalars["se_price"]),
            fixtures["fe_se_cluster_price"], "cluster se[price]"
        )
        #expect(result.scalars["df_r"] == fixtures["fe_G"] - 1)
    }

    @Test("xtreg without xtset or fe errors clearly")
    func xtregValidation() 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("xtreg revenue price, fe")   // no xtset
        }
        await #expect(throws: ZQEngineError.self) {
            _ = try await session.execute("xtreg log_rev price")     // no fe
        }
    }

    @Test("ivregress 2sls matches the R manual 2SLS")
    func twoStageLeastSquares() async throws {
        let result = try await session.execute(
            "ivregress 2sls log_rev (price = z1 z2)"
        )
        #expect(result.scalars["N"] == fixtures["iv_N"])
        expectClose(try #require(result.scalars["b_price"]), fixtures["iv_b_price"], "b[price]")
        expectClose(try #require(result.scalars["b__cons"]), fixtures["iv_b_cons"], "b[_cons]")
        expectClose(
            try #require(result.scalars["se_price"]), fixtures["iv_se_price"], "se[price]"
        )
        expectClose(
            try #require(result.scalars["se__cons"]), fixtures["iv_se_cons"], "se[_cons]"
        )
        expectClose(try #require(result.scalars["r2"]), fixtures["iv_r2"], "R²")
        expectClose(try #require(result.scalars["rmse"]), fixtures["iv_rmse"], "root MSE")
        #expect(result.text.contains("Instrumented: price"))
    }

    @Test("ivregress robust SE matches R HC1-small")
    func ivRobust() async throws {
        let result = try await session.execute(
            "ivregress 2sls log_rev (price = z1 z2), robust"
        )
        expectClose(
            try #require(result.scalars["se_price"]),
            fixtures["iv_se_hc1_price"], "robust se[price]"
        )
    }

    @Test("ivregress order condition is enforced")
    func orderCondition() async throws {
        await #expect(throws: (any Error).self) {
            _ = try await session.execute(
                "ivregress 2sls log_rev (price orders = z1)"
            )
        }
    }
}