sergio/humanoid

1package com.rizoma.trazo.data
2
3import android.content.Context
4import androidx.datastore.core.DataStore
5import androidx.datastore.preferences.core.Preferences
6import androidx.datastore.preferences.core.edit
7import androidx.datastore.preferences.core.intPreferencesKey
8import androidx.datastore.preferences.preferencesDataStore
9import kotlinx.coroutines.CoroutineScope
10import kotlinx.coroutines.Dispatchers
11import kotlinx.coroutines.flow.first
12import kotlinx.coroutines.launch
13import java.util.concurrent.ConcurrentHashMap
14
15private val Context.trazoLearnedStore: DataStore<Preferences> by preferencesDataStore(name = "trazo_learned")
16
17/**
18 * Predicción por prefijo, bilingüe (es + en). Dos fuentes:
19 *  - [base]: diccionarios empaquetados (assets), peso fijo por frecuencia.
20 *  - [learned]: palabras que el usuario escribe, contadas y persistidas.
21 *
22 * Mecanismo "buenas sugerencias": una palabra del diccionario se refuerza al instante,
23 * pero una palabra NUEVA solo entra a las sugerencias tras escribirse [PROMOTE] veces,
24 * así los errores escritos una sola vez nunca contaminan. Las sugerencias salen en el
25 * idioma del prefijo que escribes (ambos diccionarios conviven).
26 */
27class SuggestionEngine(context: Context, private val scope: CoroutineScope) {
28
29    private val appContext = context.applicationContext
30    private val store = appContext.trazoLearnedStore
31
32    // Diccionarios separados por idioma para poder activarlos/desactivarlos.
33    private val baseEs = ConcurrentHashMap<String, Int>()
34    private val baseEn = ConcurrentHashMap<String, Int>()
35    private val learned = ConcurrentHashMap<String, Int>()
36    // Bigramas: palabra previa -> {palabra siguiente -> veces}. Alimenta "palabra siguiente".
37    private val bigrams = ConcurrentHashMap<String, ConcurrentHashMap<String, Int>>()
38
39    init {
40        scope.launch(Dispatchers.IO) {
41            loadFile("dict_es.txt", baseEs)
42            loadFile("dict_en.txt", baseEn)
43            loadPersisted()
44        }
45    }
46
47    private fun loadFile(name: String, into: ConcurrentHashMap<String, Int>) {
48        runCatching {
49            appContext.assets.open(name).bufferedReader().useLines { lines ->
50                var i = 0
51                lines.forEach { line ->
52                    val w = line.trim().lowercase()
53                    if (w.isNotEmpty()) {
54                        val weight = maxOf(1, BASE_TOP - i)
55                        into.merge(w, weight) { a, b -> maxOf(a, b) }
56                        i++
57                    }
58                }
59            }
60        }
61    }
62
63    private suspend fun loadPersisted() {
64        runCatching {
65            store.data.first().asMap().forEach { (key, value) ->
66                val name = key.name
67                if (value !is Int) return@forEach
68                when {
69                    name.startsWith(W) -> learned[name.removePrefix(W)] = value
70                    name.startsWith(B) -> {
71                        val rest = name.removePrefix(B)
72                        val sep = rest.indexOf(SEP)
73                        if (sep > 0 && sep < rest.length - 1) {
74                            val prev = rest.substring(0, sep)
75                            val next = rest.substring(sep + 1)
76                            bigrams.getOrPut(prev) { ConcurrentHashMap() }[next] = value
77                        }
78                    }
79                }
80            }
81        }
82    }
83
84    private fun inBase(w: String) = baseEs.containsKey(w) || baseEn.containsKey(w)
85
86    /** Completa el prefijo en los idiomas activos; mayor peso primero, a igual peso la más corta. */
87    fun suggest(prefix: String, useEs: Boolean, useEn: Boolean, limit: Int = 3): List<String> {
88        val p = prefix.trim().lowercase()
89        if (p.isEmpty()) return emptyList()
90
91        val candidates = HashMap<String, Int>()
92        fun scan(map: ConcurrentHashMap<String, Int>) {
93            map.forEach { (w, weight) ->
94                if (w.length > p.length && w.startsWith(p)) {
95                    candidates[w] = maxOf(candidates[w] ?: 0, weight)
96                }
97            }
98        }
99        if (useEs) scan(baseEs)
100        if (useEn) scan(baseEn)
101        learned.forEach { (w, count) ->
102            // Solo palabras nuevas ya "promovidas"; las del diccionario suman refuerzo.
103            if (w.length > p.length && w.startsWith(p) && (count >= PROMOTE || inBase(w))) {
104                candidates[w] = (candidates[w] ?: 0) + count * LEARN_INC
105            }
106        }
107        return candidates.entries
108            .sortedWith(compareByDescending<Map.Entry<String, Int>> { it.value }.thenBy { it.key.length })
109            .map { it.key }
110            .take(limit)
111    }
112
113    /** Palabras más probables tras [prev], según los bigramas aprendidos. */
114    fun nextWords(prev: String, limit: Int = 3): List<String> {
115        val a = clean(prev) ?: return emptyList()
116        val map = bigrams[a] ?: return emptyList()
117        return map.entries
118            // Evita propagar erratas: la siguiente debe ser palabra conocida.
119            .filter { inBase(it.key) || (learned[it.key] ?: 0) >= PROMOTE }
120            .sortedByDescending { it.value }
121            .map { it.key }
122            .take(limit)
123    }
124
125    /** Aprende/refuerza una palabra escrita por el usuario (se promueve sola al llegar a PROMOTE). */
126    fun learn(raw: String) {
127        val w = clean(raw) ?: return
128        learned.merge(w, 1) { a, _ -> a + 1 }
129        scope.launch(Dispatchers.IO) {
130            runCatching {
131                store.edit { prefs ->
132                    val key = intPreferencesKey(W + w)
133                    prefs[key] = (prefs[key] ?: 0) + 1
134                }
135            }
136        }
137    }
138
139    /** Registra que [next] siguió a [prev] (para la predicción de palabra siguiente). */
140    fun learnBigram(prev: String, next: String) {
141        val a = clean(prev) ?: return
142        val b = clean(next) ?: return
143        bigrams.getOrPut(a) { ConcurrentHashMap() }.merge(b, 1) { x, _ -> x + 1 }
144        scope.launch(Dispatchers.IO) {
145            runCatching {
146                store.edit { prefs ->
147                    val key = intPreferencesKey(B + a + SEP + b)
148                    prefs[key] = (prefs[key] ?: 0) + 1
149                }
150            }
151        }
152    }
153
154    private fun clean(raw: String): String? {
155        val w = raw.trim().lowercase()
156        return if (w.length >= 2 && w.all { it.isLetter() }) w else null
157    }
158
159    private companion object {
160        const val BASE_TOP = 5000
161        const val LEARN_INC = 400
162        const val PROMOTE = 2
163        const val W = "w:"
164        const val B = "b:"
165        const val SEP = '>'   // separador prev/next; seguro porque las palabras son solo letras
166    }
167}