KokoroPhonemeConverter.kt

import android.content.Context
import android.content.res.Resources
import android.util.Log
import com.github.medavox.ipa_transcribers.Language
import java.io.IOException

class KokoroPhonemeConverter(context: Context) {
    private val phonemeMap = mutableMapOf<String, String>()

    init {
        loadDictionary(context.applicationContext)
    }

    fun phonemize(text: String, lang: String = "en-us", normalize: Boolean = true): String {
        val normalizedText = if (normalize) normalizeText(text) else text
        val wordsAndPunctuation = normalizedText
            .split(Regex("(?<=\\W)|(?=\\W)"))
            .filter { it.isNotBlank() }

        val phonemes = StringBuilder()
        for ((index, word) in wordsAndPunctuation.withIndex()) {
            val ipaPhonemes = if (word.matches(nonWordRegex)) {
                word
            } else {
                val compactPhonemes = convertToPhonemes(word)
                    .replace(" ", "")
                    .replace("ˌ", "")
                adjustStressMarkers(compactPhonemes)
            }

            if (index > 0 && !word.matches(nonWordRegex)) {
                phonemes.append(" ")
            }
            phonemes.append(ipaPhonemes)
        }

        return postProcessPhonemes(phonemes.toString(), lang)
    }

    private fun loadDictionary(context: Context) {
        try {
            context.resources.openRawResource(R.raw.cmudict_ipa).bufferedReader()
                .useLines { lines ->
                    lines.filter { !it.startsWith(";;;") }.forEach { line ->
                        val parts = line.split("\t", limit = 2)
                        if (parts.size == 2) {
                            phonemeMap[parts[0]] = parts[1]
                        }
                    }
                }
        } catch (e: IOException) {
            Log.e(TAG, "Error loading CMU IPA dictionary", e)
        } catch (e: Resources.NotFoundException) {
            Log.e(TAG, "CMU IPA dictionary not found", e)
        }
    }

    private fun convertToPhonemes(word: String): String {
        if (word.matches(nonWordRegex)) {
            return word
        }

        val cleanWord = word.replace(Regex("[^a-zA-Z']"), "").uppercase()
        val dictionaryKey = cleanWord.replace(Regex("[0-9]"), "ˈ")
        return phonemeMap[dictionaryKey]
            ?.split(",")
            ?.firstOrNull()
            ?.trim()
            ?: Language.ENGLISH.transcriber.transcribe(word)
    }

    private fun adjustStressMarkers(input: String): String {
        val builder = StringBuilder(input)
        var i = 0

        while (i < builder.length) {
            if (builder[i] == 'ˈ' || builder[i] == 'ˌ') {
                val stressIndex = i
                val stressChar = builder[i]

                for (j in stressIndex + 1 until builder.length) {
                    if (builder[j] in vowels) {
                        builder.deleteCharAt(stressIndex)
                        builder.insert(j - 1, stressChar)
                        i = j
                        break
                    }
                }
            }
            i++
        }

        return builder.toString()
    }

    private fun normalizeText(text: String): String {
        var normalizedText = text
            .lines()
            .joinToString("\n") { it.trim() }
            .replace("[‘’]".toRegex(), "'")
            .replace("[“”«»]".toRegex(), "\"")
            .replace("[、。！，：；？]".toRegex()) { match ->
                when (match.value) {
                    "、" -> ","
                    "。" -> "."
                    "！" -> "!"
                    "，" -> ","
                    "：" -> ":"
                    "；" -> ";"
                    "？" -> "?"
                    else -> match.value
                } + " "
            }

        normalizedText = normalizedText
            .replace(Regex("\\bD[Rr]\\.(?= [A-Z])"), "Doctor")
            .replace(Regex("\\b(?:Mr\\.|MR\\.(?= [A-Z]))"), "Mister")
            .replace(Regex("\\b(?:Ms\\.|MS\\.(?= [A-Z]))"), "Miss")
            .replace(Regex("\\b(?:Mrs\\.|MRS\\.(?= [A-Z]))"), "Mrs")
            .replace(Regex("\\betc\\.(?! [A-Z])"), "etc")
            .replace(Regex("(?<=\\d),(?=\\d)"), "")
            .replace(Regex("(?<=\\d)-(?=\\d)"), " to ")

        return KokoroNumberNormalizer.expandNumbers(normalizedText).trim()
    }

    private fun postProcessPhonemes(phonemes: String, lang: String): String {
        var result = phonemes
            .replace("r", "ɹ")
            .replace("x", "k")
            .replace("ʲ", "j")
            .replace("ɬ", "l")
            .replace("kəkˈoːɹoʊ", "kˈoʊkəɹoʊ")
            .replace("kəkˈɔːɹəʊ", "kˈəʊkəɹəʊ")

        if (lang == "en-us") {
            result = result.replace("ti", "di")
        }

        return result
            .map { char -> if (char in vocab) char else ' ' }
            .joinToString(separator = "")
            .replace(Regex("\\s+"), " ")
            .trim()
    }

    private companion object {
        private const val TAG = "KokoroPhonemes"
        private val nonWordRegex = Regex("[^a-zA-Z']+")
        private val vowels = setOf(
            'a', 'e', 'i', 'o', 'u', 'ɑ', 'ɐ', 'ɒ', 'æ', 'ɔ', 'ə', 'ɘ', 'ɚ', 'ɛ',
            'ɜ', 'ɝ', 'ɞ', 'ɪ', 'ɨ', 'ø', 'ɵ', 'œ', 'ɶ', 'ʉ', 'ʊ', 'ʌ', 'A', 'E',
            'I', 'O', 'U', 'ː', 'ˑ'
        )
        private val vocab = run {
            val pad = '$'
            val punctuation = ";:,.!?¡¿—…\"«»“” "
            val letters = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
            val lettersIpa =
                "ɑɐɒæɓʙβɔɕçɗɖðʤəɘɚɛɜɝɞɟʄɡɠɢʛɦɧħɥʜɨɪʝɭɬɫɮʟɱɯɰŋɳɲɴøɵɸθœɶʘɹɺɾɻʀʁɽʂʃʈʧʉʊʋⱱʌɣɤʍχʎʏʑʐʒʔʡʕʢǀǁǂǃˈˌːˑʼʴʰʱʲʷˠˤ˞↓↑→↗↘'̩'ᵻ"

            (listOf(pad) + punctuation.toList() + letters.toList() + lettersIpa.toList()).toSet()
        }
    }
}