import os, sys, json, time, hashlib
from pathlib import Path
os.environ['HF_HUB_OFFLINE']='1'
os.environ['TRANSFORMERS_OFFLINE']='1'
os.environ['TOKENIZERS_PARALLELISM']='false'
import numpy as np
import mlx.core as mx
from scipy.io import wavfile
from scipy.signal import stft
from mlx_audio.tts.utils import load_model

ROOT=Path(os.environ.get('RESULTS_DIR', '.'))
BASE=Path(os.environ['MODEL_ROOT'])
variant=sys.argv[1]
paths={
 'original':Path(os.environ['ORIGINAL_MODEL']),
 'affine6':Path(os.environ['AFFINE6_MODEL']),
}
path=paths.get(variant,BASE/variant)
mx.set_memory_limit(8*2**30);mx.set_cache_limit(128*2**20)
model=load_model(str(path));model.eval();mx.eval(model.parameters());mx.synchronize()
assert model.tokenizer is not None and model._conds is not None
texts=[
 'This is the oMLX Chatterbox Turbo checkpoint. It was quantized from the original fp16 model.',
 'Hello mate, the training checkpoint is saved, and this is Chatterbox speaking locally on the Mac.',
 'Please check the appointment on Thursday at three thirty. Bring the blue folder, your keys, and a bottle of water.',
]
def log_mel(audio, sr):
    _, _, z=stft(audio,fs=sr,nperseg=1024,noverlap=768,boundary='zeros')
    power=np.abs(z)**2
    mel=lambda hz:2595*np.log10(1+hz/700)
    hz=lambda m:700*(10**(m/2595)-1)
    edges=hz(np.linspace(mel(40),mel(12000),82))
    freqs=np.linspace(0,sr/2,power.shape[0]);bank=np.zeros((80,power.shape[0]))
    for j in range(80):
        bank[j]=np.maximum(0,np.minimum((freqs-edges[j])/(edges[j+1]-edges[j]),(edges[j+2]-freqs)/(edges[j+2]-edges[j+1])))
    return bank@power

rows=[]
for i,text in enumerate(texts):
    tokens_file=ROOT/f'original-speech-tokens-{i}.npy'
    if variant=='original':
        capture={};original=model.s3gen.inference
        def traced(*args,**kwargs):
            capture['tokens']=np.array(kwargs['speech_tokens'])
            return original(*args,**kwargs)
        model.s3gen.inference=traced
        mx.random.seed(20261009);np.random.seed(20261009)
        for result in model.generate(text,temperature=.8,top_p=.95,top_k=1000,repetition_penalty=1.2,split_pattern=None,max_tokens=400):mx.eval(result.audio)
        model.s3gen.inference=original
        np.save(tokens_file,capture['tokens'])
    tokens=np.load(tokens_file)
    mx.random.seed(9182);mx.synchronize();start=time.perf_counter()
    audio,_=model.s3gen.inference(speech_tokens=mx.array(tokens),ref_dict=model._conds.gen,n_cfm_timesteps=2)
    mx.eval(audio);mx.synchronize();elapsed=time.perf_counter()-start
    audio=np.array(audio).reshape(-1).astype(np.float32)
    assert audio.size and np.isfinite(audio).all()
    np.save(ROOT/f'{variant}-decoder-{i}.npy',audio)
    wavfile.write(ROOT/f'{variant}-decoder-{i}.wav',24000,(np.clip(audio,-1,1)*32767).astype(np.int16))
    reference=np.load(ROOT/f'original-decoder-{i}.npy')
    assert audio.shape==reference.shape,(audio.shape,reference.shape)
    a=log_mel(reference,24000);b=log_mel(audio,24000)
    floor=max(float(a.max())*1e-6,1e-12)
    active=(a>=max(float(a.max())*1e-4,1e-12))
    delta=10*np.log10(np.maximum(a,floor))-10*np.log10(np.maximum(b,floor))
    rms_delta=float(np.sqrt(np.mean(delta[active]**2)))
    row={'text':text,'speech_tokens':int(tokens.size),'audio_seconds':audio.size/24000,'decoder_seconds':elapsed,'log_mel_rms_difference_db':rms_delta,'peak_amplitude':float(np.abs(audio).max()),'clipped_fraction':float(np.mean(np.abs(audio)>=1))}
    rows.append(row);print('DECODER',variant,json.dumps(row),flush=True)
report={'variant':variant,'method':'Identical original speech tokens, shared bundled conditioning, identical decoder RNG seed. RMS log-mel difference over reference bins above -40 dB, floor -60 dB. Lower means spectral output closer to original; not a subjective quality score.','decoder_seed':9182,'clips':rows,'mean_log_mel_rms_difference_db':float(np.mean([r['log_mel_rms_difference_db'] for r in rows]))}
(ROOT/f'{variant}-decoder-fidelity.json').write_text(json.dumps(report,indent=2)+'\n')