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

import ZQData
import ZQParser

/// Execution backend. The planner chooses automatically; the user never
/// does (CLAUDE.md §5).
public enum ZQBackend: String, Equatable, Sendable {
    case cpu, gpu, hybrid
}

/// A planned command: the parsed AST plus dispatch decision. Factor
/// variables expand here (plan time), not at parse time.
public struct ZQExecutionPlan: Equatable, Sendable {
    public var command: ZQCommand
    public var backend: ZQBackend

    public init(command: ZQCommand, backend: ZQBackend) {
        self.command = command
        self.backend = backend
    }
}

public struct ZQPlanner: Sendable {
    /// Single-estimation row-count threshold below which the CPU LAPACK
    /// path wins (heuristic 1, tuned by Tests/Bench).
    public var cpuRowThreshold: Int
    /// Independent-replicate threshold at which batched GPU solves win
    /// (heuristic 2).
    public var gpuReplicateThreshold: Int
    /// Whether a GPU backend is available at all. The pure-Swift Philox
    /// reference keeps plans reproducible even when this is false.
    public var gpuAvailable: Bool

    public init(
        cpuRowThreshold: Int = 5_000_000,
        gpuReplicateThreshold: Int = 500,
        gpuAvailable: Bool = false
    ) {
        self.cpuRowThreshold = cpuRowThreshold
        self.gpuReplicateThreshold = gpuReplicateThreshold
        self.gpuAvailable = gpuAvailable
    }

    public func plan(_ command: ZQCommand, rowCount: Int) -> ZQExecutionPlan {
        ZQExecutionPlan(command: command, backend: backend(for: command, rowCount: rowCount))
    }

    private func backend(for command: ZQCommand, rowCount: Int) -> ZQBackend {
        guard gpuAvailable else { return .cpu }

        // Heuristic 2: ≥ threshold independent replicates → batched GPU.
        // (permute joins once its GPU argsort path lands; the permutation
        // definition is already argsort-of-Philox-keys for that reason.)
        if command.verb == "bootstrap" {
            if let reps = command.option("reps")?.firstArgument,
               let count = Int(reps), count >= gpuReplicateThreshold {
                return .gpu
            }
        }

        // Heuristic 1: single estimation under the row threshold → CPU.
        if rowCount < cpuRowThreshold { return .cpu }

        // Large single estimations stay on CPU until the batched GPU
        // estimators land (v0.2); MCMC-style workloads will plan .hybrid.
        return .cpu
    }
}

/// Plan-time factor-variable expansion (`i.region` → one indicator column
/// per level, base level omitted).
public enum ZQFactorExpansion {

    /// Distinct sorted levels of a numeric variable, ignoring missing.
    public static func levels(values: [Double], missing: [Bool]) -> [Double] {
        var seen = Set<Double>()
        for i in 0..<values.count where !missing[i] {
            seen.insert(values[i])
        }
        return seen.sorted()
    }

    /// Expands one factor spec into named indicator columns against the
    /// given data. The lowest level is the omitted base (Stata default).
    /// The numeric level rides along so post-estimation commands can
    /// recompute the indicator on other observations.
    public static func expand(
        name: String, values: [Double], missing: [Bool]
    ) -> [(name: String, level: Double, values: [Double])] {
        let allLevels = levels(values: values, missing: missing)
        guard allLevels.count > 1 else { return [] }
        return allLevels.dropFirst().map { level in
            let rendered = level == level.rounded()
                ? String(Int(level))
                : String(level)
            let indicator = values.enumerated().map { index, value in
                (!missing[index] && value == level) ? 1.0 : 0.0
            }
            return (name: "\(rendered).\(name)", level: level, values: indicator)
        }
    }
}