import Foundation

/// 3D LUT lattice. Table layout matches .cube: red index fastest, then green, then blue.
/// Entry (r,g,b) at index b*size² + g*size + r, three floats per entry.
public struct Lut3D: Sendable {
    public let size: Int
    /// size³ * 3 floats.
    public let table: [Float]

    public init(size: Int, table: [Float]) {
        precondition(size >= 2, "LUT size must be >= 2")
        precondition(table.count == size * size * size * 3, "table count mismatch")
        self.size = size
        self.table = table
    }

    public static func identity(size: Int) -> Lut3D {
        build(size: size) { $0 }
    }

    /// Build lattice by evaluating a function at each grid point.
    public static func build(size: Int, function: (SIMD3<Float>) -> SIMD3<Float>) -> Lut3D {
        var table = [Float]()
        table.reserveCapacity(size * size * size * 3)
        let denom = Float(size - 1)
        for b in 0..<size {
            for g in 0..<size {
                for r in 0..<size {
                    let out = function(SIMD3(Float(r) / denom, Float(g) / denom, Float(b) / denom))
                    table.append(out.x)
                    table.append(out.y)
                    table.append(out.z)
                }
            }
        }
        return Lut3D(size: size, table: table)
    }

    @inline(__always)
    private func entry(_ r: Int, _ g: Int, _ b: Int) -> SIMD3<Float> {
        let i = ((b * size + g) * size + r) * 3
        return SIMD3(table[i], table[i + 1], table[i + 2])
    }

    /// True if lattice ≈ identity within tolerance (sampled grid corners).
    /// Used to detect "nothing left after CDL split" flatten results.
    public func isApproximatelyIdentity(tolerance: Float = 1.0 / 512) -> Bool {
        let denom = Float(size - 1)
        let step = max(1, size / 4)
        for b in stride(from: 0, to: size, by: step) {
            for g in stride(from: 0, to: size, by: step) {
                for r in stride(from: 0, to: size, by: step) {
                    let idx = ((b * size + g) * size + r) * 3
                    if abs(table[idx] - Float(r) / denom) > tolerance ||
                        abs(table[idx + 1] - Float(g) / denom) > tolerance ||
                        abs(table[idx + 2] - Float(b) / denom) > tolerance {
                        return false
                    }
                }
            }
        }
        return true
    }

    /// Trilinear interpolation; input clamped to [0,1].
    public func sample(_ rgb: SIMD3<Float>) -> SIMD3<Float> {
        let maxIdx = Float(size - 1)
        let clamped = SIMD3(
            min(max(rgb.x, 0), 1),
            min(max(rgb.y, 0), 1),
            min(max(rgb.z, 0), 1))
        let pos = clamped * maxIdx
        let i0 = SIMD3<Int>(Int(pos.x), Int(pos.y), Int(pos.z))
        let i1 = SIMD3<Int>(
            min(i0.x + 1, size - 1),
            min(i0.y + 1, size - 1),
            min(i0.z + 1, size - 1))
        let f = pos - SIMD3(Float(i0.x), Float(i0.y), Float(i0.z))

        // 8 corners; lerp r, then g, then b.
        let c000 = entry(i0.x, i0.y, i0.z), c100 = entry(i1.x, i0.y, i0.z)
        let c010 = entry(i0.x, i1.y, i0.z), c110 = entry(i1.x, i1.y, i0.z)
        let c001 = entry(i0.x, i0.y, i1.z), c101 = entry(i1.x, i0.y, i1.z)
        let c011 = entry(i0.x, i1.y, i1.z), c111 = entry(i1.x, i1.y, i1.z)

        let c00 = c000 + (c100 - c000) * f.x
        let c10 = c010 + (c110 - c010) * f.x
        let c01 = c001 + (c101 - c001) * f.x
        let c11 = c011 + (c111 - c011) * f.x
        let c0 = c00 + (c10 - c00) * f.y
        let c1 = c01 + (c11 - c01) * f.y
        return c0 + (c1 - c0) * f.z
    }
}

// MARK: - .cube I/O

public enum CubeError: Error, Equatable {
    case missingSize
    case invalidFloat(line: Int)
    case entryCountMismatch(expected: Int, got: Int)
    case invalidSize
}

public enum CubeFile {
    /// Serialize to Resolve/IRIDAS .cube text.
    public static func write(_ lut: Lut3D, title: String = "Forge") -> String {
        var out = "TITLE \"\(title)\"\nLUT_3D_SIZE \(lut.size)\n"
        out.reserveCapacity(lut.table.count * 10)
        var i = 0
        while i < lut.table.count {
            out += String(format: "%.6f %.6f %.6f\n", lut.table[i], lut.table[i + 1], lut.table[i + 2])
            i += 3
        }
        return out
    }

    /// Parse .cube text. Ignores TITLE/DOMAIN_MIN/DOMAIN_MAX/comments/blank lines.
    public static func parse(_ text: String) throws -> Lut3D {
        var size: Int?
        var values = [Float]()
        var lineNo = 0
        for rawLine in text.split(separator: "\n", omittingEmptySubsequences: false) {
            lineNo += 1
            let line = rawLine.trimmingCharacters(in: .whitespaces)
            if line.isEmpty || line.hasPrefix("#") { continue }
            if line.hasPrefix("TITLE") || line.hasPrefix("DOMAIN_MIN") || line.hasPrefix("DOMAIN_MAX")
                || line.hasPrefix("LUT_1D_SIZE") {
                continue
            }
            if line.hasPrefix("LUT_3D_SIZE") {
                let parts = line.split(separator: " ")
                guard parts.count == 2, let s = Int(parts[1]), s >= 2 else { throw CubeError.invalidSize }
                size = s
                continue
            }
            // Data row: three floats.
            let comps = line.split(separator: " ", omittingEmptySubsequences: true)
            guard comps.count == 3 else { throw CubeError.invalidFloat(line: lineNo) }
            for c in comps {
                guard let v = Float(c) else { throw CubeError.invalidFloat(line: lineNo) }
                values.append(v)
            }
        }
        guard let s = size else { throw CubeError.missingSize }
        let expected = s * s * s * 3
        guard values.count == expected else {
            throw CubeError.entryCountMismatch(expected: expected, got: values.count)
        }
        return Lut3D(size: s, table: values)
    }
}