"""
Try the model on your own voice, any time during training.

    python listen.py --record 10                 # record 10 seconds from the microphone
    python listen.py my_recitation.mp3           # or a recording (mp3, wav, flac, ogg)
    python listen.py --record 10 --ayah 2:255    # compare with an ayah (Warsh text by default)

Uses the best checkpoint so far (runs/<name>/best.pt, the lowest validation error) on the CPU, so
it doesn't slow down training on the GPU. Without --ayah it shows the closest ayahs. Recording
needs `pip install sounddevice`; saving each recording to my_recordings/ lets you replay the
same audio on later checkpoints (python listen.py my_recordings/<file>.wav).
"""
import argparse
import os
import sys
import time
from types import SimpleNamespace

import numpy as np
import soundfile as sf
import torch

from quran_asr import audio, config as config_mod, features, model as model_mod, quran, text

HERE = os.path.dirname(os.path.abspath(__file__))


def load_model(cfg, checkpoint):
    path = checkpoint or os.path.join(cfg.paths.run_dir, "best.pt")
    if not os.path.exists(path):
        path = os.path.join(cfg.paths.run_dir, "last.pt")
    for _ in range(5):  # train.py may be writing it right now
        try:
            ckpt = torch.load(path, map_location="cpu", weights_only=False)
            break
        except (RuntimeError, EOFError, OSError):
            time.sleep(1)
    else:
        sys.exit(f"could not read {path}")
    model = model_mod.build(SimpleNamespace(**ckpt["model_config"]), len(text.VOCAB))
    model.load_state_dict(ckpt["model"])
    model.eval()
    cer = f", validation letter error {ckpt['valid_cer'] * 100:.1f}%" if "valid_cer" in ckpt else ""
    print(f"model: {os.path.relpath(path, HERE)} (step {ckpt.get('step', '?')}{cer})")
    return model


def record(seconds):
    try:
        import sounddevice as sd
    except ImportError:
        sys.exit("recording needs: pip install sounddevice   (or pass an audio file instead)")
    for k in (3, 2, 1):
        print(f"recording in {k}...", end="\r", flush=True)
        time.sleep(1)
    print(f"recording {seconds} s: recite now", flush=True)
    wave = sd.rec(int(seconds * audio.SAMPLE_RATE), samplerate=audio.SAMPLE_RATE, channels=1, dtype="float32")
    sd.wait()
    wave = wave[:, 0]
    os.makedirs(os.path.join(HERE, "my_recordings"), exist_ok=True)
    path = os.path.join(HERE, "my_recordings", time.strftime("%Y%m%d-%H%M%S") + ".wav")
    sf.write(path, wave, audio.SAMPLE_RATE)
    print(f"saved {os.path.relpath(path, HERE)}  (peak level {np.abs(wave).max():.2f}"
          + (": very quiet, move closer to the microphone)" if np.abs(wave).max() < 0.05 else ")"))
    return wave


def transcribe(model, wave):
    feats = features.fbank(torch.from_numpy(np.ascontiguousarray(wave, dtype=np.float32)))
    with torch.no_grad():
        log_probs, lens = model(feats[None], torch.tensor([feats.shape[0]]))
    return text.decode(model_mod.greedy_decode(log_probs, lens)[0])


def word_diff(heard, ref):
    """The reference with words the model got wrong in [brackets] and extra heard words in {braces}."""
    import difflib

    h, r = heard.split(), ref.split()
    out = []
    for op, i1, i2, j1, j2 in difflib.SequenceMatcher(None, r, h, autojunk=False).get_opcodes():
        if op == "equal":
            out += r[i1:i2]
        else:
            out += [f"[{w}]" for w in r[i1:i2]] + [f"{{{w}}}" for w in h[j1:j2]]
    return " ".join(out)


def closest(heard, texts, n=3):
    """Ayahs whose text best matches what was heard: a quick letter-pair filter, then edit distance."""
    def pairs(s):
        s = s.replace(" ", "")
        return {s[i:i + 2] for i in range(len(s) - 1)}

    hp = pairs(heard)
    if not hp:
        return []
    shortlist = sorted(texts.items(), key=lambda kv: -len(hp & pairs(kv[1])) / (len(hp | pairs(kv[1])) or 1))[:40]
    scored = []
    for key, ref in shortlist:
        err, total = text.error_counts(heard, ref)[:2]
        scored.append((err / max(1, total), key, ref))
    return sorted(scored)[:n]


def main():
    parser = config_mod.from_args(__doc__)
    parser.set_defaults(config=os.path.join(HERE, "config", "full.yaml"))
    parser.add_argument("file", nargs="?", help="a recording to transcribe")
    parser.add_argument("--record", type=float, metavar="SECONDS", help="record from the microphone")
    parser.add_argument("--ayah", help="what you recited, e.g. 2:255, to compare with")
    parser.add_argument("--riwaya", choices=["warsh", "hafs"], default="warsh", help="reference text (default warsh)")
    parser.add_argument("--checkpoint", help="a specific checkpoint instead of best.pt")
    args = parser.parse_args()
    if not args.file and not args.record:
        parser.error("give a recording or --record SECONDS")
    if hasattr(sys.stdout, "reconfigure"):
        sys.stdout.reconfigure(encoding="utf-8")
    torch.set_num_threads(4)  # leave the rest of the CPU to training

    cfg = config_mod.load(args.config)
    model = load_model(cfg, args.checkpoint)
    wave = record(args.record) if args.record else audio.load(args.file)

    started = time.time()
    heard = transcribe(model, wave)
    print(f"\nheard ({len(wave) / audio.SAMPLE_RATE:.1f} s of audio in {time.time() - started:.1f} s):\n  {heard or '(nothing)'}")

    if args.riwaya == "hafs":
        hafs = quran.load_hafs(cfg.paths.data_dir)
        texts = {k: quran.hafs_label(hafs, *k) for k in hafs}
    else:
        texts = quran.warsh_by_ayah(quran.load_warsh(cfg.paths.warsh_text))

    if args.ayah:
        sura, aya = (int(v) for v in args.ayah.split(":"))
        ref = texts.get((sura, aya))
        if not ref:
            sys.exit(f"no ayah {args.ayah} in the {args.riwaya} text")
        counts = text.error_counts(heard, ref)
        print(f"\n{args.riwaya} {sura}:{aya}:\n  {ref}")
        print(f"\nletter error {counts[0] / max(1, counts[1]) * 100:.1f}%   word error {counts[2] / max(1, counts[3]) * 100:.1f}%"
              f"   skeleton error {counts[4] / max(1, counts[5]) * 100:.1f}%")
        print(f"[missed or wrong]  {{extra}}:\n  {word_diff(heard, ref)}")
    elif heard:
        print(f"\nclosest ayahs ({args.riwaya}):")
        for err, (sura, aya), ref in closest(heard, texts):
            print(f"  {sura}:{aya}  letter error {err * 100:.0f}%   {ref[:90]}")


if __name__ == "__main__":
    main()