essentiaAnalysis.kt

@file:Suppress("unused")
@file:afield.runtime.completion.CompletionCategory("Audio")

package afield.audio.analysis

import afield.data.collections.s
import afield.data.collections.times
import afield.document.graph.minusAssign
import afield.document.graph.plusAssign
import afield.document.properties.compareTo
import afield.document.workspace.files
import afield.graphics.render.rangeTo
import afield.lang.python.npy
import afield.math.autodiff.compareTo
import afield.math.autodiff.div
import afield.math.autodiff.minus
import afield.math.autodiff.plus
import afield.math.autodiff.times
import afield.math.linalg.remAssign
import afield.math.linalg.set
import afield.math.linalg.times
import afield.math.linalg.x
import afield.math.units.div
import afield.math.units.minus
import afield.math.units.plus
import afield.math.units.times

import afield.audio.synth.*
import afield.data.collections.get
import afield.data.collections.invoke
import afield.document.workspace.absolutePath
import afield.document.workspace.file
import afield.document.workspace.path
import afield.geometry.mesh.contains
import afield.geometry.mesh.normalized
import afield.math.autodiff.minOf
import afield.math.autodiff.product
import afield.math.linalg.Vec2
import afield.math.linalg.Vec3
import afield.math.linalg.i
import afield.math.linalg.invoke
import afield.media.image.normalized
import afield.media.image.split
import kotlinx.serialization.json.Json
import kotlinx.serialization.json.booleanOrNull
import kotlinx.serialization.json.floatOrNull
import kotlinx.serialization.json.intOrNull
import kotlinx.serialization.json.jsonObject
import kotlinx.serialization.json.jsonPrimitive
import org.yaml.snakeyaml.LoaderOptions
import org.yaml.snakeyaml.Yaml
import afield.math.linalg.Vector2f
import afield.math.linalg.Vector3f
import java.io.BufferedInputStream
import java.io.DataInputStream
import java.io.File
import java.nio.ByteBuffer
import java.nio.ByteOrder
import kotlin.math.absoluteValue
import kotlin.math.roundToInt

fun loadFramesSig(path: String, options: FramesSigLoadOptions = FramesSigLoadOptions()): FramesSigDocument =
    loadFramesSig(File(path), options)

private fun FramesSigProduct.toInfo(metadata: FramesSigMetadata): FramesSigProductInfo =
    FramesSigProductInfo(
        path = path,
        frameCount = frameCount,
        vectorWidth = when (this) {
            is FramesSigFloatProduct -> 1
            is FramesSigFloatVectorProduct -> vectorWidth
        },
        startTimeSeconds = metadata.analysisStartTimeSeconds,
        endTimeSeconds = frameTimeAtIndex((frameCount - 1).coerceAtLeast(0), metadata, frameCount),
    )

private fun configureLoaderOptions(file: File, options: FramesSigLoadOptions): LoaderOptions {
    val loaderOptions = LoaderOptions()
    val expandedLimit = (file.length().coerceAtMost(Int.MAX_VALUE.toLong()).toInt() + options.codePointLimitPadding)
        .coerceAtLeast(options.minimumCodePointLimit)

    runCatching {
        LoaderOptions::class.java.getMethod("setCodePointLimit", Int::class.javaPrimitiveType)
            .invoke(loaderOptions, expandedLimit)
    }

    return loaderOptions
}

private fun parseFramesSigMetadata(raw: Map<String, Any?>): FramesSigMetadata {
    val version = raw.stringMap("version")
    val audioProperties = raw.stringMap("audio_properties")
    val analysis = audioProperties.stringMap("analysis")
    val tags = raw.stringMap("tags")

    return FramesSigMetadata(
        raw = raw,
        version = version.entries.associate { (key, value) -> key to value.toString() },
        fileName = tags["file_name"]?.toString(),
        codec = audioProperties["codec"]?.toString(),
        sampleRate = audioProperties["sample_rate"].asIntOrNull(),
        numberChannels = audioProperties["number_channels"].asIntOrNull(),
        bitRate = audioProperties["bit_rate"].asIntOrNull(),
        audioLengthSeconds = audioProperties["length"].asDoubleOrNull(),
        analysisLengthSeconds = analysis["length"].asDoubleOrNull(),
        analysisStartTimeSeconds = analysis["start_time"].asDoubleOrNull() ?: 0.0,
        replayGain = audioProperties["replay_gain"].asDoubleOrNull(),
    )
}

private fun collectProducts(
    path: String,
    value: Any?,
    scalarProducts: MutableMap<String, Double>,
    floatProducts: MutableMap<String, FramesSigFloatProduct>,
    floatVectorProducts: MutableMap<String, FramesSigFloatVectorProduct>,
) {
    when (value) {
        is Map<*, *> -> {
            value.forEach { (childKey, childValue) ->
                collectProducts(
                    path = "$path.${childKey.toString()}",
                    value = childValue,
                    scalarProducts = scalarProducts,
                    floatProducts = floatProducts,
                    floatVectorProducts = floatVectorProducts,
                )
            }
        }

        is Number -> scalarProducts[path] = value.toDouble()

        is List<*> -> {
            val flat = value.asFloatArrayOrNull()
            if (flat != null) {
                floatProducts[path] = FramesSigFloatProduct(path, flat)
                return
            }

            val vectors = value.asFloatVectorListOrNull()
            if (vectors != null) {
                floatVectorProducts[path] = FramesSigFloatVectorProduct(path, vectors)
            }
        }
    }
}

private fun Map<String, Any?>.stringMap(key: String): Map<String, Any?> =
    (this[key] as? Map<*, *>)
        ?.entries
        ?.associate { (childKey, childValue) -> childKey.toString() to childValue }
        ?: emptyMap()

private fun Any?.asDoubleOrNull(): Double? = when (this) {
    is Number -> toDouble()
    is String -> toDoubleOrNull()
    else -> null
}

private fun Any?.asIntOrNull(): Int? = when (this) {
    is Number -> toInt()
    is String -> toIntOrNull()
    else -> null
}

private fun List<*>.asFloatArrayOrNull(): FloatArray? {
    if (isEmpty()) return FloatArray(0)
    val out = FloatArray(size)
    for (i in indices) {
        val value = this[i].asDoubleOrNull() ?: return null
        out[i] = value.toFloat()
    }
    return out
}

private fun List<*>.asFloatVectorListOrNull(): List<FloatArray>? {
    if (isEmpty()) return emptyList()
    val out = ArrayList<FloatArray>(size)
    for (entry in this) {
        val vector = (entry as? List<*>)?.asFloatArrayOrNull() ?: return null
        out += vector
    }
    return out
}

private fun frameIndexAtTime(timeSeconds: Double, metadata: FramesSigMetadata, frameCount: Int): Int {
    if (frameCount <= 1) return 0

    val duration = metadata.analysisLengthSeconds ?: metadata.audioLengthSeconds ?: 0.0
    if (duration <= 0.0) return 0

    val normalized = ((timeSeconds - metadata.analysisStartTimeSeconds) / duration).coerceIn(0.0, 1.0)
    return (normalized * (frameCount - 1)).roundToInt().coerceIn(0, frameCount - 1)
}

private fun frameTimeAtIndex(index: Int, metadata: FramesSigMetadata, frameCount: Int): Double {
    if (frameCount <= 1) return metadata.analysisStartTimeSeconds

    val duration = metadata.analysisLengthSeconds ?: metadata.audioLengthSeconds ?: 0.0
    if (duration <= 0.0) return metadata.analysisStartTimeSeconds

    return metadata.analysisStartTimeSeconds + duration * (index.toDouble() / (frameCount - 1).toDouble())
}

fun loadSound2Vec(dir: File, options: Sound2VecLoadOptions = Sound2VecLoadOptions()): Sound2VecDataset {
    require(dir.isDirectory) { "loadSound2Vec: directory does not exist: ${dir.absolutePath}" }

    val metadataFile = File(dir, "metadata.csv")
    val embeddingsFile = File(dir, "embeddings.npy")
    val manifestFile = File(dir, "manifest.json")

    require(metadataFile.isFile) { "loadSound2Vec: missing metadata.csv in ${dir.absolutePath}" }
    require(embeddingsFile.isFile) { "loadSound2Vec: missing embeddings.npy in ${dir.absolutePath}" }

    val manifest = when {
        manifestFile.isFile -> parseSound2VecManifest(manifestFile)
        options.requireManifest -> error("loadSound2Vec: missing manifest.json in ${dir.absolutePath}")
        else -> null
    }

    val rows = readCsvRows(metadataFile)
    val embeddings = readNpyFloatMatrix(embeddingsFile)

    if (options.requireRowAlignment) {
        require(rows.size == embeddings.size) {
            "loadSound2Vec: metadata row count ${rows.size} does not match embeddings row count ${embeddings.size}"
        }
    }

    if (manifest?.numChunks != null && options.requireRowAlignment) {
        require(manifest.numChunks == rows.size) {
            "loadSound2Vec: manifest num_chunks ${manifest.numChunks} does not match metadata row count ${rows.size}"
        }
    }

    if (manifest?.embeddingDim != null && embeddings.isNotEmpty() && options.requireRowAlignment) {
        require(manifest.embeddingDim == embeddings.first().size) {
            "loadSound2Vec: manifest embedding_dim ${manifest.embeddingDim} does not match embedding width ${embeddings.first().size}"
        }
    }

    val clipCount = minOf(rows.size, embeddings.size)
    val clips = List(clipCount) { index ->
        val row = rows[index]
        val embedding = embeddings[index]
        Sound2VecClip(
            rowIndex = index,
            chunkId = row.int("chunk_id") ?: index,
            filePath = row.string("file_path") ?: error("loadSound2Vec: metadata row $index missing file_path"),
            relPath = row.string("rel_path") ?: "",
            fileName = row.string("file_name") ?: "",
            chunkIndex = row.int("chunk_index") ?: index,
            startSec = row.double("start_sec") ?: 0.0,
            endSec = row.double("end_sec") ?: 0.0,
            durationSec = row.double("duration_sec") ?: 0.0,
            sampleRate = row.int("sample_rate") ?: manifest?.sampleRate ?: 0,
            numSamples = row.int("num_samples") ?: 0,
            tsne2dX = row.float("tsne_2d_x") ?: row.float("x"),
            tsne2dY = row.float("tsne_2d_y") ?: row.float("y"),
            tsne3dX = row.float("tsne_3d_x") ?: row.float("x"),
            tsne3dY = row.float("tsne_3d_y") ?: row.float("y"),
            tsne3dZ = row.float("tsne_3d_z") ?: row.float("z"),
            embedding = embedding,
        )
    }

    return Sound2VecDataset(
        sourceDir = dir,
        manifest = manifest,
        clips = clips,
    )
}

fun loadSound2Vec(path: String, options: Sound2VecLoadOptions = Sound2VecLoadOptions()): Sound2VecDataset =
    loadSound2Vec(File(path), options)

private fun parseSound2VecManifest(file: File): Sound2VecManifest {
    val element = Json.parseToJsonElement(file.readText()).jsonObject
    val raw = element.mapValues { it.value.toString() }

    return Sound2VecManifest(
        raw = raw,
        modelId = element["model_id"]?.jsonPrimitive?.contentOrNullCompat(),
        device = element["device"]?.jsonPrimitive?.contentOrNullCompat(),
        sampleRate = element["sample_rate"]?.jsonPrimitive?.intOrNull,
        clipSeconds = element["clip_seconds"]?.jsonPrimitive?.doubleOrNullCompat(),
        hopSeconds = element["hop_seconds"]?.jsonPrimitive?.doubleOrNullCompat(),
        minSeconds = element["min_seconds"]?.jsonPrimitive?.doubleOrNullCompat(),
        numFiles = element["num_files"]?.jsonPrimitive?.intOrNull,
        numChunks = element["num_chunks"]?.jsonPrimitive?.intOrNull,
        embeddingDim = element["embedding_dim"]?.jsonPrimitive?.intOrNull,
        l2Normalized = element["l2_normalized"]?.jsonPrimitive?.booleanOrNull,
    )
}

private fun readCsvRows(file: File): List<Map<String, String>> {
    val lines = file.readLines().filter { it.isNotBlank() }
    if (lines.isEmpty()) return emptyList()

    val header = parseCsvLine(lines.first())
    return lines.drop(1).map { line ->
        val values = parseCsvLine(line)
        buildMap(header.size) {
            for (index in header.indices) {
                put(header[index], values.getOrElse(index) { "" })
            }
        }
    }
}

private fun parseCsvLine(line: String): List<String> {
    val out = ArrayList<String>()
    val current = StringBuilder()
    var inQuotes = false
    var index = 0

    while (index < line.length) {
        val char = line[index]
        when {
            char == '"' && inQuotes && index + 1 < line.length && line[index + 1] == '"' -> {
                current.append('"')
                index += 2
                continue
            }

            char == '"' -> inQuotes = !inQuotes
            char == ',' && !inQuotes -> {
                out += current.toString()
                current.setLength(0)
            }

            else -> current.append(char)
        }
        index += 1
    }

    out += current.toString()
    return out
}

private fun readNpyFloatMatrix(file: File): List<FloatArray> =
    DataInputStream(BufferedInputStream(file.inputStream())).use { input ->
        val magic = ByteArray(6)
        input.readFully(magic)
        require(magic.contentEquals(byteArrayOf(0x93.toByte(), 'N'.code.toByte(), 'U'.code.toByte(), 'M'.code.toByte(), 'P'.code.toByte(), 'Y'.code.toByte()))) {
            "loadSound2Vec: ${file.absolutePath} is not a valid .npy file"
        }

        val major = input.readUnsignedByte()
        val minor = input.readUnsignedByte()
        val headerLength = when (major) {
            1 -> readLittleEndianShort(input)
            2, 3 -> readLittleEndianInt(input)
            else -> error("loadSound2Vec: unsupported .npy version $major.$minor")
        }

        val headerBytes = ByteArray(headerLength)
        input.readFully(headerBytes)
        val header = headerBytes.toString(Charsets.US_ASCII)

        val descr = Regex("'descr'\\s*:\\s*'([^']+)'").find(header)?.groupValues?.get(1)
            ?: error("loadSound2Vec: could not read dtype from ${file.absolutePath}")
        val fortranOrder = Regex("'fortran_order'\\s*:\\s*(True|False)").find(header)?.groupValues?.get(1) == "True"
        val shapeText = Regex("'shape'\\s*:\\s*\\(([^)]*)\\)").find(header)?.groupValues?.get(1)
            ?: error("loadSound2Vec: could not read shape from ${file.absolutePath}")

        require(!fortranOrder) { "loadSound2Vec: Fortran-ordered arrays are not supported in ${file.absolutePath}" }
        require(descr == "<f4" || descr == "|f4") { "loadSound2Vec: only float32 embeddings are supported, found $descr" }

        val dims = shapeText.split(',')
            .map { it.trim() }
            .filter { it.isNotEmpty() }
            .map { it.toInt() }

        require(dims.isNotEmpty() && dims.size <= 2) {
            "loadSound2Vec: expected a 1D or 2D array in ${file.absolutePath}, found shape (${dims.joinToString()})"
        }

        val rows = if (dims.size == 1) 1 else dims[0]
        val cols = dims.last()
        val payload = input.readBytes()
        val expectedBytes = rows * cols * Float.SIZE_BYTES
        require(payload.size >= expectedBytes) {
            "loadSound2Vec: expected $expectedBytes bytes of embedding data in ${file.absolutePath}, found ${payload.size}"
        }

        val buffer = ByteBuffer.wrap(payload, 0, expectedBytes).order(ByteOrder.LITTLE_ENDIAN)
        List(rows) {
            FloatArray(cols) { buffer.float }
        }
    }

private fun readLittleEndianShort(input: DataInputStream): Int {
    val bytes = ByteArray(2)
    input.readFully(bytes)
    return ByteBuffer.wrap(bytes).order(ByteOrder.LITTLE_ENDIAN).short.toInt() and 0xFFFF
}

private fun readLittleEndianInt(input: DataInputStream): Int {
    val bytes = ByteArray(4)
    input.readFully(bytes)
    return ByteBuffer.wrap(bytes).order(ByteOrder.LITTLE_ENDIAN).int
}

private fun Map<String, String>.string(key: String): String? = this[key]?.takeIf { it.isNotBlank() }

private fun Map<String, String>.int(key: String): Int? = this[key]?.toIntOrNull()

private fun Map<String, String>.float(key: String): Float? = this[key]?.toFloatOrNull()

private fun Map<String, String>.double(key: String): Double? = this[key]?.toDoubleOrNull()

private fun normalizeSound2VecPath(path: String): String = File(path).absoluteFile.normalize().path

private fun kotlinx.serialization.json.JsonPrimitive.contentOrNullCompat(): String? =
    toString()
        .trim()
        .removeSurrounding("\"")
        .takeIf { it != "null" }

private fun kotlinx.serialization.json.JsonPrimitive.doubleOrNullCompat(): Double? =
    contentOrNullCompat()?.toDoubleOrNull()