#!/usr/bin/env python3
"""Entrena el reconocimiento solo, sin que nadie hable.

    python3 ~/COFRE/CODERS/JARVIS/config/entrenar.py            # tanda normal
    python3 ~/COFRE/CODERS/JARVIS/config/entrenar.py --rapido   # una voz
    python3 ~/COFRE/CODERS/JARVIS/config/entrenar.py --aplicar  # guarda alias

Como funciona: por cada accion del catalogo se dice su frase en voz alta con
Piper, en varias voces y velocidades y con algo de ruido y sala, se transcribe
con el mismo whisper que usa el asistente, y se mira si la transcripcion llega
a la accion correcta. Lo que no llega se apunta como alias candidato.

Por que sirve: los fallos del reconocedor en castellano son sistematicos, no
aleatorios. "abre el correo" se oye "a ver el correo" siempre, y una tanda de
varios cientos de frases los saca casi todos de una sentada.

Que NO sustituye: esto es voz sintetica, no la suya. Sirve para las colisiones
foneticas del idioma —que son la mayoria de los fallos— pero no para su acento,
su microfono ni su habitacion. Para eso esta el registro de lo que no entiende
en el uso real, que se revisa con aprender.py. Las dos cosas se suman.

Un alias solo se propone si no choca con otra accion: si "a ver el correo"
pudiera ser dos cosas distintas, se descarta y se avisa.
"""
import argparse
import importlib.util
import json
import os
import random
import shutil
import subprocess
import sys
import tempfile
import time
from collections import defaultdict

CASA = os.path.expanduser("~")
JARVIS = os.path.join(CASA, "COFRE/CODERS/JARVIS")
CONFIG = os.path.join(CASA, ".var/app/io.github.qwersyk.Newelle/config")
APP = "io.github.qwersyk.Newelle//master"

PIPER = os.path.join(JARVIS, "voz/piper/piper")
MODELOS_VOZ = os.path.join(JARVIS, "voz/modelos")
WHISPER = os.path.join(CONFIG, "models/whisper/whisper.cpp/build/bin/whisper-cli")
MODELOS_STT = os.path.join(CONFIG, "models/whisper/whisper.cpp/models")

CATALOGO = os.path.expanduser(
    "~/.local/share/flatpak/app/io.github.qwersyk.Newelle/x86_64/master/"
    "active/files/share/newelle/acciones/acciones.json")
CATALOGO_FUENTE = os.path.join(JARVIS, "newelle/data/acciones/acciones.json")
PROPIO = os.path.join(CONFIG, "jarvis-acciones.json")
INFORME = os.path.join(JARVIS, "muestras/entrenamiento.json")

MOTOR = os.path.join(JARVIS, "newelle/src/utility/acciones.py")

# Voces distintas para que no aprenda los defectos de una sola, y velocidades
# distintas porque uno no habla igual con prisa que sin ella.
VOCES = ["es_ES-sharvard-medium", "es_ES-davefx-medium", "es_ES-carlfm-x_low"]
VELOCIDADES = [1.0, 0.88, 1.15]

# Un poco de sala y de ruido de fondo: en seco el reconocedor acierta mas de lo
# que acertara nunca en una habitacion de verdad, y entrenar con eso da una
# falsa sensacion de que todo va bien.
RUIDOS = [
    "",                                              # limpio
    "aecho=0.8:0.7:22:0.12",                         # algo de sala
    "volume=0.75,highpass=f=120",                    # mas lejos del microfono
]


def carga_motor():
    spec = importlib.util.spec_from_file_location("acciones", MOTOR)
    m = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(m)
    return m


def dentro(orden, limite=900):
    """Ejecuta algo dentro del flatpak, que es donde vive whisper."""
    return subprocess.run(
        ["flatpak", "run", "--command=bash", APP, "-c", orden],
        capture_output=True, text=True, timeout=limite)


def genera(texto, voz, velocidad, ruido, destino):
    crudo = tempfile.mktemp(suffix=".wav")
    try:
        r = subprocess.run(
            [PIPER, "--model", os.path.join(MODELOS_VOZ, voz + ".onnx"),
             "--length_scale", str(1 / velocidad), "--output_file", crudo],
            input=texto, capture_output=True, text=True, timeout=60)
        if r.returncode != 0 or not os.path.exists(crudo):
            return False
        filtros = "aresample=16000" + ("," + ruido if ruido else "")
        subprocess.run(
            ["ffmpeg", "-hide_banner", "-loglevel", "error", "-i", crudo,
             "-ac", "1", "-af", filtros, "-y", destino],
            check=True, timeout=60)
        return os.path.exists(destino)
    except (OSError, subprocess.SubprocessError):
        return False
    finally:
        if os.path.exists(crudo):
            os.remove(crudo)


# Por tandas y no de una: mil frases son un cuarto de hora, y una sola llamada
# con limite de tiempo se corta a la mitad y se pierde todo lo transcrito. Por
# tandas ademas se puede ver el avance, que en algo que tarda tanto importa.
POR_TANDA = 40


def transcribe_lote(carpeta, ficheros, modelo="base", hilos=8):
    """Transcribe una lista de wav. Devuelve {fichero: texto}."""
    salida = {}
    ini = time.time()
    for i in range(0, len(ficheros), POR_TANDA):
        tanda = ficheros[i:i + POR_TANDA]
        lista = " ".join(f'"{carpeta}/{n}"' for n in tanda)
        orden = (
            f'for f in {lista}; do '
            f'  echo "### $(basename "$f")"; '
            f'  {WHISPER} -m {MODELOS_STT}/ggml-{modelo}.bin -l es -nt -t {hilos} '
            f'    -f "$f" 2>/dev/null | tr "\\n" " "; echo; '
            f'done')
        try:
            r = dentro(orden)
        except subprocess.SubprocessError as e:
            print(f"  tanda {i//POR_TANDA + 1} fallo: {e}", flush=True)
            continue

        actual = None
        for linea in r.stdout.splitlines():
            if linea.startswith("### "):
                actual = linea[4:].strip()
            elif actual:
                salida[actual] = linea.strip()
                actual = None

        hechas = min(i + POR_TANDA, len(ficheros))
        queda = (time.time() - ini) / hechas * (len(ficheros) - hechas)
        print(f"  transcritas {hechas}/{len(ficheros)}  "
              f"({queda/60:.0f} min restantes)", flush=True)
    return salida


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--rapido", action="store_true", help="una voz y una velocidad")
    ap.add_argument("--aplicar", action="store_true", help="guardar los alias encontrados")
    ap.add_argument("--tope", type=int, default=0, help="limitar acciones (para probar)")
    ap.add_argument("--reanudar", action="store_true",
                    help="usar los wav ya generados en vez de rehacerlos")
    args = ap.parse_args()

    for ruta, que in ((PIPER, "piper"), (WHISPER, "whisper")):
        if not os.path.exists(ruta):
            print(f"falta {que}: {ruta}")
            return 1

    catalogo = CATALOGO if os.path.exists(CATALOGO) else CATALOGO_FUENTE
    with open(catalogo) as f:
        acciones = json.load(f)["acciones"]
    if args.tope:
        acciones = acciones[:args.tope]

    voces = VOCES[:1] if args.rapido else VOCES
    velocidades = VELOCIDADES[:1] if args.rapido else VELOCIDADES
    ruidos = RUIDOS[:1] if args.rapido else RUIDOS

    motor = carga_motor()
    cat = motor.Catalogo([catalogo])

    # Las acciones que esperan argumento hay que decirlas CON argumento. Sin el,
    # el motor las rechaza a proposito —"busca el fichero" a secas no es una
    # orden— y contarlo como fallo del reconocedor es medir mal: en la primera
    # tanda hundio el resultado del 48 %, con una cuarta parte de las frases
    # condenadas de antemano.
    RELLENOS = ["notas", "el informe de ayer", "documentos"]

    # Solo la frase principal de cada accion: las variantes ya estan escritas a
    # mano y lo que se busca aqui es como se oye la que uno diria de verdad.
    trabajos = []
    for a in acciones:
        if not a["frases"]:
            continue
        frase = a["frases"][0]
        if a.get("captura"):
            frase = frase + " " + random.choice(RELLENOS)
        for voz in voces:
            for vel in velocidades:
                trabajos.append((a["id"], frase, voz, vel,
                                 random.choice(ruidos)))

    total = len(trabajos)
    print(f"{len(acciones)} acciones · {len(voces)} voces · "
          f"{len(velocidades)} velocidades = {total} frases habladas")
    print("esto tarda; se puede dejar corriendo\n")

    carpeta = os.path.join(CONFIG, "entrenamiento")

    indice = {f"{n:05d}.wav": (ident, frase, voz, vel)
              for n, (ident, frase, voz, vel, _) in enumerate(trabajos, 1)}

    if args.reanudar and os.path.isdir(carpeta):
        indice = {n: v for n, v in indice.items()
                  if os.path.exists(os.path.join(carpeta, n))}
        print(f"reanudando con {len(indice)} frases ya generadas\n", flush=True)
    else:
        shutil.rmtree(carpeta, ignore_errors=True)
        os.makedirs(carpeta, exist_ok=True)
        vivos, ini = {}, time.time()
        for n, (ident, frase, voz, vel, ruido) in enumerate(trabajos, 1):
            nombre = f"{n:05d}.wav"
            if genera(frase, voz, vel, ruido, os.path.join(carpeta, nombre)):
                vivos[nombre] = (ident, frase, voz, vel)
            if n % 25 == 0 or n == total:
                queda = (time.time() - ini) / n * (total - n)
                print(f"  generadas {n}/{total}  ({queda/60:.0f} min restantes)",
                      flush=True)
        indice = vivos

    print(f"\ntranscribiendo {len(indice)} frases...", flush=True)
    textos = transcribe_lote(carpeta, sorted(indice))
    print(f"transcritas {len(textos)}\n", flush=True)

    aciertos = 0
    candidatos = defaultdict(lambda: defaultdict(int))
    for nombre, (ident, frase, voz, vel) in indice.items():
        oido = textos.get(nombre, "").strip()
        if not oido:
            continue
        encontrada, _ = cat.busca(oido)
        if encontrada is not None and encontrada.id == ident:
            aciertos += 1
        else:
            candidatos[ident][motor.normaliza(oido)] += 1

    evaluadas = sum(1 for n in indice if textos.get(n))
    print(f"aciertos: {aciertos}/{evaluadas} "
          f"({100*aciertos/max(1,evaluadas):.0f} %)\n")

    # Un alias solo vale si no puede ser otra cosa. Si la misma frase mal oida
    # aparece en dos acciones, elegir una seria peor que no hacer nada.
    duenos = defaultdict(set)
    for ident, frases in candidatos.items():
        for f in frases:
            duenos[f].add(ident)

    propuestos, descartados = defaultdict(list), []
    for ident, frases in candidatos.items():
        for f, veces in sorted(frases.items(), key=lambda x: -x[1]):
            if len(duenos[f]) > 1:
                descartados.append((f, sorted(duenos[f])))
                continue
            if len(f.split()) < 2:          # una palabra suelta engancha de todo
                continue
            # Si lo oido no se parece a lo dicho, es que el reconocedor se
            # perdio del todo, no que haya una confusion fonetica estable.
            # Guardarlo como alias meteria una frase arbitraria en el catalogo
            # y algun dia enganchara una conversacion.
            original = next((a["frases"][0] for a in acciones if a["id"] == ident), "")
            if motor.parecido(f, motor.normaliza(original)) < 0.55:
                descartados.append((f, [ident + " (irreconocible)"]))
                continue
            if cat.busca(f)[0] is not None:  # ya lo coge otra cosa
                continue
            propuestos[ident].append((f, veces))

    if propuestos:
        print("alias que faltan (frase mal oida -> accion):")
        for ident, lista in sorted(propuestos.items()):
            for f, veces in lista[:3]:
                print(f'  "{f}"  ->  {ident}  ({veces}x)')
    else:
        print("no falta ningun alias: todo lo dicho llega a su accion")

    if descartados:
        print(f"\n{len(descartados)} descartados por ambiguos "
              f"(la misma frase encaja en dos acciones):")
        for f, ids in descartados[:5]:
            print(f'  "{f}" -> {", ".join(ids)}')

    os.makedirs(os.path.dirname(INFORME), exist_ok=True)
    with open(INFORME, "w") as f:
        json.dump({
            "frases": evaluadas, "aciertos": aciertos,
            "propuestos": {k: v for k, v in propuestos.items()},
            "ambiguos": descartados,
        }, f, indent=2, ensure_ascii=False)
    print(f"\ninforme en {INFORME}")

    if args.aplicar and propuestos:
        guarda(propuestos, catalogo)
    elif propuestos:
        print("para guardarlos: vuelve a lanzarlo con --aplicar")

    print(f"los wav se quedan en {carpeta} por si hay que repasar algo")
    return 0


def guarda(propuestos, catalogo):
    with open(catalogo) as f:
        original = {a["id"]: a for a in json.load(f)["acciones"]}
    try:
        with open(PROPIO) as f:
            propio = json.load(f)
    except (OSError, json.JSONDecodeError):
        propio = {"version": 1, "acciones": []}

    porid = {a["id"]: a for a in propio.get("acciones", [])}
    n = 0
    for ident, lista in propuestos.items():
        entrada = porid.get(ident)
        if entrada is None:
            entrada = dict(original[ident])
            porid[ident] = entrada
            propio.setdefault("acciones", []).append(entrada)
        for frase, _ in lista:
            if frase not in entrada["frases"]:
                entrada["frases"].append(frase)
                n += 1

    with open(PROPIO, "w") as f:
        json.dump(propio, f, indent=2, ensure_ascii=False)
    print(f"{n} alias guardados en {PROPIO}")
    print("se aplican al abrir la pestaña JARVIS otra vez")


if __name__ == "__main__":
    sys.exit(main())