rainbow-dragon/Tests/ForgeGradeTests/FlattenerTests.swift

import XCTest
import ForgeColor
@testable import ForgeGrade

final class FlattenerTests: XCTestCase {

    /// Camera-domain stack for flattening: LUT input domain is camera log signal.
    /// Flattened LUT must reproduce stack output for any log-encoded input.
    func makeStack() -> GradeStack {
        var set = CurveSet.identity
        set.master = Curve(points: [.init(x: 0, y: 0.02), .init(x: 0.5, y: 0.45), .init(x: 1, y: 0.98)])
        return GradeStack(nodes: [
            GradeNode(kind: .cdl(CDL(
                slope: SIMD3(1.1, 0.95, 1.02),
                offset: SIMD3(0.01, -0.005, 0),
                power: SIMD3(0.98, 1.03, 1.0),
                saturation: 1.15))),
            GradeNode(kind: .curves(set)),
            GradeNode(kind: .saturation(0.9)),
        ])
    }

    // Flattened 33³ LUT matches direct stack eval on random in-domain samples.
    // Tolerance: trilinear error bound is max|f''|·h²/8 per axis (h=1/32). The
    // spline curve node has |f''| up to ~40 near segment junctions -> ~0.005
    // worst case, cross-terms add. 1% (10/1024) is the honest 33³ bound for
    // production-strength grades; 65³ tightens 4x (covered below).
    func testFlattenMatchesDirectEval() {
        let stack = makeStack()
        let look = Flattener.flatten(stack: stack, latticeSize: 33, splitLeadingCDL: false)
        XCTAssertNil(look.cdl)
        guard let lut = look.lut else { return XCTFail("no lut") }
        var rng = SystemRandomNumberGenerator()
        for _ in 0..<300 {
            let px = SIMD3<Float>(
                Float.random(in: 0...1, using: &rng),
                Float.random(in: 0...1, using: &rng),
                Float.random(in: 0...1, using: &rng))
            let direct = stack.evaluate(px)
            let viaLut = lut.sample(px)
            XCTAssertEqual(viaLut.x, direct.x, accuracy: 10.0 / 1024)
            XCTAssertEqual(viaLut.y, direct.y, accuracy: 10.0 / 1024)
            XCTAssertEqual(viaLut.z, direct.z, accuracy: 10.0 / 1024)
        }
    }

    // 65³ lattice: ~2x tighter than 33³. Error scales with h (not h²) because the
    // stack has derivative kinks — curve endpoint pinning and CDL pre-power clamp
    // are C0 surfaces that trilinear crosses. Observed max ~0.004 at 65³.
    func testFlatten65MatchesTighter() {
        let stack = makeStack()
        let look = Flattener.flatten(stack: stack, latticeSize: 65, splitLeadingCDL: false)
        guard let lut = look.lut else { return XCTFail("no lut") }
        var rng = SystemRandomNumberGenerator()
        for _ in 0..<300 {
            let px = SIMD3<Float>(
                Float.random(in: 0...1, using: &rng),
                Float.random(in: 0...1, using: &rng),
                Float.random(in: 0...1, using: &rng))
            let direct = stack.evaluate(px)
            let viaLut = lut.sample(px)
            XCTAssertEqual(viaLut.x, direct.x, accuracy: 5.0 / 1024)
            XCTAssertEqual(viaLut.y, direct.y, accuracy: 5.0 / 1024)
            XCTAssertEqual(viaLut.z, direct.z, accuracy: 5.0 / 1024)
        }
    }

    // Split mode: leading CDL excluded from bake, returned separately.
    // CDL.apply(then lut.sample) must equal full stack.
    func testSplitLeadingCDL() {
        let stack = makeStack()
        let look = Flattener.flatten(stack: stack, latticeSize: 33, splitLeadingCDL: true)
        guard let cdl = look.cdl, let lut = look.lut else { return XCTFail("missing parts") }

        // Returned CDL is stack's first node CDL.
        if case .cdl(let first) = stack.nodes[0].kind {
            XCTAssertEqual(cdl, first)
        } else {
            XCTFail("expected cdl first")
        }

        var rng = SystemRandomNumberGenerator()
        for _ in 0..<300 {
            let px = SIMD3<Float>(
                Float.random(in: 0...1, using: &rng),
                Float.random(in: 0...1, using: &rng),
                Float.random(in: 0...1, using: &rng))
            let direct = stack.evaluate(px)
            let recombined = lut.sample(cdl.apply(px))
            XCTAssertEqual(recombined.x, direct.x, accuracy: 10.0 / 1024)
            XCTAssertEqual(recombined.y, direct.y, accuracy: 10.0 / 1024)
            XCTAssertEqual(recombined.z, direct.z, accuracy: 10.0 / 1024)
        }
    }

    // Split requested but first node not CDL -> no split, full bake.
    func testSplitFallsBackWhenNoLeadingCDL() {
        let stack = GradeStack(nodes: [GradeNode(kind: .saturation(0.8))])
        let look = Flattener.flatten(stack: stack, latticeSize: 17, splitLeadingCDL: true)
        XCTAssertNil(look.cdl)
        XCTAssertNotNil(look.lut)
    }

    // Split requested but leading CDL bypassed -> not split.
    func testSplitIgnoresBypassedCDL() {
        var node = GradeNode(kind: .cdl(CDL(slope: SIMD3(2, 2, 2), offset: .zero, power: .one, saturation: 1)))
        node.isEnabled = false
        let stack = GradeStack(nodes: [node, GradeNode(kind: .saturation(0.8))])
        let look = Flattener.flatten(stack: stack, latticeSize: 17, splitLeadingCDL: true)
        XCTAssertNil(look.cdl)
    }

    // Deterministic: same stack -> byte-identical tables.
    func testDeterministic() {
        let stack = makeStack()
        let a = Flattener.flatten(stack: stack, latticeSize: 17, splitLeadingCDL: false)
        let b = Flattener.flatten(stack: stack, latticeSize: 17, splitLeadingCDL: false)
        XCTAssertEqual(a.lut?.table, b.lut?.table)
    }

    // Empty stack flattens to identity LUT.
    func testEmptyStackIdentityLut() {
        let look = Flattener.flatten(stack: GradeStack(nodes: []), latticeSize: 9, splitLeadingCDL: false)
        guard let lut = look.lut else { return XCTFail() }
        let px = SIMD3<Float>(0.3, 0.6, 0.9)
        let out = lut.sample(px)
        XCTAssertEqual(out.x, px.x, accuracy: 1e-4)
        XCTAssertEqual(out.y, px.y, accuracy: 1e-4)
        XCTAssertEqual(out.z, px.z, accuracy: 1e-4)
    }

    // Lattice size respected.
    func testLatticeSize() {
        let look = Flattener.flatten(stack: GradeStack(nodes: []), latticeSize: 65, splitLeadingCDL: false)
        XCTAssertEqual(look.lut?.size, 65)
    }
}