import XCTest
@testable import MarkBase

final class MathTest: XCTestCase {

    // MARK: - GELU (tanh approximation)

    func testGELUZero() {
        let result = gelu(0.0)
        XCTAssertEqual(result, 0.0, accuracy: 1e-6)
    }

    func testGELUPositive() {
        let result = gelu(1.0)
        XCTAssertEqual(result, 0.841192, accuracy: 1e-4)
    }

    func testGELUNegative() {
        let result = gelu(-1.0)
        XCTAssertEqual(result, -0.158808, accuracy: 1e-4)
    }

    func testGELULargePositive() {
        let result = gelu(5.0)
        XCTAssertEqual(result, 5.0, accuracy: 1e-4)
    }

    func testGELULargeNegative() {
        let result = gelu(-5.0)
        XCTAssertEqual(result, 0.0, accuracy: 1e-4)
    }

    // MARK: - RMS Normalization

    func testRMSNormBasic() {
        let input: [Float] = [1.0, 2.0, 3.0, 4.0]
        let weight: [Float] = [0.5, 0.5, 0.5, 0.5]

        let rms = sqrt(input.map { $0 * $0 }.reduce(0, +) / Float(input.count))
        let scale: Float = 1.0 / (rms + 1e-6)
        let expected = input.map { $0 * scale * 0.5 }

        let output = rmsNorm(input: input, weight: weight, eps: 1e-6)
        for (o, e) in zip(output, expected) {
            XCTAssertEqual(o, e, accuracy: 1e-5)
        }
    }

    func testRMSNormAllZeros() {
        let input: [Float] = [0, 0, 0]
        let weight: [Float] = [1, 1, 1]

        let output = rmsNorm(input: input, weight: weight, eps: 1e-6)
        XCTAssertEqual(output, [0, 0, 0])
    }

    func testRMSNormSingleElement() {
        let input: [Float] = [3.0]
        let weight: [Float] = [2.0]

        let expected: Float = 2.0
        let output = rmsNorm(input: input, weight: weight, eps: 1e-6)
        XCTAssertEqual(output[0], expected, accuracy: 1e-5)
    }

    // MARK: - Element-wise Addition

    func testEltwiseAdd() {
        let a: [Float] = [1, 2, 3]
        let b: [Float] = [4, 5, 6]
        let result = eltwiseAdd(a: a, b: b)
        XCTAssertEqual(result, [5, 7, 9])
    }

    func testEltwiseAddDifferentLengths() {
        let a: [Float] = [1, 2]
        let b: [Float] = [3, 4, 5]
        let result = eltwiseAdd(a: a, b: b)
        XCTAssertEqual(result.count, min(a.count, b.count))
        XCTAssertEqual(result, [4, 6])
    }

    // MARK: - Element-wise Multiplication

    func testEltwiseMul() {
        let a: [Float] = [1, 2, 3]
        let b: [Float] = [4, 5, 6]
        let result = eltwiseMul(a: a, b: b)
        XCTAssertEqual(result, [4, 10, 18])
    }

    func testEltwiseMulByZero() {
        let a: [Float] = [1, 2, 3]
        let b: [Float] = [0, 0, 0]
        let result = eltwiseMul(a: a, b: b)
        XCTAssertEqual(result, [0, 0, 0])
    }

    func testEltwiseMulNegative() {
        let a: [Float] = [2, -3]
        let b: [Float] = [-4, 5]
        let result = eltwiseMul(a: a, b: b)
        XCTAssertEqual(result, [-8, -15])
    }
}

// MARK: - Helper implementations for CPU-based tests

func gelu(_ x: Float) -> Float {
    let tanhVal = tanh(0.79788456 * (x + 0.044715 * x * x * x))
    return 0.5 * x * (1.0 + tanhVal)
}

func rmsNorm(input: [Float], weight: [Float], eps: Float) -> [Float] {
    let count = input.count
    let rms = sqrt(input.map { $0 * $0 }.reduce(0, +) / Float(count) + eps)
    let scale: Float = 1.0 / rms
    return zip(input, weight).map { $0 * scale * $1 }
}

func eltwiseAdd(a: [Float], b: [Float]) -> [Float] {
    return zip(a, b).map(+)
}

func eltwiseMul(a: [Float], b: [Float]) -> [Float] {
    return zip(a, b).map(*)
}