"""Compare H3 audio-latent and lossless-waveform boundaries."""

import argparse
import json
import math
import wave
from pathlib import Path

import numpy as np
import torch


def dbfs(value: float) -> float:
    return 20.0 * math.log10(max(value, 1e-20))


def load_audio_latent(path: Path) -> torch.Tensor:
    state = torch.load(path, map_location="cpu", weights_only=False)
    latent = state.get("audio_latent") if isinstance(state, dict) else state
    if latent is None or latent.ndim != 4:
        raise ValueError(f"{path} does not contain a [B,C,S,T] audio latent")
    return latent.float()


def latent_metrics(latent: torch.Tensor) -> dict:
    frames = latent.movedim(-1, 0).flatten(1)
    frame_rms = frames.square().mean(1).sqrt()
    frame_mean = frames.mean(1)
    frame_max = frames.abs().amax(1)
    deltas = frames[1:] - frames[:-1]
    delta_rms = deltas.square().mean(1).sqrt()
    adjacent_cosine = torch.nn.functional.cosine_similarity(frames[:-1], frames[1:], dim=1)

    block_frames = min(4, frames.shape[0] // 2)
    first = frames[:block_frames].flatten()
    last = frames[-block_frames:].flatten()
    first_last_cosine = torch.nn.functional.cosine_similarity(first, last, dim=0)

    first_count = min(20, frames.shape[0])
    return {
        "shape": list(latent.shape),
        "dtype": str(latent.dtype),
        "first_20_frame_rms": frame_rms[:first_count].tolist(),
        "first_20_frame_mean": frame_mean[:first_count].tolist(),
        "first_20_frame_max_abs": frame_max[:first_count].tolist(),
        "first_19_delta_rms": delta_rms[: max(0, first_count - 1)].tolist(),
        "first_19_adjacent_cosine": adjacent_cosine[: max(0, first_count - 1)].tolist(),
        "first_4_rms": float(frames[:block_frames].square().mean().sqrt()),
        "frames_4_20_rms": float(frames[block_frames:first_count].square().mean().sqrt()),
        "remaining_rms": float(frames[first_count:].square().mean().sqrt()),
        "first_4_vs_last_4_cosine": float(first_last_cosine),
        "largest_delta_frame": int(delta_rms.argmax().item() + 1),
        "largest_delta_rms": float(delta_rms.max()),
    }


def load_wav(path: Path) -> tuple[np.ndarray, int]:
    with wave.open(str(path), "rb") as source:
        if source.getsampwidth() != 2:
            raise ValueError(f"{path} must be PCM S16")
        channels = source.getnchannels()
        sample_rate = source.getframerate()
        samples = np.frombuffer(source.readframes(source.getnframes()), dtype="<i2")
    return samples.reshape(-1, channels).astype(np.float32) / 32768.0, sample_rate


def waveform_metrics(samples: np.ndarray, sample_rate: int) -> dict:
    first_half_second = samples[: sample_rate // 2]
    mono = first_half_second.mean(1)
    window_samples = sample_rate // 100
    windows = []
    for start in range(0, len(first_half_second), window_samples):
        block = first_half_second[start : start + window_samples]
        if len(block) == 0:
            continue
        windows.append({
            "start_ms": start * 1000.0 / sample_rate,
            "peak_dbfs": dbfs(float(np.max(np.abs(block)))),
            "rms_dbfs": dbfs(float(np.sqrt(np.mean(block * block)))),
            "mean": float(block.mean()),
        })

    spectrum_samples = min(sample_rate // 10, len(mono))
    windowed = mono[:spectrum_samples] * np.hanning(spectrum_samples)
    magnitudes = np.abs(np.fft.rfft(windowed))
    frequencies = np.fft.rfftfreq(spectrum_samples, 1.0 / sample_rate)
    dominant = np.argsort(magnitudes[1:])[-8:][::-1] + 1
    derivatives = np.max(np.abs(np.diff(first_half_second, axis=0)), axis=1)
    return {
        "sample_rate": sample_rate,
        "samples": len(samples),
        "first_sample": samples[0].tolist(),
        "first_sample_dbfs": [dbfs(float(abs(value))) for value in samples[0]],
        "first_500ms_peak_dbfs": dbfs(float(np.max(np.abs(first_half_second)))),
        "first_500ms_rms_dbfs": dbfs(float(np.sqrt(np.mean(first_half_second**2)))),
        "largest_derivative": float(derivatives.max()),
        "largest_derivative_ms": float((derivatives.argmax() + 1) * 1000.0 / sample_rate),
        "dominant_first_100ms_hz": [float(frequencies[index]) for index in dominant],
        "windows_10ms": windows,
    }


def analyze(latent_path: Path, wav_path: Path) -> dict:
    samples, sample_rate = load_wav(wav_path)
    return {
        "latent_path": str(latent_path),
        "wav_path": str(wav_path),
        "latent": latent_metrics(load_audio_latent(latent_path)),
        "waveform": waveform_metrics(samples, sample_rate),
    }


parser = argparse.ArgumentParser()
parser.add_argument("--affected-latent", type=Path, required=True)
parser.add_argument("--affected-wav", type=Path, required=True)
parser.add_argument("--clean-latent", type=Path, required=True)
parser.add_argument("--clean-wav", type=Path, required=True)
args = parser.parse_args()

print(json.dumps({
    "affected": analyze(args.affected_latent, args.affected_wav),
    "clean": analyze(args.clean_latent, args.clean_wav),
}, indent=2))