"""Recreate the public JVS bandwidth variants and Koenami measurements.

Requires a prepared Koenami checkout, its native engine and JVS samples.
Does not download or redistribute audio. Refuses an existing output directory.
"""
import argparse, csv, hashlib, json, math, os, subprocess
from pathlib import Path
import numpy as np
import soundfile as sf
from scipy import signal

parser=argparse.ArgumentParser(description=__doc__)
parser.add_argument('--app',type=Path,required=True,help='Prepared Koenami repository root')
parser.add_argument('--csv',type=Path,default=Path(__file__).with_name('measurements.csv'))
parser.add_argument('--out',type=Path,required=True)
parser.add_argument('--engine',type=Path)
parser.add_argument('--limit',type=int,default=0,help='Optional smoke test: first N speakers')
args=parser.parse_args()
assert not args.out.exists(), 'Use a fresh output directory'
engine=args.engine or args.app/'measure/target/release/koenami-measure.exe'
if not engine.is_file():
    engine=args.engine or args.app/'measure/target/release/koenami-measure'
assert engine.is_file(), 'Build or specify the native engine'
raw=(args.app/'data/native-ja.json').read_bytes()
try: library=json.loads(raw.decode('utf-8-sig'))
except UnicodeDecodeError: library=json.loads(raw.decode('cp932'))
clips={r['id']:r for r in library['clips']}
with args.csv.open(encoding='utf-8-sig',newline='') as f:
    reps=[r for r in csv.DictReader(f) if r['condition']=='original']
assert len(reps)==98 and len({r['speaker'] for r in reps})==98
assert all(r['group'] in ('female','male') and r['id'].startswith('jvs') for r in reps)
if args.limit:reps=reps[:args.limit]
cases=[]
for r in reps:
    source=args.app/'data/samples'/Path(clips[r['id']]['audio']).name
    assert source.is_file(), f'Missing public corpus sample: {r["id"]}'
    cases.append((r,source))
args.out.mkdir(parents=True)
audio=args.out/'audio';audio.mkdir()
sos=signal.butter(8,[300,3400],btype='bandpass',fs=16000,output='sos')
records=[];files=[]
for r,source in cases:
    x,sr=sf.read(source,dtype='float64')
    if x.ndim==2:x=x.mean(axis=1)
    if sr!=16000:
        d=math.gcd(sr,16000);x=signal.resample_poly(x,16000//d,sr//d)
    variants=dict(original=x,
        roundtrip_8k=signal.resample_poly(signal.resample_poly(x,1,2),2,1)[:len(x)],
        bandpass_300_3400=signal.sosfiltfilt(sos,x))
    for condition,y in variants.items():
        path=audio/f'{r["id"]}__{condition}.wav'
        assert len(y)==len(x)
        sf.write(path,y,16000,subtype='FLOAT')
        run=subprocess.run([str(engine.resolve()),str(path.resolve())],check=True,
            capture_output=True,env={**os.environ,'RAYON_NUM_THREADS':'4'})
        result=json.loads(run.stdout.decode('utf-8'))['measurement']
        assert result['version']=='4.1.0', 'This experiment requires engine 4.1.0'
        records.append(dict(id=r['id'],speaker=r['speaker'],group=r['group'],
            condition=condition,measurement=result))
        files.append(dict(file=path.name,sha256=hashlib.sha256(path.read_bytes()).hexdigest()))
    print('Measured',r['id'],flush=True)
with (args.out/'measurements.json').open('x',encoding='utf-8') as f:
    json.dump(records,f,ensure_ascii=False,indent=2,allow_nan=False)
with (args.out/'audio-hashes.json').open('x',encoding='utf-8') as f:
    json.dump(files,f,indent=2)
print('Completed',len(records),'measurements')