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

import Foundation
import ZQData
import ZQParser

/// Vectorized evaluation of a `ZQExpression` against the working dataset.
///
/// Missing-value semantics: any arithmetic or comparison touching a
/// missing operand yields missing; undefined operations (division by
/// zero, log of a non-positive number) yield missing. The `missing(x)`
/// function itself never returns missing. In `if` qualifiers a missing
/// condition excludes the observation.
struct ExpressionEvaluator {
    struct EvaluationError: Error, CustomStringConvertible {
        let message: String
        var description: String { message }
    }

    /// A column vector or a broadcastable scalar.
    enum Value {
        case numeric(values: [Double], missing: [Bool])
        case strings([String?])
        case numericScalar(Double)
        case stringScalar(String)
        case missingScalar
    }

    let frame: ZQDataFrame

    func evaluate(_ expression: ZQExpression) throws -> Value {
        switch expression {
        case .number(let value):
            return .numericScalar(value)
        case .string(let value):
            return .stringScalar(value)
        case .missing:
            return .missingScalar
        case .variable(let name):
            guard let column = frame.column(named: name) else {
                throw EvaluationError(message: "variable '\(name)' not found")
            }
            switch column.data {
            case .float64(let values, let missing):
                return .numeric(values: values, missing: missing)
            case .string(let values):
                return .strings(values)
            }
        case .unary(let op, let operand):
            return try evaluateUnary(op: op, operand: operand)
        case .binary(let op, let lhs, let rhs):
            return try evaluateBinary(op: op, lhs: lhs, rhs: rhs)
        case .call(let name, let arguments):
            return try evaluateCall(name: name, arguments: arguments)
        }
    }

    /// Evaluates to a full-length numeric column (broadcasting scalars).
    func evaluateNumericColumn(
        _ expression: ZQExpression
    ) throws -> (values: [Double], missing: [Bool]) {
        switch try evaluate(expression) {
        case .numeric(let values, let missing):
            return (values, missing)
        case .numericScalar(let value):
            return (
                [Double](repeating: value, count: frame.rowCount),
                [Bool](repeating: false, count: frame.rowCount)
            )
        case .missingScalar:
            return (
                [Double](repeating: .nan, count: frame.rowCount),
                [Bool](repeating: true, count: frame.rowCount)
            )
        case .strings, .stringScalar:
            throw EvaluationError(message: "type mismatch: expected a numeric expression")
        }
    }

    /// Evaluates an `if` qualifier into a keep-mask. Missing → excluded.
    func evaluateCondition(_ expression: ZQExpression) throws -> [Bool] {
        let (values, missing) = try evaluateNumericColumn(expression)
        return (0..<values.count).map { !missing[$0] && values[$0] != 0 }
    }

    // MARK: - Operators

    private func evaluateUnary(op: String, operand: ZQExpression) throws -> Value {
        let value = try evaluate(operand)
        switch op {
        case "-":
            return try mapNumeric(value) { -$0 }
        case "!":
            return try mapNumeric(value) { $0 == 0 ? 1 : 0 }
        default:
            throw EvaluationError(message: "unknown unary operator '\(op)'")
        }
    }

    private func evaluateBinary(
        op: String, lhs: ZQExpression, rhs: ZQExpression
    ) throws -> Value {
        let left = try evaluate(lhs)
        let right = try evaluate(rhs)

        // String equality against a literal or another string column.
        if isString(left) || isString(right) {
            guard op == "==" || op == "!=" else {
                throw EvaluationError(
                    message: "operator '\(op)' is not defined for strings"
                )
            }
            return try compareStrings(left, right, equal: op == "==")
        }

        let combine: (Double, Double) -> Double
        switch op {
        case "+": combine = (+)
        case "-": combine = (-)
        case "*": combine = (*)
        case "/": combine = { $1 == 0 ? .nan : $0 / $1 }
        case "^": combine = { pow($0, $1) }
        case "==": combine = { $0 == $1 ? 1 : 0 }
        case "!=": combine = { $0 != $1 ? 1 : 0 }
        case "<": combine = { $0 < $1 ? 1 : 0 }
        case "<=": combine = { $0 <= $1 ? 1 : 0 }
        case ">": combine = { $0 > $1 ? 1 : 0 }
        case ">=": combine = { $0 >= $1 ? 1 : 0 }
        case "&": combine = { ($0 != 0 && $1 != 0) ? 1 : 0 }
        case "|": combine = { ($0 != 0 || $1 != 0) ? 1 : 0 }
        default:
            throw EvaluationError(message: "unknown operator '\(op)'")
        }
        return try zipNumeric(left, right, combine)
    }

    private func evaluateCall(name: String, arguments: [ZQExpression]) throws -> Value {
        // missing(x): 1/0 indicator, itself never missing.
        if name == "missing" || name == "mi" {
            guard arguments.count == 1 else {
                throw EvaluationError(message: "missing() takes exactly one argument")
            }
            switch try evaluate(arguments[0]) {
            case .numeric(_, let missing):
                return .numeric(
                    values: missing.map { $0 ? 1.0 : 0.0 },
                    missing: [Bool](repeating: false, count: missing.count)
                )
            case .strings(let values):
                return .numeric(
                    values: values.map { $0 == nil ? 1.0 : 0.0 },
                    missing: [Bool](repeating: false, count: values.count)
                )
            case .missingScalar:
                return .numericScalar(1)
            case .numericScalar, .stringScalar:
                return .numericScalar(0)
            }
        }

        let unaryFunctions: [String: @Sendable (Double) -> Double] = [
            "ln": { $0 > 0 ? Foundation.log($0) : .nan },
            "log": { $0 > 0 ? Foundation.log($0) : .nan },
            "log10": { $0 > 0 ? Foundation.log10($0) : .nan },
            "exp": Foundation.exp,
            "sqrt": { $0 >= 0 ? $0.squareRoot() : .nan },
            "abs": abs,
            "floor": { $0.rounded(.down) },
            "ceil": { $0.rounded(.up) },
            "round": { $0.rounded(.toNearestOrAwayFromZero) },
            "int": { $0.rounded(.towardZero) },
        ]
        if let function = unaryFunctions[name] {
            guard arguments.count == 1 else {
                throw EvaluationError(message: "\(name)() takes exactly one argument")
            }
            return try mapNumeric(try evaluate(arguments[0]), function)
        }

        if name == "max" || name == "min" {
            guard arguments.count >= 2 else {
                throw EvaluationError(message: "\(name)() takes at least two arguments")
            }
            var result = try evaluate(arguments[0])
            for argument in arguments.dropFirst() {
                let next = try evaluate(argument)
                result = try zipNumeric(result, next, name == "max" ? max : min)
            }
            return result
        }

        throw EvaluationError(message: "unknown function '\(name)()'")
    }

    // MARK: - Helpers

    private func isString(_ value: Value) -> Bool {
        switch value {
        case .strings, .stringScalar: return true
        default: return false
        }
    }

    private func compareStrings(
        _ lhs: Value, _ rhs: Value, equal: Bool
    ) throws -> Value {
        func stringAt(_ value: Value, _ index: Int) -> String?? {
            switch value {
            case .strings(let values): return values[index]
            case .stringScalar(let value): return value
            default: return Optional<String?>.none
            }
        }
        let count: Int
        switch (lhs, rhs) {
        case (.strings(let values), _), (_, .strings(let values)):
            count = values.count
        default:
            // scalar vs scalar
            if case .stringScalar(let a) = lhs, case .stringScalar(let b) = rhs {
                return .numericScalar((a == b) == equal ? 1 : 0)
            }
            throw EvaluationError(message: "type mismatch in string comparison")
        }
        var values = [Double](repeating: 0, count: count)
        let missing = [Bool](repeating: false, count: count)
        for i in 0..<count {
            guard let a = stringAt(lhs, i), let b = stringAt(rhs, i) else {
                throw EvaluationError(message: "type mismatch in string comparison")
            }
            switch (a, b) {
            case (nil, nil), (nil, _), (_, nil):
                // Stata treats "" as the string missing value; comparisons
                // against it still evaluate.
                let equalValues = (a ?? "") == (b ?? "")
                values[i] = equalValues == equal ? 1 : 0
            default:
                values[i] = (a! == b!) == equal ? 1 : 0
            }
        }
        return .numeric(values: values, missing: missing)
    }

    private func mapNumeric(
        _ value: Value, _ transform: (Double) -> Double
    ) throws -> Value {
        switch value {
        case .numeric(let values, var missing):
            var result = [Double](repeating: .nan, count: values.count)
            for i in 0..<values.count where !missing[i] {
                result[i] = transform(values[i])
                if result[i].isNaN { missing[i] = true }
            }
            return .numeric(values: result, missing: missing)
        case .numericScalar(let value):
            let result = transform(value)
            return result.isNaN ? .missingScalar : .numericScalar(result)
        case .missingScalar:
            return .missingScalar
        case .strings, .stringScalar:
            throw EvaluationError(message: "type mismatch: expected a numeric operand")
        }
    }

    private func zipNumeric(
        _ lhs: Value, _ rhs: Value, _ combine: (Double, Double) -> Double
    ) throws -> Value {
        switch (lhs, rhs) {
        case (.numericScalar(let a), .numericScalar(let b)):
            let result = combine(a, b)
            return result.isNaN ? .missingScalar : .numericScalar(result)
        case (.missingScalar, _), (_, .missingScalar):
            if case .numeric(let values, _) = lhs {
                return .numeric(
                    values: [Double](repeating: .nan, count: values.count),
                    missing: [Bool](repeating: true, count: values.count)
                )
            }
            if case .numeric(let values, _) = rhs {
                return .numeric(
                    values: [Double](repeating: .nan, count: values.count),
                    missing: [Bool](repeating: true, count: values.count)
                )
            }
            return .missingScalar
        default:
            let (leftValues, leftMissing) = try broadcast(lhs)
            let (rightValues, rightMissing) = try broadcast(rhs)
            guard leftValues.count == rightValues.count else {
                throw EvaluationError(message: "operand length mismatch")
            }
            var values = [Double](repeating: .nan, count: leftValues.count)
            var missing = [Bool](repeating: true, count: leftValues.count)
            for i in 0..<leftValues.count where !leftMissing[i] && !rightMissing[i] {
                let result = combine(leftValues[i], rightValues[i])
                if !result.isNaN {
                    values[i] = result
                    missing[i] = false
                }
            }
            return .numeric(values: values, missing: missing)
        }
    }

    private func broadcast(
        _ value: Value
    ) throws -> (values: [Double], missing: [Bool]) {
        switch value {
        case .numeric(let values, let missing):
            return (values, missing)
        case .numericScalar(let scalar):
            return (
                [Double](repeating: scalar, count: frame.rowCount),
                [Bool](repeating: false, count: frame.rowCount)
            )
        case .missingScalar:
            return (
                [Double](repeating: .nan, count: frame.rowCount),
                [Bool](repeating: true, count: frame.rowCount)
            )
        case .strings, .stringScalar:
            throw EvaluationError(message: "type mismatch: expected a numeric operand")
        }
    }
}