from __future__ import annotations

import argparse
import csv
import hashlib
import random
from pathlib import Path
from typing import Iterable

import numpy as np
from scipy import signal
from tqdm import tqdm

CLASS_DIRS = {
    "DAUTEL EVO nano": ["DAUTEL EVO NANO", "DAUTEL EVO nano"],
    "DJI MAVIC3 PRO": ["DJI MAVIC3 PRO"],
    "DJI MINI3": ["DJI MINI3"],
    "DJI MINI4 PRO": ["DJI MINI4 PRO"],
}


def load_iq(path: Path, dtype: str) -> np.ndarray:
    if dtype == "complex64":
        return np.fromfile(path, dtype=np.complex64)

    if dtype == "float32_iq":
        raw = np.fromfile(path, dtype=np.float32)
        raw = raw[: raw.size - (raw.size % 2)]
        return raw[0::2] + 1j * raw[1::2]

    if dtype == "int16_iq":
        raw = np.fromfile(path, dtype=np.int16).astype(np.float32)
        raw = raw[: raw.size - (raw.size % 2)]
        iq = raw[0::2] + 1j * raw[1::2]
        return iq / 32768.0

    raise ValueError(f"Unsupported dtype: {dtype}")


def find_class_files(root: Path, label: str) -> list[Path]:
    files: list[Path] = []
    for dirname in CLASS_DIRS[label]:
        files.extend((root / dirname).rglob("*.iq"))
    if not files:
        # Fallback: match class name anywhere in path, case-insensitive.
        target = label.replace(" ", "").lower()
        for path in root.rglob("*.iq"):
            compact = "".join(path.parts).replace(" ", "").lower()
            if target in compact:
                files.append(path)
    return sorted(set(files))


def iter_windows(iq: np.ndarray, window_size: int, stride: int) -> Iterable[np.ndarray]:
    for start in range(0, max(0, len(iq) - window_size + 1), stride):
        yield iq[start : start + window_size]


def spectrogram_db(iq_window: np.ndarray, sample_rate: float, nperseg: int, noverlap: int, nfft: int) -> np.ndarray:
    _, _, sxx = signal.spectrogram(
        iq_window,
        fs=sample_rate,
        window="hann",
        nperseg=nperseg,
        noverlap=noverlap,
        nfft=nfft,
        return_onesided=False,
        mode="magnitude",
        scaling="density",
    )
    sxx = np.fft.fftshift(sxx, axes=0)
    sxx_db = 20.0 * np.log10(sxx + 1e-12)
    return sxx_db.astype(np.float32)


def split_for_key(key: str, train_ratio: float, val_ratio: float) -> str:
    value = int(hashlib.sha1(key.encode("utf-8")).hexdigest(), 16) / float(16**40)
    if value < train_ratio:
        return "train"
    if value < train_ratio + val_ratio:
        return "val"
    return "test"


def main() -> None:
    parser = argparse.ArgumentParser(description="Convert RFUAV .iq files into spectrogram .npz samples.")
    parser.add_argument("--extracted-dir", default="/data/RFUAV_extracted")
    parser.add_argument("--out-dir", default="/data/RFUAV_processed")
    parser.add_argument("--dtype", choices=["complex64", "float32_iq", "int16_iq"], default="complex64")
    parser.add_argument("--sample-rate", type=float, default=20e6, help="Used for STFT axes. Update after XML inspection if needed.")
    parser.add_argument("--window-size", type=int, default=8192)
    parser.add_argument("--stride", type=int, default=8192)
    parser.add_argument("--nperseg", type=int, default=256)
    parser.add_argument("--noverlap", type=int, default=128)
    parser.add_argument("--nfft", type=int, default=256)
    parser.add_argument("--max-windows-per-class", type=int, default=1000)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--train-ratio", type=float, default=0.70)
    parser.add_argument("--val-ratio", type=float, default=0.15)
    args = parser.parse_args()

    random.seed(args.seed)
    root = Path(args.extracted_dir)
    out_dir = Path(args.out_dir)
    sample_dir = out_dir / "samples"
    sample_dir.mkdir(parents=True, exist_ok=True)

    labels = list(CLASS_DIRS.keys())
    label_to_id = {label: idx for idx, label in enumerate(labels)}
    rows: list[dict[str, str | int]] = []

    for label in labels:
        iq_files = find_class_files(root, label)
        print(f"{label}: found {len(iq_files)} .iq files")
        if not iq_files:
            continue

        windows_written = 0
        for iq_path in tqdm(iq_files, desc=label):
            iq = load_iq(iq_path, args.dtype)
            if len(iq) < args.window_size:
                continue

            for window_idx, iq_window in enumerate(iter_windows(iq, args.window_size, args.stride)):
                spec = spectrogram_db(iq_window, args.sample_rate, args.nperseg, args.noverlap, args.nfft)
                key = f"{label}/{iq_path.name}/{window_idx}"
                split = split_for_key(key, args.train_ratio, args.val_ratio)
                file_stem = hashlib.sha1(key.encode("utf-8")).hexdigest()[:16]
                rel_path = Path(split) / label.replace(" ", "_") / f"{file_stem}.npz"
                out_path = sample_dir / rel_path
                out_path.parent.mkdir(parents=True, exist_ok=True)
                np.savez_compressed(out_path, x=spec, y=label_to_id[label], label=label)
                rows.append({"path": str(rel_path), "label": label, "label_id": label_to_id[label], "split": split})

                windows_written += 1
                if windows_written >= args.max_windows_per_class:
                    break
            if windows_written >= args.max_windows_per_class:
                break

        print(f"{label}: wrote {windows_written} spectrogram windows")

    with (out_dir / "manifest.csv").open("w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=["path", "label", "label_id", "split"])
        writer.writeheader()
        writer.writerows(rows)

    with (out_dir / "labels.txt").open("w") as f:
        for label in labels:
            f.write(f"{label_to_id[label]}\t{label}\n")

    print(f"Wrote {len(rows)} samples to {out_dir}")
    print(f"Manifest: {out_dir / 'manifest.csv'}")


if __name__ == "__main__":
    main()