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

import Foundation

/// Gradient-boosted regression trees with squared-error loss, following
/// xgboost's exact-greedy algorithm precisely so results cross-validate:
///
///   gain  = ½[G_L²/(H_L+λ) + G_R²/(H_R+λ) − (G_L+G_R)²/(H_L+H_R+λ)] − γ
///   leaf  = −G/(H+λ)
///
/// with g_i = ŷ_i − y_i and h_i = 1 for squared loss. Splits are placed
/// at midpoints between consecutive distinct feature values (xgboost's
/// convention), `x < split` goes left, and missing features follow the
/// default-left rule. No subsampling — training is deterministic.
public struct ZQBoostModel: Equatable, Sendable {
    public struct Node: Equatable, Sendable {
        /// Feature index for internal nodes; nil marks a leaf.
        public var feature: Int?
        public var split: Double
        public var left: Int
        public var right: Int
        public var value: Double   // leaf weight
    }

    public struct Tree: Equatable, Sendable {
        public var nodes: [Node]

        public func predict(_ features: [Double]) -> Double {
            var index = 0
            while let feature = nodes[index].feature {
                let value = features[feature]
                index = value.isNaN || value < nodes[index].split
                    ? nodes[index].left
                    : nodes[index].right
            }
            return nodes[index].value
        }
    }

    public var featureNames: [String]
    public var baseScore: Double
    public var learningRate: Double
    public var trees: [Tree]

    /// Prediction for one observation's feature vector (ordered as
    /// `featureNames`; NaN = missing).
    public func predict(_ features: [Double]) -> Double {
        var result = baseScore
        for tree in trees {
            result += learningRate * tree.predict(features)
        }
        return result
    }
}

public enum ZQGradientBoosting {

    /// - Parameters:
    ///   - rounds: number of trees.
    ///   - learningRate: shrinkage η applied to every leaf.
    ///   - maxDepth: maximum tree depth (1 = stumps).
    ///   - lambda: L2 regularization on leaf weights (xgboost default 1).
    ///   - gamma: minimum gain to split (default 0).
    ///   - minChildWeight: minimum hessian sum per child (= observation
    ///     count under squared loss; default 1).
    ///   - baseScore: initial prediction; nil = mean of y.
    public static func fit(
        y: [Double],
        features: [(name: String, values: [Double])],
        rounds: Int,
        learningRate: Double = 0.3,
        maxDepth: Int = 6,
        lambda: Double = 1,
        gamma: Double = 0,
        minChildWeight: Double = 1,
        baseScore: Double? = nil
    ) throws -> ZQBoostModel {
        let n = y.count
        let p = features.count
        guard p > 0 else { throw ZQStatsError("boost: features required") }
        guard n > 1 else { throw ZQStatsError("boost: insufficient observations") }
        guard rounds > 0 else { throw ZQStatsError("boost: rounds() must be positive") }
        guard maxDepth >= 1 else { throw ZQStatsError("boost: maxdepth() must be ≥ 1") }
        for column in features where column.values.count != n {
            throw ZQStatsError("boost: feature '\(column.name)' has wrong length")
        }

        // Pre-sorted feature indices (missing excluded; they default left).
        var sortedIndices = [[Int]](repeating: [], count: p)
        for j in 0..<p {
            let values = features[j].values
            sortedIndices[j] = (0..<n)
                .filter { !values[$0].isNaN }
                .sorted { values[$0] < values[$1] }
        }

        let base = baseScore ?? y.reduce(0, +) / Double(n)
        var predictions = [Double](repeating: base, count: n)
        var trees: [ZQBoostModel.Tree] = []

        for _ in 0..<rounds {
            // Squared loss: g = ŷ − y, h = 1.
            let gradients = (0..<n).map { predictions[$0] - y[$0] }
            var tree = ZQBoostModel.Tree(nodes: [])
            buildNode(
                rows: Array(0..<n),
                depth: 0,
                features: features,
                sortedIndices: sortedIndices,
                gradients: gradients,
                maxDepth: maxDepth,
                lambda: lambda,
                gamma: gamma,
                minChildWeight: minChildWeight,
                into: &tree.nodes
            )
            for i in 0..<n {
                let vector = (0..<p).map { features[$0].values[i] }
                predictions[i] += learningRate * tree.predict(vector)
            }
            trees.append(tree)
        }

        return ZQBoostModel(
            featureNames: features.map(\.name),
            baseScore: base,
            learningRate: learningRate,
            trees: trees
        )
    }

    /// Recursively grows one node; returns its index in `nodes`.
    @discardableResult
    private static func buildNode(
        rows: [Int],
        depth: Int,
        features: [(name: String, values: [Double])],
        sortedIndices: [[Int]],
        gradients: [Double],
        maxDepth: Int,
        lambda: Double,
        gamma: Double,
        minChildWeight: Double,
        into nodes: inout [ZQBoostModel.Node]
    ) -> Int {
        let rowSet = Set(rows)
        let gTotal = rows.reduce(0.0) { $0 + gradients[$1] }
        let hTotal = Double(rows.count)

        let index = nodes.count
        nodes.append(ZQBoostModel.Node(
            feature: nil, split: 0, left: -1, right: -1,
            value: -gTotal / (hTotal + lambda)
        ))
        guard depth < maxDepth, rows.count > 1 else { return index }

        // Exact greedy split search over every feature.
        let parentScore = gTotal * gTotal / (hTotal + lambda)
        var bestGain = 0.0
        var bestFeature = -1
        var bestSplit = 0.0
        for j in 0..<features.count {
            let values = features[j].values
            let ordered = sortedIndices[j].filter { rowSet.contains($0) }
            guard ordered.count > 1 else { continue }
            let gMissing = gTotal - ordered.reduce(0.0) { $0 + gradients[$1] }
            let hMissing = hTotal - Double(ordered.count)

            var gLeft = gMissing   // missing rows follow the left child
            var hLeft = hMissing
            for position in 0..<(ordered.count - 1) {
                gLeft += gradients[ordered[position]]
                hLeft += 1
                let current = values[ordered[position]]
                let next = values[ordered[position + 1]]
                guard next > current else { continue }
                let gRight = gTotal - gLeft
                let hRight = hTotal - hLeft
                guard hLeft >= minChildWeight, hRight >= minChildWeight else { continue }
                let gain = 0.5 * (
                    gLeft * gLeft / (hLeft + lambda)
                        + gRight * gRight / (hRight + lambda)
                        - parentScore
                ) - gamma
                if gain > bestGain {
                    bestGain = gain
                    bestFeature = j
                    bestSplit = (current + next) / 2
                }
            }
        }
        guard bestFeature >= 0 else { return index }

        let values = features[bestFeature].values
        let leftRows = rows.filter { values[$0].isNaN || values[$0] < bestSplit }
        let rightRows = rows.filter { !values[$0].isNaN && values[$0] >= bestSplit }
        guard !leftRows.isEmpty, !rightRows.isEmpty else { return index }

        nodes[index].feature = bestFeature
        nodes[index].split = bestSplit
        nodes[index].left = buildNode(
            rows: leftRows, depth: depth + 1, features: features,
            sortedIndices: sortedIndices, gradients: gradients,
            maxDepth: maxDepth, lambda: lambda, gamma: gamma,
            minChildWeight: minChildWeight, into: &nodes
        )
        nodes[index].right = buildNode(
            rows: rightRows, depth: depth + 1, features: features,
            sortedIndices: sortedIndices, gradients: gradients,
            maxDepth: maxDepth, lambda: lambda, gamma: gamma,
            minChildWeight: minChildWeight, into: &nodes
        )
        return index
    }
}