#!/usr/bin/env python3
"""
Per-channel correlation between two PCM WAV files.

Finds best lag, then (by default) a 1:1 channel mapping that maximizes
Pearson correlation before printing per-channel stats.

Example:

    python tools/compare_channel_correlation.py \\
        test_files/correlation/orig5.1.flac \\
        tmp/auro2d_decoded.wav

    python tools/compare_channel_correlation.py ref.wav dec.wav --matrix --max-lag-sec 1
"""

from __future__ import annotations

import argparse
import re
import struct
import subprocess
import sys
import tempfile
import wave
from pathlib import Path

import numpy as np

CHANNEL_NAMES = {
    1: ["M"],
    2: ["FL", "FR"],
    3: ["FL", "FR", "LFE"],
    4: ["FL", "FR", "LS", "RS"],
    5: ["FL", "FR", "C", "LS", "RS"],
    6: ["FL", "FR", "C", "LFE", "LS", "RS"],
    8: ["FL", "FR", "C", "LFE", "LB", "RB", "LS", "RS"],
}

# Microsoft WAVEFORMATEXTENSIBLE speaker bits (low 18 used here).
WAVE_SPEAKER_BITS = [
    (0, "FL"),
    (1, "FR"),
    (2, "C"),
    (3, "LFE"),
    (4, "BL"),
    (5, "BR"),
    (6, "FLC"),
    (7, "FRC"),
    (8, "BC"),
    (9, "SL"),
    (10, "SR"),
    (11, "TC"),
    (12, "TFL"),
    (13, "TFC"),
    (14, "TFR"),
    (15, "TBL"),
    (16, "TBC"),
    (17, "TBR"),
]


def channel_label(index: int, channel_count: int, names: list[str] | None = None) -> str:
    if names is not None and index < len(names):
        return names[index]
    preset = CHANNEL_NAMES.get(channel_count)
    if preset and index < len(preset):
        return preset[index]
    return f"ch{index + 1}"


def wav_channel_names(path: Path, channel_count: int) -> list[str] | None:
    """Return speaker names from dwChannelMask when present and bit count matches."""
    try:
        with open(path, "rb") as f:
            header = f.read(256)
    except OSError:
        return None
    if len(header) < 36 or header[8:12] != b"WAVE":
        return None
    pos = 12
    while pos + 8 <= len(header):
        tag = header[pos : pos + 4]
        size = struct.unpack_from("<I", header, pos + 4)[0]
        body = pos + 8
        if tag == b"fmt " and body + size <= len(header):
            fmt = header[body : body + size]
            if len(fmt) < 16:
                return None
            audio_format = struct.unpack_from("<H", fmt, 0)[0]
            channels = struct.unpack_from("<H", fmt, 2)[0]
            if channels != channel_count:
                return None
            if audio_format == 0xFFFE and len(fmt) >= 40:
                mask = struct.unpack_from("<I", fmt, 20)[0]
                names = [name for bit, name in WAVE_SPEAKER_BITS if mask & (1 << bit)]
                if len(names) == channel_count:
                    return names
            return None
        pos = body + size + (size & 1)
    return None


def auro_xml_channel_names(wav_path: Path, channel_count: int) -> list[str] | None:
    """Read channel names from orua3d-decode sidecar XML if present."""
    xml_path = wav_path.with_suffix(".xml")
    if not xml_path.is_file():
        return None
    try:
        text = xml_path.read_text(encoding="utf-8", errors="replace")
    except OSError:
        return None
    found: dict[int, str] = {}
    for match in re.finditer(
        r'<channel\b[^>]*\bindex="(\d+)"[^>]*\bname="([^"]+)"',
        text,
    ):
        found[int(match.group(1))] = match.group(2)
    if len(found) != channel_count:
        return None
    return [found[i] for i in range(channel_count)]


def resolve_channel_names(
    path: Path,
    channel_count: int,
    wav_path: Path | None = None,
) -> list[str] | None:
    # Prefer WAVE mask on the actual WAV bytes; fall back to sidecar XML near
    # the original path (temp ffmpeg WAVs have no XML).
    names = wav_channel_names(wav_path or path, channel_count)
    if names is not None:
        return names
    return auro_xml_channel_names(path, channel_count)


def safe_unlink(path: Path | None) -> None:
    if path is None:
        return
    try:
        path.unlink(missing_ok=True)
    except OSError:
        pass


def ensure_wav(path: Path) -> tuple[Path, Path | None]:
    """Return a WAV path; if input is not WAV, decode via ffmpeg to a temp file."""
    suffix = path.suffix.lower()
    if suffix == ".wav":
        return path, None
    tmp = Path(tempfile.mkstemp(prefix="corr_", suffix=".wav")[1])
    cmd = [
        "ffmpeg",
        "-y",
        "-v",
        "error",
        "-i",
        str(path),
        "-map",
        "0:a:0",
        "-c:a",
        "pcm_s24le",
        str(tmp),
    ]
    try:
        subprocess.check_call(cmd)
    except (OSError, subprocess.CalledProcessError) as exc:
        safe_unlink(tmp)
        raise SystemExit(f"ffmpeg failed to decode {path}: {exc}") from exc
    return tmp, tmp


def read_wav(path: Path) -> tuple[np.ndarray, int]:
    with wave.open(str(path), "rb") as wav:
        channels = wav.getnchannels()
        sample_width = wav.getsampwidth()
        sample_rate = wav.getframerate()
        frames = wav.getnframes()
        raw = wav.readframes(frames)

    if sample_width == 2:
        data = np.frombuffer(raw, dtype="<i2").astype(np.float64)
    elif sample_width == 3:
        bytes_ = np.frombuffer(raw, dtype=np.uint8).reshape(-1, 3)
        data = (
            bytes_[:, 0].astype(np.int32)
            | (bytes_[:, 1].astype(np.int32) << 8)
            | (bytes_[:, 2].astype(np.int32) << 16)
        )
        data = ((data << 8) >> 8).astype(np.float64)
    elif sample_width == 4:
        data = np.frombuffer(raw, dtype="<i4").astype(np.float64)
    else:
        raise SystemExit(f"{path}: unsupported sample width {sample_width}")

    if data.size % channels != 0:
        raise SystemExit(f"{path}: PCM size is not a multiple of channel count")
    return data.reshape(-1, channels), sample_rate


def pearson(a: np.ndarray, b: np.ndarray) -> float:
    a = a - a.mean()
    b = b - b.mean()
    da = float(np.sqrt(np.dot(a, a)))
    db = float(np.sqrt(np.dot(b, b)))
    if da == 0.0 or db == 0.0:
        return float("nan")
    return float(np.dot(a, b) / (da * db))


def correlation_matrix(ref: np.ndarray, cand: np.ndarray) -> np.ndarray:
    n_ref = ref.shape[1]
    n_cand = cand.shape[1]
    matrix = np.full((n_ref, n_cand), np.nan, dtype=np.float64)
    for i in range(n_ref):
        for j in range(n_cand):
            matrix[i, j] = pearson(ref[:, i], cand[:, j])
    return matrix


def best_channel_mapping(
    matrix: np.ndarray,
    min_corr: float = 0.25,
) -> list[tuple[int, int, float]]:
    """
    1:1 assignment maximizing sum of Pearson correlations.

    Only pairs with corr >= min_corr are kept (avoids forcing BL/BR onto
    leftover height slots). Uses SciPy Hungarian when available; otherwise
    greedy by descending corr. Returns (ref_index, cand_index, corr),
    sorted by ref_index.
    """
    n_ref, n_cand = matrix.shape
    # Eligible scores: below threshold become very bad so they are not chosen
    # when a better eligible partner exists; after solve we still filter.
    finite = np.where(np.isnan(matrix), -1.0, matrix)
    eligible = finite.copy()
    eligible[finite < min_corr] = -1.0

    pairs: list[tuple[int, int, float]] = []
    try:
        from scipy.optimize import linear_sum_assignment  # type: ignore

        n = max(n_ref, n_cand)
        # Dummy padded cells must be worse than any real eligible pair.
        cost = np.full((n, n), 2.0, dtype=np.float64)
        cost[:n_ref, :n_cand] = -eligible
        row_ind, col_ind = linear_sum_assignment(cost)
        for r, c in zip(row_ind, col_ind):
            if r >= n_ref or c >= n_cand:
                continue
            corr = float(matrix[r, c])
            if corr == corr and corr >= min_corr:
                pairs.append((int(r), int(c), corr))
    except Exception:
        candidates: list[tuple[float, int, int]] = []
        for i in range(n_ref):
            for j in range(n_cand):
                value = float(eligible[i, j])
                if value < min_corr:
                    continue
                candidates.append((value, i, j))
        candidates.sort(reverse=True)

        used_ref: set[int] = set()
        used_cand: set[int] = set()
        for value, i, j in candidates:
            if i in used_ref or j in used_cand:
                continue
            used_ref.add(i)
            used_cand.add(j)
            pairs.append((i, j, float(matrix[i, j])))

    pairs.sort(key=lambda item: item[0])
    return pairs


# Prefer these pairings when Pearson scores are effectively tied.
CHANNEL_AFFINITY: dict[tuple[str, str], int] = {
    ("FL", "FL"): 2,
    ("FR", "FR"): 2,
    ("C", "C"): 2,
    ("LFE", "LFE"): 2,
    ("SL", "LS"): 2,
    ("SR", "RS"): 2,
    ("BL", "LS"): 1,
    ("BR", "RS"): 1,
    ("LS", "LS"): 2,
    ("RS", "RS"): 2,
    ("TFL", "HL"): 2,
    ("TFR", "HR"): 2,
    ("TBL", "HLS"): 2,
    ("TBR", "HRS"): 2,
    ("HL", "HL"): 2,
    ("HR", "HR"): 2,
    ("HLS", "HLS"): 2,
    ("HRS", "HRS"): 2,
}


def channel_affinity(ref_name: str, cand_name: str) -> int:
    if ref_name == cand_name:
        return 3
    return CHANNEL_AFFINITY.get((ref_name, cand_name), 0)


def refine_mapping_by_names(
    pairs: list[tuple[int, int, float]],
    matrix: np.ndarray,
    ref_names: list[str] | None,
    cand_names: list[str] | None,
    ref_channels: int,
    cand_channels: int,
    corr_eps: float = 1e-4,
) -> list[tuple[int, int, float]]:
    """Swap candidate partners when names fit better and corr loss is tiny."""
    if len(pairs) < 2:
        return pairs
    refined = list(pairs)
    changed = True
    while changed:
        changed = False
        for a in range(len(refined)):
            for b in range(a + 1, len(refined)):
                ref_a, cand_a, corr_a = refined[a]
                ref_b, cand_b, corr_b = refined[b]
                swap_corr_a = float(matrix[ref_a, cand_b])
                swap_corr_b = float(matrix[ref_b, cand_a])
                if swap_corr_a != swap_corr_a or swap_corr_b != swap_corr_b:
                    continue
                before = corr_a + corr_b
                after = swap_corr_a + swap_corr_b
                if after < before - corr_eps:
                    continue
                name_ref_a = channel_label(ref_a, ref_channels, ref_names)
                name_ref_b = channel_label(ref_b, ref_channels, ref_names)
                name_cand_a = channel_label(cand_a, cand_channels, cand_names)
                name_cand_b = channel_label(cand_b, cand_channels, cand_names)
                aff_before = channel_affinity(name_ref_a, name_cand_a) + channel_affinity(
                    name_ref_b, name_cand_b
                )
                aff_after = channel_affinity(name_ref_a, name_cand_b) + channel_affinity(
                    name_ref_b, name_cand_a
                )
                if aff_after > aff_before or (
                    aff_after == aff_before and after > before + corr_eps
                ):
                    refined[a] = (ref_a, cand_b, swap_corr_a)
                    refined[b] = (ref_b, cand_a, swap_corr_b)
                    changed = True
    refined.sort(key=lambda item: item[0])
    return refined


def best_lag(ref: np.ndarray, dec: np.ndarray, max_lag: int) -> tuple[int, float]:
    n = min(len(ref), len(dec))
    if n <= 1:
        return 0, float("nan")
    r = ref[:n] - ref[:n].mean()
    d = dec[:n] - dec[:n].mean()
    nr = float(np.sqrt(np.dot(r, r)))
    nd = float(np.sqrt(np.dot(d, d)))
    if nr == 0.0 or nd == 0.0:
        return 0, float("nan")

    size = 2 * n
    fft_size = 1 << (size - 1).bit_length()
    xc = np.fft.irfft(np.fft.rfft(r, fft_size) * np.conj(np.fft.rfft(d, fft_size)), fft_size)

    best = 0
    best_val = -1e300
    limit = min(max_lag, n - 1)
    for lag in range(0, limit + 1):
        value = float(xc[lag])
        if value > best_val:
            best_val = value
            best = lag
    for lag in range(1, limit + 1):
        value = float(xc[-lag])
        if value > best_val:
            best_val = value
            best = -lag
    return best, best_val / (nr * nd)


def align(ref: np.ndarray, dec: np.ndarray, lag: int) -> tuple[np.ndarray, np.ndarray]:
    if lag > 0:
        dec = dec[lag:]
        ref = ref[: len(dec)]
    elif lag < 0:
        ref = ref[-lag:]
        dec = dec[: len(ref)]
    n = min(len(ref), len(dec))
    return ref[:n], dec[:n]


def rms(x: np.ndarray) -> float:
    return float(np.sqrt(np.mean(x * x))) if x.size else 0.0


def fmt_corr(value: float) -> str:
    return "nan" if value != value else f"{value:.6f}"


def fmt_db(value: float) -> str:
    return "nan" if value != value else f"{value:.3f}"


def build_arg_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description=(
            "Compare two audio files with per-channel Pearson correlation. "
            "By default searches lag and best channel mapping."
        )
    )
    parser.add_argument("reference", type=Path, help="reference WAV/FLAC/...")
    parser.add_argument("candidate", type=Path, help="decoded/candidate WAV/FLAC/...")
    parser.add_argument(
        "--max-lag-sec",
        type=float,
        default=1.0,
        help="maximum absolute lag to search, seconds (default: 1)",
    )
    parser.add_argument(
        "--align-offset-sec",
        type=float,
        default=2.0,
        help="start of lag-estimation window on the reference, seconds (default: 2)",
    )
    parser.add_argument(
        "--align-duration-sec",
        type=float,
        default=4.0,
        help="lag-estimation window duration, seconds (default: 4)",
    )
    parser.add_argument(
        "--align-channels",
        type=str,
        default="1,2",
        help="1-based reference channels averaged for lag search (default: 1,2)",
    )
    parser.add_argument(
        "--matrix",
        action="store_true",
        help="also print full reference-vs-candidate correlation matrix",
    )
    parser.add_argument(
        "--no-align",
        action="store_true",
        help="skip lag search and compare from sample 0",
    )
    parser.add_argument(
        "--no-remap",
        action="store_true",
        help="compare channels by index only (disable auto channel mapping)",
    )
    parser.add_argument(
        "--min-corr",
        type=float,
        default=0.25,
        help=(
            "minimum Pearson corr to accept a mapped pair (default: 0.25); "
            "use 0 to force a full 1:1 assignment"
        ),
    )
    return parser


def parse_channel_list(value: str, channel_count: int) -> list[int]:
    channels: list[int] = []
    for part in value.split(","):
        part = part.strip()
        if not part:
            continue
        index = int(part)
        if index < 1 or index > channel_count:
            raise SystemExit(
                f"align channel {index} out of range 1..{channel_count}"
            )
        channels.append(index - 1)
    if not channels:
        raise SystemExit("at least one --align-channels value is required")
    return channels


def print_stats_table(
    pairs: list[tuple[str, np.ndarray, np.ndarray]],
) -> None:
    full_scale = float(1 << 23)
    print(f"{'ch':<14} {'corr':>10} {'ref_rms':>12} {'cand_rms':>12} {'gain_db':>10}")
    for label, ref_ch, cand_ch in pairs:
        corr = pearson(ref_ch, cand_ch)
        ref_rms = rms(ref_ch) / full_scale
        cand_rms = rms(cand_ch) / full_scale
        if ref_rms > 0.0 and cand_rms > 0.0:
            gain_db = 20.0 * np.log10(cand_rms / ref_rms)
        else:
            gain_db = float("nan")
        print(
            f"{label:<14} {fmt_corr(corr):>10} {ref_rms:12.6f} "
            f"{cand_rms:12.6f} {fmt_db(gain_db):>10}"
        )


def main(argv: list[str] | None = None) -> int:
    args = build_arg_parser().parse_args(argv)
    if not args.reference.is_file():
        raise SystemExit(f"reference not found: {args.reference}")
    if not args.candidate.is_file():
        raise SystemExit(f"candidate not found: {args.candidate}")

    ref_wav, ref_tmp = ensure_wav(args.reference)
    cand_wav, cand_tmp = ensure_wav(args.candidate)
    try:
        ref, ref_rate = read_wav(ref_wav)
        cand, cand_rate = read_wav(cand_wav)
        ref_names = resolve_channel_names(args.reference, ref.shape[1], ref_wav)
        cand_names = resolve_channel_names(args.candidate, cand.shape[1], cand_wav)
    finally:
        safe_unlink(ref_tmp)
        safe_unlink(cand_tmp)

    if ref_rate != cand_rate:
        raise SystemExit(
            f"sample rate mismatch: reference {ref_rate} Hz, candidate {cand_rate} Hz"
        )
    if args.no_remap and ref.shape[1] != cand.shape[1]:
        raise SystemExit(
            f"channel count mismatch: reference {ref.shape[1]}, candidate {cand.shape[1]} "
            f"(omit --no-remap to allow auto mapping)"
        )

    ref_channels = ref.shape[1]
    cand_channels = cand.shape[1]
    sample_rate = ref_rate
    n = min(len(ref), len(cand))
    ref = ref[:n]
    cand = cand[:n]

    lag = 0
    lag_score = float("nan")
    if not args.no_align:
        align_ref = parse_channel_list(args.align_channels, ref_channels)
        # Pair align channels by position for lag estimate when counts differ.
        align_cand = [
            min(idx, cand_channels - 1) for idx in align_ref
        ]
        offset = max(0, int(round(args.align_offset_sec * sample_rate)))
        duration = max(1, int(round(args.align_duration_sec * sample_rate)))
        if offset >= n:
            offset = 0
        duration = min(duration, n - offset)
        ref_win = ref[offset : offset + duration, align_ref].mean(axis=1)
        cand_win = cand[offset : offset + duration, align_cand].mean(axis=1)
        max_lag = max(0, int(round(args.max_lag_sec * sample_rate)))
        lag, lag_score = best_lag(ref_win, cand_win, max_lag)

    ref_a, cand_a = align(ref, cand, lag)
    frames = len(ref_a)

    print(f"reference: {args.reference}")
    print(f"candidate: {args.candidate}")
    print(
        f"rate={sample_rate} ref_channels={ref_channels} cand_channels={cand_channels} "
        f"compared_frames={frames} ({frames / sample_rate:.3f}s)"
    )
    print(
        f"lag={lag} samples ({1000.0 * lag / sample_rate:.2f} ms) "
        f"score={fmt_corr(lag_score)}"
    )

    matrix = correlation_matrix(ref_a, cand_a)

    if args.no_remap:
        if ref_channels != cand_channels:
            raise SystemExit("internal: --no-remap requires equal channel counts")
        print("mapping: identity (same channel index)")
        print()
        pairs = [
            (
                channel_label(i, ref_channels, ref_names),
                ref_a[:, i],
                cand_a[:, i],
            )
            for i in range(ref_channels)
        ]
        print_stats_table(pairs)
    else:
        mapping = best_channel_mapping(matrix, min_corr=args.min_corr)
        mapping = refine_mapping_by_names(
            mapping,
            matrix,
            ref_names,
            cand_names,
            ref_channels,
            cand_channels,
        )
        print(f"mapping: ref <- cand  (min_corr={args.min_corr:g})")
        for ref_i, cand_j, corr in mapping:
            ref_l = channel_label(ref_i, ref_channels, ref_names)
            cand_l = channel_label(cand_j, cand_channels, cand_names)
            print(
                f"  [{ref_i}]{ref_l} <- [{cand_j}]{cand_l}  corr={fmt_corr(corr)}"
            )
        mapped_ref = {ref_i for ref_i, _, _ in mapping}
        mapped_cand = {cand_j for _, cand_j, _ in mapping}
        unused_ref = [i for i in range(ref_channels) if i not in mapped_ref]
        unused_cand = [j for j in range(cand_channels) if j not in mapped_cand]
        if unused_ref:
            labels = ", ".join(
                f"[{i}]{channel_label(i, ref_channels, ref_names)}" for i in unused_ref
            )
            print(f"unmapped reference: {labels}")
        if unused_cand:
            labels = ", ".join(
                f"[{j}]{channel_label(j, cand_channels, cand_names)}" for j in unused_cand
            )
            print(f"unused candidate: {labels}")
        print()
        pairs = []
        for ref_i, cand_j, _ in mapping:
            ref_l = channel_label(ref_i, ref_channels, ref_names)
            cand_l = channel_label(cand_j, cand_channels, cand_names)
            label = ref_l if ref_l == cand_l else f"{ref_l}<-{cand_l}"
            pairs.append((label, ref_a[:, ref_i], cand_a[:, cand_j]))
        if pairs:
            print_stats_table(pairs)
        else:
            print("no channel pairs above --min-corr")

    if args.matrix:
        print()
        print("cross-corr matrix (rows=reference, cols=candidate):")
        row_labels = [
            channel_label(i, ref_channels, ref_names) for i in range(ref_channels)
        ]
        col_labels = [
            channel_label(j, cand_channels, cand_names) for j in range(cand_channels)
        ]
        print("      " + " ".join(f"{name:>8}" for name in col_labels))
        for i, row_name in enumerate(row_labels):
            values = []
            for j in range(cand_channels):
                value = float(matrix[i, j])
                values.append(0.0 if value != value else value)
            print(f"{row_name:<4}  " + " ".join(f"{v:8.4f}" for v in values))

    return 0


if __name__ == "__main__":
    sys.exit(main())