#!/usr/bin/env python3
"""对比 A/B 两版的双耳空间分离度。

指标: 左右声道去均值后的相关系数。
  接近 1.0 = 两耳听到的几乎一样 -> 声音"钉在头中间"(普通耳机听 5.1 的听感)
  明显偏低 = 两耳有差异     -> 携带了方位信息(HRTF 起了作用)
"""
import os
import struct
import wave


def read_stereo(path: str) -> tuple[list[float], list[float]]:
    """读 wav 的前两个声道。"""
    with wave.open(path) as w:
        channels = w.getnchannels()
        raw = w.readframes(w.getnframes())
    frames = struct.unpack("<%dh" % (len(raw) // 2), raw)
    left: list[float] = []
    right: list[float] = []
    for i in range(0, len(frames) - channels + 1, channels):
        left.append(float(frames[i]))
        right.append(float(frames[i + 1]) if channels > 1 else float(frames[i]))
    return left, right


def centered(xs: list[float]) -> list[float]:
    """去均值。"""
    if not xs:
        return []
    mean = sum(xs) / len(xs)
    return [x - mean for x in xs]


def correlation(a: list[float], b: list[float]) -> float:
    """皮尔逊相关系数。"""
    num = sum(x * y for x, y in zip(a, b))
    da = sum(x * x for x in a) ** 0.5
    db = sum(y * y for y in b) ** 0.5
    return num / (da * db) if da and db else 0.0


def rms(xs: list[float]) -> float:
    """均方根。"""
    return (sum(x * x for x in xs) / len(xs)) ** 0.5 if xs else 0.0


def main() -> int:
    os.chdir(os.path.dirname(os.path.abspath(__file__)))
    print("%-22s %10s %10s %10s" % ("文件", "相关性", "L 均方根", "R 均方根"))
    print("-" * 58)
    for name in ("A-普通下混.wav", "B-HRTF环绕.wav"):
        if not os.path.isfile(name):
            continue
        left, right = read_stereo(name)
        print("%-22s %10.3f %10.1f %10.1f"
              % (name, correlation(centered(left), centered(right)), rms(left), rms(right)))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())