//
//  BoostTests.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 ZQGPU
import ZQStats

/// Gradient boosting vs R xgboost (exact greedy, no subsampling —
/// deterministic). xgboost computes in float32, so per-observation
/// prediction agreement is asserted at 1e-4 relative.
@Suite("Gradient boosting", .serialized)
struct BoostTests {
    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("matches xgboost predictions observation by observation")
    func xgboostParity() async throws {
        let result = try await session.execute(
            "boost log_rev z1 z2 orders, rounds(20) eta(0.3) maxdepth(3) lambda(1)"
        )
        #expect(result.scalars["N"] == 60)
        expectClose(
            try #require(result.scalars["rmse"]),
            fixtures["gbm_rmse"], rtol: 1e-4, "training RMSE"
        )

        _ = try await session.execute("predict bhat")
        let frame = await session.frame
        let (predictions, _) = try frame.requireNumeric("bhat")
        for (row, key) in [(0, "gbm_pred_1"), (16, "gbm_pred_17"),
                           (41, "gbm_pred_42"), (56, "gbm_pred_57")] {
            expectClose(
                predictions[row], fixtures[key], rtol: 1e-4,
                "prediction[\(row + 1)]"
            )
        }
        _ = try await session.execute("drop bhat")
    }

    @Test("a single stump finds the obvious split")
    func stumpSplit() throws {
        // Step function: y = 0 for x < 5, y = 10 for x ≥ 5.
        let x = (0..<20).map(Double.init)
        let y = x.map { $0 < 5 ? 0.0 : 10.0 }
        let model = try ZQGradientBoosting.fit(
            y: y, features: [("x", x)],
            rounds: 1, learningRate: 1, maxDepth: 1, lambda: 0
        )
        let tree = model.trees[0]
        #expect(tree.nodes[0].feature == 0)
        #expect(tree.nodes[0].split == 4.5)
        // Leaf values are exact means of the residuals under lambda 0.
        #expect(abs(model.predict([2]) - 0) < 1e-12)
        #expect(abs(model.predict([9]) - 10) < 1e-12)
    }

    @Test("training loss decreases every round")
    func monotoneTrainingLoss() throws {
        var stream = PhiloxStream(seed: 5)
        let n = 150
        let x1 = (0..<n).map { _ in stream.nextUniform() * 10 }
        let x2 = (0..<n).map { _ in stream.nextUniform() * 10 }
        let y = (0..<n).map { i in
            Foundation.sin(x1[i]) + 0.3 * x2[i] + 0.2 * stream.nextNormal()
        }

        var previous = Double.infinity
        for rounds in [1, 5, 20, 80] {
            let model = try ZQGradientBoosting.fit(
                y: y, features: [("x1", x1), ("x2", x2)],
                rounds: rounds, learningRate: 0.3, maxDepth: 3
            )
            var rss = 0.0
            for i in 0..<n {
                let predicted = model.predict([x1[i], x2[i]])
                rss += (y[i] - predicted) * (y[i] - predicted)
            }
            #expect(rss < previous, "rounds \(rounds): rss \(rss)")
            previous = rss
        }
        // Deep boosting on smooth data should fit very well in-sample.
        #expect(previous / Double(n) < 0.05)
    }

    @Test("boost requires rounds() and margins refuses afterwards")
    func validation() async throws {
        await #expect(throws: ZQEngineError.self) {
            _ = try await session.execute("boost log_rev z1")
        }
        _ = try await session.execute("boost log_rev z1 z2, rounds(5)")
        await #expect(throws: ZQEngineError.self) {
            _ = try await session.execute("margins, dydx(z1)")
        }
        await #expect(throws: ZQEngineError.self) {
            _ = try await session.execute("predict p, pr")
        }
    }
}