#!/usr/bin/env python
"""Probe FOA azimuth conventions using a coarse active-intensity estimate.

This script is intended for debugging Spatial-BEATs training when azimuth
learning stalls near random. It reads manifest entries, crops each source to
its weak active window, computes a coarse FOA active-intensity vector from the
mixture waveform, and compares several azimuth conventions against the GT.

The goal is not to produce a perfect DOA estimator. The goal is to answer:
    "Is the current FOA / azimuth coordinate convention obviously flipped,
    swapped, or rotated before I even train the model?"
"""

from __future__ import annotations

import argparse
import math
from collections import defaultdict
from pathlib import Path
from typing import Dict, Iterable, List, Optional, Sequence, Tuple

import torch
from tqdm.auto import tqdm

from spatial_dataset import _load_audio_file, _load_manifest_entries


def circular_distance_deg(a_deg: float, b_deg: float) -> float:
    """Return the wrapped absolute angular distance in degrees."""
    return abs(((a_deg - b_deg + 180.0) % 360.0) - 180.0)


def normalize_deg(angle_deg: float) -> float:
    """Normalize an angle to [0, 360)."""
    return angle_deg % 360.0


def resolve_clip_path(entry: Dict[str, object]) -> str:
    """Resolve the FOA waveform path for one manifest entry."""
    for key in ("output_foa_path", "waveform_path", "audio_path", "foa_path"):
        value = entry.get(key)
        if value:
            return str(value)
    raise KeyError("Manifest entry is missing an FOA waveform path.")


def resolve_clip_duration_seconds(entry: Dict[str, object], waveform: torch.Tensor, sample_rate: int) -> float:
    """Resolve clip duration, falling back to waveform length when needed."""
    for key in ("clip_duration_seconds", "output_duration_seconds", "duration"):
        value = entry.get(key)
        if value is not None:
            return float(value)
    return float(waveform.size(-1)) / float(sample_rate)


def resolve_entry_id(entry: Dict[str, object], default_index: int) -> str:
    """Resolve a human-readable sample identifier for logging."""
    for key in ("scene_id", "pair_id", "sample_id", "id"):
        value = entry.get(key)
        if value is not None:
            return str(value)
    return str(default_index)


def resolve_source_times(source_entry: Dict[str, object], clip_duration_seconds: float) -> Tuple[float, float]:
    """Resolve the weak active window used by the supervision pipeline."""
    active_time = source_entry.get("active_time")
    full_time = source_entry.get("full_time")
    if isinstance(active_time, Sequence) and len(active_time) >= 2:
        return float(active_time[0]), float(active_time[1])
    if isinstance(full_time, Sequence) and len(full_time) >= 2:
        return float(full_time[0]), float(full_time[1])
    return 0.0, float(clip_duration_seconds)


def resolve_source_azimuth_deg(entry: Dict[str, object], source_entry: Dict[str, object]) -> float:
    """Resolve GT azimuth in degrees from source-level or top-level fields."""
    doa = source_entry.get("doa")
    if isinstance(doa, dict) and doa.get("azimuth_deg") is not None:
        return float(doa["azimuth_deg"])
    for key in ("azimuth_deg", "azimuth"):
        value = source_entry.get(key)
        if value is not None:
            return float(value)
    if entry.get("rir_doa_azimuth_deg") is not None:
        return float(entry["rir_doa_azimuth_deg"])
    raise KeyError("Unable to resolve GT azimuth from manifest entry.")


def resolve_source_label(source_entry: Dict[str, object]) -> str:
    """Resolve a readable label for debugging output."""
    for key in ("mono_target_label", "mono_primary_label", "final_label", "label"):
        value = source_entry.get(key)
        if value:
            return str(value)
    return "<unknown>"


def iter_sources(entry: Dict[str, object], clip_duration_seconds: float) -> List[Dict[str, object]]:
    """Return source dicts in a unified shape for ov1/ov2/ov3 manifests."""
    sources = entry.get("sources")
    if isinstance(sources, list) and sources:
        return [dict(source) for source in sources if isinstance(source, dict)]

    return [
        {
            "mono_target_label": entry.get("mono_target_label", entry.get("mono_primary_label")),
            "doa": {
                "azimuth_deg": entry.get("rir_doa_azimuth_deg"),
                "elevation_deg": entry.get("rir_doa_elevation_deg"),
            },
            "active_time": [0.0, clip_duration_seconds],
            "full_time": [0.0, clip_duration_seconds],
        }
    ]


def is_isolated_window(source_index: int, sources: Sequence[Dict[str, object]], clip_duration_seconds: float) -> bool:
    """Check whether a source weak window overlaps with any other source window."""
    start_a, end_a = resolve_source_times(sources[source_index], clip_duration_seconds)
    for other_index, other_source in enumerate(sources):
        if other_index == source_index:
            continue
        start_b, end_b = resolve_source_times(other_source, clip_duration_seconds)
        if min(end_a, end_b) > max(start_a, start_b):
            return False
    return True


def crop_waveform_to_window(
    waveform: torch.Tensor,
    sample_rate: int,
    start_time_seconds: float,
    end_time_seconds: float,
) -> torch.Tensor:
    """Crop one FOA waveform to a weak source activity window."""
    total_num_samples = waveform.size(-1)
    start_sample = max(0, min(int(math.floor(start_time_seconds * sample_rate)), total_num_samples - 1))
    end_sample = max(start_sample + 1, min(int(math.ceil(end_time_seconds * sample_rate)), total_num_samples))
    return waveform[:, start_sample:end_sample].contiguous()


def reorder_dcase_wyzx_to_wxyz(waveform: torch.Tensor) -> torch.Tensor:
    """Convert stored DCASE FOA waveform order [W, Y, Z, X] to [W, X, Y, Z]."""
    if waveform.ndim != 2 or waveform.size(0) != 4:
        raise ValueError(f"Expected waveform [4, T], got {tuple(waveform.shape)}")
    return waveform[[0, 3, 1, 2], :]


def estimate_active_intensity_vector(
    waveform: torch.Tensor,
    sample_rate: int,
    n_fft: int,
    win_length: int,
    hop_length: int,
    frame_energy_quantile: float,
) -> Tuple[float, float, float]:
    """Estimate a coarse FOA active-intensity vector from one cropped waveform.

    Returns:
        Tuple[float, float, float]:
            Mean active-intensity components (Ix, Iy, Iz).
    """
    if waveform.ndim != 2 or waveform.size(0) != 4:
        raise ValueError(f"Expected waveform [4, T], got {tuple(waveform.shape)}")

    waveform = reorder_dcase_wyzx_to_wxyz(waveform)
    window = torch.hann_window(win_length, dtype=waveform.dtype, device=waveform.device)
    stft = torch.stft(
        waveform,
        n_fft=n_fft,
        hop_length=hop_length,
        win_length=win_length,
        window=window,
        center=True,
        pad_mode="reflect",
        return_complex=True,
    )
    w = stft[0]
    x = stft[1]
    y = stft[2]
    z = stft[3]
    power = w.abs().pow(2.0)

    frame_energy = power.sum(dim=0)
    if frame_energy.numel() == 0:
        return 0.0, 0.0, 0.0

    threshold = torch.quantile(frame_energy, q=float(frame_energy_quantile))
    active_frame_mask = frame_energy >= threshold
    if not bool(active_frame_mask.any()):
        active_frame_mask = torch.ones_like(frame_energy, dtype=torch.bool)

    power = power[:, active_frame_mask]
    ix = torch.real(w[:, active_frame_mask] * torch.conj(x[:, active_frame_mask]))
    iy = torch.real(w[:, active_frame_mask] * torch.conj(y[:, active_frame_mask]))
    iz = torch.real(w[:, active_frame_mask] * torch.conj(z[:, active_frame_mask]))

    weight = power
    denom = torch.clamp(weight.sum(), min=1e-8)
    ix_mean = float((ix * weight).sum().item() / denom.item())
    iy_mean = float((iy * weight).sum().item() / denom.item())
    iz_mean = float((iz * weight).sum().item() / denom.item())
    return ix_mean, iy_mean, iz_mean


def azimuth_from_components(x_comp: float, y_comp: float) -> float:
    """Convert x/y Cartesian components to azimuth degrees."""
    return normalize_deg(math.degrees(math.atan2(y_comp, x_comp)))


def build_convention_predictions(ix: float, iy: float) -> Dict[str, float]:
    """Evaluate several common FOA azimuth sign / axis conventions."""
    return {
        "atan2(+y,+x)": azimuth_from_components(+ix, +iy),
        "atan2(-y,+x)": azimuth_from_components(+ix, -iy),
        "atan2(+y,-x)": azimuth_from_components(-ix, +iy),
        "atan2(-y,-x)": azimuth_from_components(-ix, -iy),
        "atan2(+x,+y)": azimuth_from_components(+iy, +ix),
        "atan2(-x,+y)": azimuth_from_components(+iy, -ix),
        "atan2(+x,-y)": azimuth_from_components(-iy, +ix),
        "atan2(-x,-y)": azimuth_from_components(-iy, -ix),
    }


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Probe FOA IV azimuth alignment against GT.")
    parser.add_argument("--manifest", type=str, required=True, help="Path to ov*.jsonl manifest.")
    parser.add_argument("--split", type=str, default=None, help="Optional split filter, e.g. train/valid/test.")
    parser.add_argument("--limit", type=int, default=200, help="Maximum number of usable source windows to evaluate.")
    parser.add_argument("--sample-rate", type=int, default=16000, help="Expected FOA sample rate.")
    parser.add_argument("--n-fft", type=int, default=400, help="STFT FFT size.")
    parser.add_argument("--win-length", type=int, default=400, help="STFT window length.")
    parser.add_argument("--hop-length", type=int, default=160, help="STFT hop length.")
    parser.add_argument("--min-window-seconds", type=float, default=0.3, help="Skip very short source windows.")
    parser.add_argument("--frame-energy-quantile", type=float, default=0.7, help="Use only high-energy frames above this quantile.")
    parser.add_argument(
        "--require-isolated-window",
        action="store_true",
        help="Only evaluate source windows that do not overlap any other source window in the same clip.",
    )
    parser.add_argument("--show-examples", type=int, default=12, help="Number of per-sample examples to print.")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    manifest_path = Path(args.manifest)
    entries = _load_manifest_entries(manifest_path, show_progress=False)
    if args.split is not None:
        entries = [entry for entry in entries if entry.get("split") == args.split]

    convention_errors: Dict[str, List[float]] = defaultdict(list)
    examples: List[Dict[str, object]] = []
    num_skipped_short = 0
    num_skipped_overlap = 0
    num_skipped_zero_vector = 0
    num_audio_failures = 0

    progress = tqdm(entries, desc=f"Probe IV azimuth {manifest_path.name}")
    usable_windows = 0
    for entry_index, entry in enumerate(progress):
        if usable_windows >= args.limit:
            break

        try:
            clip_path = resolve_clip_path(entry)
            waveform = _load_audio_file(clip_path, args.sample_rate)
        except Exception:
            num_audio_failures += 1
            continue

        clip_duration_seconds = resolve_clip_duration_seconds(entry, waveform, args.sample_rate)
        sources = iter_sources(entry, clip_duration_seconds)
        sample_id = resolve_entry_id(entry, entry_index)

        for source_index, source in enumerate(sources):
            if usable_windows >= args.limit:
                break
            if args.require_isolated_window and not is_isolated_window(source_index, sources, clip_duration_seconds):
                num_skipped_overlap += 1
                continue

            start_time_seconds, end_time_seconds = resolve_source_times(source, clip_duration_seconds)
            if end_time_seconds - start_time_seconds < args.min_window_seconds:
                num_skipped_short += 1
                continue

            segment = crop_waveform_to_window(
                waveform=waveform,
                sample_rate=args.sample_rate,
                start_time_seconds=start_time_seconds,
                end_time_seconds=end_time_seconds,
            )
            ix, iy, iz = estimate_active_intensity_vector(
                waveform=segment,
                sample_rate=args.sample_rate,
                n_fft=args.n_fft,
                win_length=args.win_length,
                hop_length=args.hop_length,
                frame_energy_quantile=args.frame_energy_quantile,
            )
            xy_norm = math.sqrt(ix * ix + iy * iy)
            if xy_norm < 1e-8:
                num_skipped_zero_vector += 1
                continue

            gt_azimuth_deg = normalize_deg(resolve_source_azimuth_deg(entry, source))
            predictions = build_convention_predictions(ix, iy)
            for convention_name, pred_azimuth_deg in predictions.items():
                convention_errors[convention_name].append(
                    circular_distance_deg(pred_azimuth_deg, gt_azimuth_deg)
                )

            examples.append(
                {
                    "sample_id": sample_id,
                    "source_index": source_index,
                    "label": resolve_source_label(source),
                    "gt_azimuth_deg": gt_azimuth_deg,
                    "ix": ix,
                    "iy": iy,
                    "iz": iz,
                    "window": (start_time_seconds, end_time_seconds),
                    "predictions": predictions,
                }
            )
            usable_windows += 1
            progress.set_postfix(usable=usable_windows)

    print()
    print(f"Manifest: {manifest_path}")
    print(f"Split: {args.split or '<all>'}")
    print(f"Usable source windows: {usable_windows}")
    print(f"Skipped short windows: {num_skipped_short}")
    print(f"Skipped overlapping windows: {num_skipped_overlap}")
    print(f"Skipped zero XY intensity: {num_skipped_zero_vector}")
    print(f"Audio load failures: {num_audio_failures}")

    if usable_windows == 0:
        print("No usable source windows found.")
        return

    summary_rows: List[Tuple[str, float, float]] = []
    for convention_name, errors in convention_errors.items():
        error_tensor = torch.tensor(errors, dtype=torch.float32)
        summary_rows.append(
            (
                convention_name,
                float(error_tensor.mean().item()),
                float(error_tensor.median().item()),
            )
        )
    summary_rows.sort(key=lambda row: row[1])

    print()
    print("Convention ranking by circular azimuth error:")
    for convention_name, mean_error, median_error in summary_rows:
        print(
            f"  {convention_name:<15} mean_abs_err={mean_error:7.3f} deg"
            f"  median_abs_err={median_error:7.3f} deg"
        )

    best_convention = summary_rows[0][0]
    print()
    print(f"Examples using best convention: {best_convention}")
    for example in examples[: args.show_examples]:
        pred = float(example["predictions"][best_convention])
        err = circular_distance_deg(pred, float(example["gt_azimuth_deg"]))
        print(
            f"  {example['sample_id']} src={example['source_index']} "
            f"label={example['label']} window={example['window'][0]:.2f}-{example['window'][1]:.2f}s "
            f"GT={example['gt_azimuth_deg']:7.2f} pred={pred:7.2f} err={err:6.2f} "
            f"IV=({example['ix']:+.4f},{example['iy']:+.4f},{example['iz']:+.4f})"
        )


if __name__ == "__main__":
    main()