#!/usr/bin/env python3
"""
fft_analyse.py  —  General-purpose IQ spectrum analyser
Python >= 3.8   Dependencies: numpy>=1.20, matplotlib>=3.4

Usage (standalone):
    python fft_analyse.py <input_file> [options]

Options:
    --i-col  NAME     Column name for I channel (auto-detected if omitted)
    --q-col  NAME     Column name for Q channel (auto-detected if omitted)
    --fs     FLOAT    Sample rate in Hz, e.g. 120e6  (default: auto from filename or 1.0)
    --bits   INT      Signed bit-width for full-scale ref: 8/12/16/24/32 (default: auto)
    --fs-val FLOAT    Full-scale amplitude override (overrides --bits)
    --xlim   F,F      Frequency display range in MHz, e.g. -10,10 (default: +-fs/2)
    --out    DIR      Output directory (default: same directory as input file)
    --prefix STR      Filename prefix for outputs (default: derived from input name)
    --skip   INT      Skip first N rows of data (e.g. FIR startup zeros)
    --title  STR      Plot title override

Supported input formats
-----------------------
CSV / TSV with header row.  Column order doesn't matter.

I-channel values can be:
  • Signed decimal integer:    -1234  or  1234
  • Unsigned hex string:       0x0bad  or  0BAD   (2's complement by bit-width)

Auto-detection priority for column names
-----------------------------------------
1. User-supplied --i-col / --q-col
2. Columns matching patterns below (case-insensitive, first match wins):
     I : data_i, iq_i, ch_i, m_i, i_dec, i_data, i[  …  followed by Q equivalent
     Q : data_q, iq_q, ch_q, m_q, q_dec, q_data, q[
3. If only two non-time numeric columns exist, treat first as I and second as Q.

Outputs  (saved to --out dir, filename = <prefix>_<YYYYMMDD_HHMMSS>_*)
--------
  *_analysis.png   Three-panel figure (waveform / magnitude / spectrum)
  *_analysis.svg   Vector copy
  *_spectrum.csv   freq_hz, freq_mhz, amp, dbfs columns
  *_summary.txt    Numeric statistics (OBW, peaks, RMS …)
"""
from __future__ import annotations

import argparse
import csv
import re
import subprocess
import sys
from datetime import datetime
from pathlib import Path

# ── dependency check / auto-install ─────────────────────────────────────────
def _ensure_deps() -> None:
    missing = []
    for pkg in ("numpy", "matplotlib"):
        try:
            __import__(pkg)
        except ImportError:
            missing.append(pkg)
    if missing:
        print(f"[fft_analyse] Installing missing packages: {missing}")
        subprocess.check_call([sys.executable, "-m", "pip", "install", "--quiet"] + missing)

_ensure_deps()

import matplotlib
matplotlib.use("Agg")          # headless — no GUI needed
import matplotlib.pyplot as plt
import numpy as np

# ── helpers ──────────────────────────────────────────────────────────────────

RUN_STAMP = datetime.now().strftime("%Y%m%d_%H%M%S")

_I_PATTERNS = re.compile(
    r"(?:^|_)(?:data_i|iq_i|ch_i|m_i|i_dec|i_data|i_val|[_\b]i[_\b]|i\[)", re.I
)
_Q_PATTERNS = re.compile(
    r"(?:^|_)(?:data_q|iq_q|ch_q|m_q|q_dec|q_data|q_val|[_\b]q[_\b]|q\[)", re.I
)


def _guess_iq_cols(headers: list[str]) -> tuple[str | None, str | None]:
    """Two-pass: prefer non-hex columns; fall back to hex columns if nothing else matches."""
    def _scan(hdrs):
        i_col = q_col = None
        for h in hdrs:
            if i_col is None and _I_PATTERNS.search(h):
                i_col = h
            if q_col is None and _Q_PATTERNS.search(h):
                q_col = h
        return i_col, q_col

    non_hex = [h for h in headers if "hex" not in h.lower()]
    i_col, q_col = _scan(non_hex)
    if i_col is None or q_col is None:
        i_col2, q_col2 = _scan(headers)
        i_col = i_col or i_col2
        q_col = q_col or q_col2
    return i_col, q_col


def _parse_value(s: str, bits: int, force_hex: bool = False) -> float:
    s = s.strip()
    if not s or s.lower() in ("x", "z", "nan", ""):
        return 0.0
    # hex?
    if force_hex or (re.match(r"^[0-9a-fA-F]+$", s) and any(c in s.lower() for c in "abcdef")):
        try:
            v = int(s, 16) & ((1 << bits) - 1)
            return float(v - (1 << bits) if v & (1 << (bits - 1)) else v)
        except ValueError:
            pass
    # 0x prefix
    if s.lower().startswith("0x"):
        v = int(s, 16) & ((1 << bits) - 1)
        return float(v - (1 << bits) if v & (1 << (bits - 1)) else v)
    # decimal (may already be signed)
    return float(s)


def _auto_bits(col_name: str, rows: list[dict], col: str) -> int:
    """Guess bit-width from column name or data range."""
    m = re.search(r"\[(\d+):0\]", col_name)
    if m:
        return int(m.group(1)) + 1
    # from data range
    try:
        vals = [float(r[col]) for r in rows[:200] if r.get(col, "").strip()]
        mx = max(abs(min(vals)), abs(max(vals)))
        for b in (8, 12, 15, 16, 23, 24, 31, 32):
            if mx < (1 << b):
                return b + 1
    except Exception:
        pass
    return 16


def rms(x: np.ndarray) -> float:
    return float(np.sqrt(np.mean(np.square(x))))


def symmetric_obw(freq_hz: np.ndarray, amp: np.ndarray, target: float) -> tuple[float, float]:
    power = np.square(amp)
    total = float(np.sum(power))
    if total <= 0.0:
        return 0.0, 0.0
    abs_freq = np.abs(freq_hz)
    for edge in np.unique(abs_freq):
        frac = float(np.sum(power[abs_freq <= edge]) / total)
        if frac >= target:
            return float(2.0 * edge / 1e6), frac
    return float(2.0 * abs_freq[-1] / 1e6), 1.0


def local_peaks(freq_hz: np.ndarray, amp: np.ndarray, dbfs: np.ndarray,
                count: int = 10) -> list[dict]:
    if len(freq_hz) < 3:
        return []
    bin_hz = abs(freq_hz[1] - freq_hz[0])
    mask = np.abs(freq_hz) > 2.0 * bin_hz
    idx = np.nonzero(mask)[0]
    cands = [int(k) for k in idx
             if 0 < k < len(dbfs) - 1
             and dbfs[k] >= dbfs[k - 1] and dbfs[k] >= dbfs[k + 1]]
    cands.sort(key=lambda k: dbfs[k], reverse=True)
    return [{"freq_hz": float(freq_hz[k]), "freq_mhz": float(freq_hz[k] / 1e6),
             "dbfs": float(dbfs[k]), "amp": float(amp[k]), "bin": int(k)}
            for k in cands[:count]]


def find_steady_start(i: np.ndarray, q: np.ndarray, skip: int) -> int:
    if skip > 0:
        return skip
    mag = np.abs(i) + np.abs(q)
    idx = np.nonzero(mag > 0)[0]
    return int(idx[0]) if len(idx) else 0


# ── core analysis ────────────────────────────────────────────────────────────

def analyse(
    input_path: Path,
    i_col: str | None,
    q_col: str | None,
    fs_hz: float,
    full_scale: float,
    skip: int,
    xlim_mhz: tuple[float, float] | None,
    out_dir: Path,
    prefix: str,
    title: str | None,
    bits: int,
) -> dict:
    # ── read file ──────────────────────────────────────────────────────────
    rows: list[dict] = []
    with input_path.open(newline="", encoding="utf-8-sig") as f:
        reader = csv.DictReader(f)
        headers = reader.fieldnames or []
        for row in reader:
            rows.append(row)

    if not rows:
        print("[fft_analyse] ERROR: no data rows found.")
        return {}

    # ── resolve column names ───────────────────────────────────────────────
    if i_col is None or q_col is None:
        gi, gq = _guess_iq_cols(headers)
        i_col = i_col or gi
        q_col = q_col or gq
    if i_col is None or q_col is None:
        print(f"[fft_analyse] Could not identify I/Q columns.")
        print(f"  Headers: {headers}")
        print("  Use --i-col and --q-col to specify them.")
        return {}

    print(f"[fft_analyse] Input : {input_path.name}  ({len(rows)} rows)")
    print(f"[fft_analyse] I col : {i_col}")
    print(f"[fft_analyse] Q col : {q_col}")
    print(f"[fft_analyse] fs    : {fs_hz/1e6:.4f} MHz   FS ref: {full_scale:.0f}")

    # ── parse IQ values ────────────────────────────────────────────────────
    # Vivado ILA CSV: column names contain "[N:0]" and values are bare hex (no 0x prefix)
    # e.g. column "AxC0_491_data_i[15:0]", value "0b6f" → must force hex parsing
    def _is_hex_col(col_name: str, sample_rows: list, col: str) -> bool:
        if "hex" in col_name.lower():
            return True
        if re.search(r"\[\d+:\d+\]", col_name):
            # Vivado ILA bit-range notation — peek at first non-empty value
            for r in sample_rows[:20]:
                v = r.get(col, "").strip()
                if v and v.lower() not in ("x", "z", "nan"):
                    # If it looks like hex (has a-f OR is exactly 4 hex digits)
                    if v.lower().startswith("0x") or (
                        re.match(r"^[0-9a-fA-F]+$", v)
                        and any(c in v.lower() for c in "abcdef")
                    ):
                        return True
                    break
        return False

    i_hex = _is_hex_col(i_col, rows, i_col)
    q_hex = _is_hex_col(q_col, rows, q_col)
    i_data = np.array([_parse_value(r.get(i_col, "0"), bits, force_hex=i_hex) for r in rows], dtype=np.float64)
    q_data = np.array([_parse_value(r.get(q_col, "0"), bits, force_hex=q_hex) for r in rows], dtype=np.float64)

    # skip startup / leading zeros
    s = find_steady_start(i_data, q_data, skip)
    if s > 0:
        print(f"[fft_analyse] Skipping {s} leading samples (zeros / startup)")
    i_data, q_data = i_data[s:], q_data[s:]

    if len(i_data) < 4:
        print("[fft_analyse] Too few samples after skip.")
        return {}

    # ── FFT ───────────────────────────────────────────────────────────────
    fft_n = 2 ** int(np.floor(np.log2(len(i_data))))
    i_mean, q_mean = float(np.mean(i_data)), float(np.mean(q_data))
    iq = (i_data[:fft_n] - i_mean) + 1j * (q_data[:fft_n] - q_mean)
    win = np.hanning(fft_n)
    cg = float(np.sum(win))               # coherent gain (for peak amplitude)
    win_power = float(np.sum(win ** 2))   # for PSD normalisation
    fft_data = np.fft.fftshift(np.fft.fft(iq * win, n=fft_n))
    amp = np.abs(fft_data) / cg           # peak amplitude (linear, for OBW / peaks)
    dbfs = 20.0 * np.log10((amp + np.finfo(np.float64).eps) / full_scale)
    # Power spectral density: dBFS/Hz  (normalised to full-scale power)
    psd = (np.abs(fft_data) ** 2) / (full_scale ** 2 * win_power * fs_hz)
    psd_dbfs_hz = 10.0 * np.log10(psd + np.finfo(np.float64).eps)
    freq_hz = (np.arange(fft_n) - fft_n / 2.0) * fs_hz / fft_n
    freq_mhz = freq_hz / 1e6

    obw90, frac90 = symmetric_obw(freq_hz, amp, 0.90)
    obw99, frac99 = symmetric_obw(freq_hz, amp, 0.99)
    peaks = local_peaks(freq_hz, amp, dbfs)
    mag_time = np.abs(i_data + 1j * q_data)

    print(f"[fft_analyse] FFT size: {fft_n}   90% OBW: {obw90:.4f} MHz   99% OBW: {obw99:.4f} MHz")
    if peaks:
        p = peaks[0]
        print(f"[fft_analyse] Top peak: {p['freq_mhz']:+.4f} MHz  {p['dbfs']:.2f} dBFS")

    result = dict(
        i_data=i_data, q_data=q_data, mag_time=mag_time, i_mean=i_mean, q_mean=q_mean,
        fft_n=fft_n, freq_hz=freq_hz, freq_mhz=freq_mhz, amp=amp, dbfs=dbfs,
        psd_dbfs_hz=psd_dbfs_hz,
        obw90=obw90, obw99=obw99, frac90=frac90, frac99=frac99, peaks=peaks,
        fs_hz=fs_hz, full_scale=full_scale, i_col=i_col, q_col=q_col,
        input_name=input_path.name, skip=s,
    )

    # ── write outputs ──────────────────────────────────────────────────────
    out_dir.mkdir(parents=True, exist_ok=True)

    csv_path = out_dir / f"{prefix}_{RUN_STAMP}_spectrum.csv"
    with csv_path.open("w", newline="") as f:
        w = csv.writer(f)
        w.writerow(["freq_hz", "freq_mhz", "amp", "dbfs"])
        for fh, a, d in zip(freq_hz, amp, dbfs):
            w.writerow([f"{fh:.6f}", f"{fh/1e6:.9f}", f"{a:.9f}", f"{d:.6f}"])

    lines = [
        f"IQ FFT analysis  —  {input_path.name}",
        f"Run: {RUN_STAMP}",
        f"I column  : {i_col}",
        f"Q column  : {q_col}",
        f"Fs        : {fs_hz/1e6:.4f} MHz",
        f"FS ref    : {full_scale:.0f}",
        f"Samples   : {len(i_data)}  (skipped {s})",
        f"FFT size  : {fft_n}",
        f"Bin spacing: {fs_hz/fft_n:.4f} Hz",
        f"I  min/max/mean/rms: {np.min(i_data):.1f} / {np.max(i_data):.1f} / "
            f"{i_mean:.4f} / {rms(i_data):.4f}",
        f"Q  min/max/mean/rms: {np.min(q_data):.1f} / {np.max(q_data):.1f} / "
            f"{q_mean:.4f} / {rms(q_data):.4f}",
        f"Mag RMS   : {rms(mag_time):.4f}",
        f"90% OBW   : {obw90:.6f} MHz  (fraction {frac90:.6f})",
        f"99% OBW   : {obw99:.6f} MHz  (fraction {frac99:.6f})",
        "",
        "Top spectral peaks:",
    ] + [
        f"  {p['freq_hz']:14.3f} Hz ({p['freq_mhz']:+.4f} MHz)  "
        f"{p['dbfs']:.2f} dBFS  amp {p['amp']:.2f}"
        for p in peaks
    ]
    (out_dir / f"{prefix}_{RUN_STAMP}_summary.txt").write_text("\n".join(lines) + "\n", encoding="utf-8")

    _plot(result, xlim_mhz, out_dir, prefix, title)

    print(f"[fft_analyse] Saved: {prefix}_{RUN_STAMP}_analysis.png  (and .svg, _spectrum.csv, _summary.txt)")
    return result


def _plot(result: dict, xlim_mhz, out_dir: Path, prefix: str, title: str | None) -> None:
    i_data = result["i_data"]
    q_data = result["q_data"]
    mag_time = result["mag_time"]
    freq_mhz = result["freq_mhz"]
    dbfs = result["dbfs"]
    fs_hz = result["fs_hz"]
    fft_n = result["fft_n"]

    if xlim_mhz is None:
        half = fs_hz / 2e6
        xlim_mhz = (-half, half)

    plt.rcParams.update({"font.size": 14, "axes.titlesize": 16, "axes.labelsize": 14,
                         "legend.fontsize": 12})
    fig, axes = plt.subplots(3, 1, figsize=(20, 14), constrained_layout=True)
    fig.patch.set_facecolor("white")
    main_title = title or f"IQ FFT Analysis — {result['input_name']}"
    fig.suptitle(main_title, fontsize=16)

    # panel 1 — IQ waveform
    count = min(512, len(i_data))
    axes[0].plot(i_data[:count], label=f"I  ({result['i_col']})", linewidth=1.4)
    axes[0].plot(q_data[:count], label=f"Q  ({result['q_col']})", linewidth=1.4)
    axes[0].set_title(f"I/Q waveform — first {count} samples (after skip={result['skip']})")
    axes[0].set_xlabel("Sample")
    axes[0].set_ylabel("Counts")
    axes[0].grid(True, alpha=0.3)
    axes[0].legend(loc="upper right")

    # panel 2 — magnitude
    axes[1].plot(mag_time, color="#28734f", linewidth=1.0, label="|I+jQ|")
    axes[1].axhline(rms(mag_time), color="#b22222", linestyle="--", linewidth=1.2,
                    label=f"RMS = {rms(mag_time):.1f}")
    axes[1].set_title("Complex magnitude |I + jQ|, all samples")
    axes[1].set_xlabel("Sample")
    axes[1].set_ylabel("Counts")
    axes[1].grid(True, alpha=0.3)
    axes[1].legend(loc="upper right")

    # panel 3 — spectrum (PSD)
    psd_dbfs_hz = result["psd_dbfs_hz"]
    axes[2].plot(freq_mhz, psd_dbfs_hz, color="#1f5a95", linewidth=1.0)
    obw90, obw99 = result["obw90"], result["obw99"]
    if obw90 > 0:
        axes[2].axvline(-obw90 / 2, color="#ff7f0e", ls="--", lw=1.3,
                        label=f"90% OBW: {obw90:.2f} MHz")
        axes[2].axvline( obw90 / 2, color="#ff7f0e", ls="--", lw=1.3)
    if obw99 > 0:
        axes[2].axvline(-obw99 / 2, color="#9467bd", ls="--", lw=1.3,
                        label=f"99% OBW: {obw99:.2f} MHz")
        axes[2].axvline( obw99 / 2, color="#9467bd", ls="--", lw=1.3)
    # annotate top peaks using PSD value at peak bin
    for p in result["peaks"][:5]:
        if xlim_mhz[0] <= p["freq_mhz"] <= xlim_mhz[1]:
            psd_val = float(psd_dbfs_hz[p["bin"]])
            axes[2].plot(p["freq_mhz"], psd_val, "o", color="#c23b22", ms=6)
            axes[2].annotate(f"{p['freq_mhz']:.2f} MHz\n{psd_val:.1f} dBFS/Hz",
                             (p["freq_mhz"], psd_val),
                             xytext=(0, 8), textcoords="offset points",
                             ha="center", fontsize=10)
    axes[2].set_title(f"{fft_n}-pt complex FFT, Hann window, DC removed, "
                      f"fs = {fs_hz/1e6:.2f} MHz, FS ref = {result['full_scale']:.0f}")
    axes[2].set_xlabel("Frequency (MHz)")
    axes[2].set_ylabel("PSD (dBFS/Hz)")
    axes[2].set_xlim(xlim_mhz)
    in_view = (freq_mhz >= xlim_mhz[0]) & (freq_mhz <= xlim_mhz[1])
    if np.any(in_view):
        y_top = max(-60.0, float(np.ceil(np.max(psd_dbfs_hz[in_view]) / 10) * 10 + 10))
    else:
        y_top = -60.0
    axes[2].set_ylim(y_top - 120, y_top)
    axes[2].grid(True, alpha=0.3)
    axes[2].legend(loc="lower right")

    fig.savefig(out_dir / f"{prefix}_{RUN_STAMP}_analysis.png", dpi=150)
    fig.savefig(out_dir / f"{prefix}_{RUN_STAMP}_analysis.svg")
    plt.close(fig)


# ── CLI ───────────────────────────────────────────────────────────────────────

def _auto_fs(path: Path) -> float:
    name = path.stem.lower()
    for pat, hz in [("491", 491.52e6), ("245", 245.76e6), ("240m", 240e6),
                    ("122", 122.88e6), ("120m", 120e6), ("61", 61.44e6),
                    ("250m", 250e6), ("125m", 125e6)]:
        if pat in name:
            return hz
    return 1.0   # unknown — caller should supply --fs


def main(argv: list[str] | None = None) -> None:
    parser = argparse.ArgumentParser(description="IQ FFT spectrum analyser")
    parser.add_argument("input", help="Input CSV / TSV / TXT file with IQ columns")
    parser.add_argument("--i-col",  default=None)
    parser.add_argument("--q-col",  default=None)
    parser.add_argument("--fs",     type=float, default=None,
                        help="Sample rate Hz, e.g. 120e6")
    parser.add_argument("--bits",   type=int,   default=None,
                        help="IQ bit width (determines full-scale ref)")
    parser.add_argument("--fs-val", type=float, default=None,
                        help="Full-scale amplitude override")
    parser.add_argument("--xlim",   default=None,
                        help="Frequency display range MHz, e.g. -10,10 or --xlim=-10,10")
    parser.add_argument("--out",    default=None,
                        help="Output directory (default: same as input)")
    parser.add_argument("--prefix", default=None,
                        help="Output filename prefix")
    parser.add_argument("--skip",   type=int, default=0,
                        help="Skip first N data rows (e.g. FIR startup)")
    parser.add_argument("--title",  default=None)
    # argparse treats "-70,70" as a flag; collapse "--xlim -N,M" → "--xlim=-N,M"
    raw = list(argv) if argv is not None else sys.argv[1:]
    for i in range(len(raw) - 1):
        if raw[i] == "--xlim" and re.match(r"^-\d", raw[i + 1]):
            raw[i:i + 2] = [f"--xlim={raw[i + 1]}"]
            break
    args = parser.parse_args(raw)

    inp = Path(args.input).resolve()
    if not inp.exists():
        sys.exit(f"[fft_analyse] File not found: {inp}")

    fs_hz    = args.fs     if args.fs     is not None else _auto_fs(inp)
    out_dir  = Path(args.out).resolve() if args.out else inp.parent
    prefix   = args.prefix or re.sub(r"[^\w]", "_", inp.stem)
    xlim_mhz = tuple(float(x) for x in args.xlim.split(",")) if args.xlim else None

    # determine full-scale
    if args.fs_val is not None:
        full_scale = args.fs_val
        bits = args.bits or 16
    elif args.bits is not None:
        bits = args.bits
        full_scale = float(1 << (bits - 1))
    else:
        bits = 16
        full_scale = 32768.0   # default: 16-bit

    analyse(
        input_path=inp,
        i_col=args.i_col,
        q_col=args.q_col,
        fs_hz=fs_hz,
        full_scale=full_scale,
        skip=args.skip,
        xlim_mhz=xlim_mhz,
        out_dir=out_dir,
        prefix=prefix,
        title=args.title,
        bits=bits,
    )


if __name__ == "__main__":
    main()