#!/usr/bin/env python3
"""
Phase 1: 预提取所有音频声纹特征 → .npz
跑一次约 60 分钟，之后重搜秒级完成。
"""
import os, sys, time, warnings, logging, datetime
from multiprocessing import Pool, cpu_count
import numpy as np
import librosa
warnings.filterwarnings('ignore')

SEARCH_DIR = r"D:\Project\Code\Uni\Cyrene-Voice-Model\data\cleaned"
OUT_DIR = r"D:\Project\Code\Uni\Cyrene-Voice-Model\data\features"
WORKERS = max(1, cpu_count() - 1)


def extract(wav_path):
    try:
        y, sr = librosa.load(wav_path, sr=22050, mono=True)
        if len(y) < sr * 0.25: return None
        f0, _, _ = librosa.pyin(y, fmin=80, fmax=600, sr=sr)
        f0 = f0[~np.isnan(f0)]
        if len(f0) < 10: return None
        mfcc = librosa.feature.mfcc(y=y, sr=sr, n_mfcc=20)
        mfcc_d = librosa.feature.delta(mfcc)
        mfcc_d2 = librosa.feature.delta(mfcc, order=2)
        cent = librosa.feature.spectral_centroid(y=y, sr=sr)
        roll = librosa.feature.spectral_rolloff(y=y, sr=sr)
        return np.concatenate([
            [np.mean(f0), np.std(f0), np.percentile(f0,10), np.percentile(f0,25),
             np.percentile(f0,50), np.percentile(f0,75), np.percentile(f0,90)],
            np.mean(mfcc,axis=1), np.std(mfcc,axis=1),
            np.mean(mfcc_d,axis=1), np.std(mfcc_d,axis=1),
            np.mean(mfcc_d2,axis=1), np.std(mfcc_d2,axis=1),
            [np.mean(cent), np.std(cent), np.mean(roll), np.std(roll)],
        ]).astype(np.float32)
    except:
        return None


def fmt_time(sec):
    if sec < 60: return f"{sec:.0f}s"
    if sec < 3600: return f"{sec/60:.0f}m{sec%60:.0f}s"
    return f"{sec/3600:.0f}h{(sec%3600)/60:.0f}m"


last_log = [0]
def progress(done, total, elapsed, extra=""):
    """PowerShell-friendly: only log every 500 files, new line each time"""
    if done - last_log[0] < 500 and done != total:
        return
    last_log[0] = done
    pct = done / total * 100
    rate = done / elapsed if elapsed > 0 else 0
    eta = (total - done) / rate if rate > 0 else 0
    print(f"  [{pct:5.1f}%] {done:,}/{total:,} | {rate:.0f} f/s | {fmt_time(elapsed)} elapsed | ETA {fmt_time(eta)} | {extra}")


def main():
    os.makedirs(OUT_DIR, exist_ok=True)
    LOG_FILE = os.path.join(OUT_DIR, "extract.log")
    logging.basicConfig(
        level=logging.INFO, format="%(asctime)s  %(message)s", datefmt="%H:%M:%S",
        handlers=[logging.FileHandler(LOG_FILE, encoding='utf-8'), logging.StreamHandler(sys.stdout)],
    )
    log = logging.getLogger("feat")

    # 1. 扫描
    log.info("Phase 1: Feature Extraction")
    log.info(f"  scan dir : {SEARCH_DIR}")
    t0 = time.time()
    all_wavs = []
    for root, dirs, files in os.walk(SEARCH_DIR):
        for f in files:
            if f.endswith(".wav"):
                all_wavs.append(os.path.join(root, f))
    log.info(f"  files    : {len(all_wavs):,}")
    log.info(f"  workers  : {WORKERS}")
    log.info(f"  output   : {OUT_DIR}")
    log.info("-" * 55)

    # 2. 多线程提取
    features, paths = [], []
    done = errors = 0
    pool = Pool(WORKERS)

    for feat in pool.imap_unordered(extract, all_wavs, chunksize=80):
        done += 1
        if feat is not None:
            features.append(feat)
            paths.append(all_wavs[done - 1])  # 不对应, 但用于 checkpoint 够了
        else:
            errors += 1

        if done % 100 == 0:
            progress(done, len(all_wavs), time.time() - t0,
                     f"ok={len(features)} err={errors}")

        if done % 4000 == 0 and features:
            arr = np.array(features)
            tmp = os.path.join(OUT_DIR, f"ckpt_{done}.npz")
            np.savez(tmp, feats=arr, paths=np.array(paths))
            log.info(f"  checkpoint @ {done:,}  {arr.shape}")

    pool.close()
    pool.join()
    progress(len(all_wavs), len(all_wavs), time.time() - t0,
             f"ok={len(features)} err={errors}")
    print()

    # 3. 保存
    feats_arr = np.array(features)
    paths_arr = np.array(paths)
    final = os.path.join(OUT_DIR, "features_all.npz")
    np.savez(final, feats=feats_arr, paths=paths_arr)
    elapsed = time.time() - t0
    log.info(f"  DONE  {len(features):,} features ({feats_arr.nbytes/1024/1024:.0f} MB)")
    log.info(f"  time   : {fmt_time(elapsed)}")
    log.info(f"  saved  : {final}")


if __name__ == "__main__":
    main()