//
//  PredictTests.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("predict", .serialized)
struct PredictTests {
    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("xb + residuals reconstruct the response after regress")
    func fittedPlusResiduals() async throws {
        _ = try await session.execute("reg log_rev price")
        _ = try await session.execute("predict yhat")
        _ = try await session.execute("predict e, residuals")
        _ = try await session.execute("gen check = yhat + e - log_rev")
        let summary = try await session.execute("summarize check")
        #expect(abs(try #require(summary.scalars["min"])) < 1e-12)
        #expect(abs(try #require(summary.scalars["max"])) < 1e-12)
        // Missing prices → missing predictions.
        #expect(summary.scalars["N"] == 57)
        _ = try await session.execute("drop yhat e check")
    }

    @Test("predict evaluates factor-variable models on all observations")
    func factorModel() async throws {
        _ = try await session.execute("reg log_rev price i.region")
        _ = try await session.execute("predict fhat")
        _ = try await session.execute("predict fe2, residuals")
        _ = try await session.execute("gen fcheck = fhat + fe2 - log_rev")
        let summary = try await session.execute("summarize fcheck")
        #expect(abs(try #require(summary.scalars["max"])) < 1e-12)
        _ = try await session.execute("drop fhat fe2 fcheck")
    }

    @Test("logit default is predicted probability; fitted mean matches ybar")
    func logitProbability() async throws {
        _ = try await session.execute("logit purchase price")
        _ = try await session.execute("predict p")
        let pSummary = try await session.execute("summarize p")
        let ySummary = try await session.execute("summarize purchase if !missing(price)")
        // With an intercept, mean fitted probability equals the sample
        // mean of the outcome (logit score equation).
        expectClose(
            try #require(pSummary.scalars["mean"]),
            try #require(ySummary.scalars["mean"]),
            rtol: 1e-9, "mean fitted probability"
        )
        #expect(try #require(pSummary.scalars["min"]) > 0)
        #expect(try #require(pSummary.scalars["max"]) < 1)
        _ = try await session.execute("drop p")
    }

    @Test("poisson default is the predicted mean count")
    func poissonMean() async throws {
        _ = try await session.execute("poisson orders price")
        _ = try await session.execute("predict mu")
        let muSummary = try await session.execute("summarize mu")
        let ySummary = try await session.execute("summarize orders if !missing(price)")
        // Poisson with intercept: mean fitted count equals ybar.
        expectClose(
            try #require(muSummary.scalars["mean"]),
            try #require(ySummary.scalars["mean"]),
            rtol: 1e-9, "mean fitted count"
        )
        _ = try await session.execute("drop mu")
    }

    @Test("predict before any estimation errors clearly")
    func requiresEstimation() 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("predict yhat")
        }
    }

    @Test("statistic options are validated against the model kind")
    func statisticValidation() async throws {
        _ = try await session.execute("reg log_rev price")
        await #expect(throws: ZQEngineError.self) {
            _ = try await session.execute("predict p, pr")   // pr needs logit/probit
        }
        await #expect(throws: ZQEngineError.self) {
            _ = try await session.execute("predict m, n")    // n needs poisson
        }
    }
}