#!/usr/bin/env python3
"""Audit timing, grouping and link quality in a router CSI capture.

The input is the raw container produced by this project's capture tools:
    <float64 timestamp> <uint16 payload-length> <payload>

The 24-byte payload header is:
    <I  timestamp_us
     6s transmitter MAC
     b  RSSI
     B  SNR
     B  bandwidth code
     B  primary-channel index
     B  receive mode
     B  sequence epoch / padding
     H  subcarrier count
     H  Tx chain
     H  Rx chain
     H  chain_info (bit 15 closes one measurement)
"""

from __future__ import annotations

import argparse
import collections
import hashlib
import json
import struct
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np


RECORD = struct.Struct("<dH")
HEADER = struct.Struct("<I6sbBBBBBHHHH")


def read_groups(path: Path):
    data = path.read_bytes()
    offset = 0
    pending = []
    groups = []
    outer_records = 0
    malformed_records = 0
    subcarrier_counts = collections.Counter()
    bandwidth_counts = collections.Counter()
    station_counts = collections.Counter()

    while offset + RECORD.size <= len(data):
        timestamp, length = RECORD.unpack_from(data, offset)
        payload_start = offset + RECORD.size
        payload_end = payload_start + length
        if payload_end > len(data) or length < HEADER.size:
            malformed_records += 1
            break

        header = HEADER.unpack_from(data, payload_start)
        payload_timestamp, station, rssi, snr, bandwidth, channel, mode, epoch, count, tx, rx, chain = header
        iq_offset = payload_start + HEADER.size
        available_pairs = max(0, (payload_end - iq_offset) // 4)
        pairs = min(count, available_pairs)
        iq = np.frombuffer(data, dtype="<i2", count=pairs * 2, offset=iq_offset)
        nonzero_iq = bool(iq.size and np.any(iq))

        outer_records += 1
        subcarrier_counts[count] += 1
        bandwidth_counts[bandwidth] += 1
        station_counts[station.hex(":")] += 1
        pending.append(
            {
                "outer_ts": timestamp,
                "payload_ts": payload_timestamp,
                "rssi": rssi,
                "snr": snr,
                "subcarriers": count,
                "tx": tx,
                "rx": rx,
                "chain": chain,
                "nonzero_iq": nonzero_iq,
            }
        )

        if chain & 0x8000:
            groups.append(pending)
            pending = []
        offset = payload_end

    return {
        "data": data,
        "groups": groups,
        "incomplete_tail_records": len(pending),
        "outer_records": outer_records,
        "malformed_records": malformed_records,
        "subcarrier_counts": subcarrier_counts,
        "bandwidth_counts": bandwidth_counts,
        "station_counts": station_counts,
    }


def rounded(value: float, digits: int = 6) -> float:
    return round(float(value), digits)


def percentile_dict(values: np.ndarray) -> dict:
    if not values.size:
        return {}
    points = np.percentile(values, [0, 10, 25, 50, 75, 90, 95, 99, 99.9, 100])
    names = ["min", "p10", "p25", "p50", "p75", "p90", "p95", "p99", "p999", "max"]
    return {name: rounded(value) for name, value in zip(names, points)}


def configure_plotting():
    plt.rcParams.update(
        {
            "figure.facecolor": "#0b0f17",
            "axes.facecolor": "#111827",
            "savefig.facecolor": "#0b0f17",
            "axes.edgecolor": "#475569",
            "axes.labelcolor": "#e5e7eb",
            "xtick.color": "#cbd5e1",
            "ytick.color": "#cbd5e1",
            "text.color": "#e5e7eb",
            "grid.color": "#334155",
            "font.size": 10,
        }
    )


def plot_distribution(intervals: np.ndarray, output: Path):
    figure, axes = plt.subplots(2, 1, figsize=(11, 8))

    clipped = intervals[intervals <= 1.0]
    axes[0].hist(clipped * 1000.0, bins=100, color="#38bdf8", alpha=0.9)
    axes[0].set_yscale("log")
    axes[0].set_xlabel("Group-to-group interval (ms; values <= 1000 ms)")
    axes[0].set_ylabel("Count (log scale)")
    axes[0].set_title(
        f"CSI measurement interval distribution, n={intervals.size:,} "
        f"({intervals.size - clipped.size:,} intervals > 1000 ms omitted)"
    )
    axes[0].grid(True, alpha=0.3)

    ordered = np.sort(intervals)
    cdf = np.arange(1, ordered.size + 1) / ordered.size
    axes[1].plot(ordered, cdf, color="#fb923c", linewidth=2)
    for percentile, color in [(95, "#a3e635"), (99, "#f472b6")]:
        value = float(np.percentile(intervals, percentile))
        axes[1].axvline(value, color=color, linestyle="--", linewidth=1.2)
        axes[1].text(
            value,
            0.08,
            f" p{percentile}={value:.3f}s",
            color=color,
            rotation=90,
            va="bottom",
        )
    axes[1].axvline(1.0, color="#f87171", linestyle=":", linewidth=1.3)
    axes[1].set_xlim(0, min(2.5, max(1.1, float(intervals.max()) + 0.05)))
    axes[1].set_ylim(0, 1.02)
    axes[1].set_xlabel("Group-to-group interval (s)")
    axes[1].set_ylabel("Empirical CDF")
    axes[1].set_title("Measurement-interval CDF and long-gap boundary")
    axes[1].grid(True, alpha=0.3)

    figure.tight_layout()
    figure.savefig(output, dpi=150)
    plt.close(figure)


def plot_timeline(
    elapsed: np.ndarray,
    intervals: np.ndarray,
    chain_counts: np.ndarray,
    rssi: np.ndarray,
    snr: np.ndarray,
    output: Path,
):
    figure, axes = plt.subplots(3, 1, figsize=(13, 10), sharex=True)

    axes[0].semilogy(elapsed[1:], intervals, color="#38bdf8", linewidth=0.8)
    axes[0].axhline(1.0, color="#f87171", linestyle="--", linewidth=1.1, label="1 s")
    axes[0].set_ylabel("Interval (s, log)")
    axes[0].set_title("CSI group timing across the 10-minute capture")
    axes[0].grid(True, alpha=0.3)
    axes[0].legend(loc="upper right")

    axes[1].step(elapsed, chain_counts, where="post", color="#a3e635", linewidth=1.0)
    axes[1].set_ylabel("Chain records / group")
    axes[1].set_ylim(2.5, 8.5)
    axes[1].set_yticks([4, 8])
    axes[1].grid(True, alpha=0.3)

    axes[2].plot(elapsed, rssi, color="#fb923c", linewidth=0.8, label="RSSI (dBm)")
    axes[2].plot(elapsed, snr, color="#c084fc", linewidth=0.8, label="SNR (dB)")
    axes[2].set_xlabel("Elapsed capture time (s)")
    axes[2].set_ylabel("RSSI / SNR")
    axes[2].grid(True, alpha=0.3)
    axes[2].legend(loc="upper right")

    figure.tight_layout()
    figure.savefig(output, dpi=150)
    plt.close(figure)


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("capture", type=Path)
    parser.add_argument("--summary", type=Path, required=True)
    parser.add_argument("--img-dir", type=Path, required=True)
    args = parser.parse_args()

    parsed = read_groups(args.capture)
    groups = parsed["groups"]
    if not groups:
        raise SystemExit("no closed measurement groups found")

    # A usable time-series group has only 256-subcarrier records and at least
    # one non-zero I/Q block. This intentionally quarantines the single 64-
    # subcarrier transition group in the public capture.
    usable = [
        group
        for group in groups
        if all(record["subcarriers"] == 256 for record in group)
        and any(record["nonzero_iq"] for record in group)
    ]
    if len(usable) < 2:
        raise SystemExit("not enough usable groups")

    elapsed = np.asarray([group[0]["outer_ts"] for group in usable], dtype=np.float64)
    elapsed -= elapsed[0]
    intervals = np.diff(elapsed)
    chain_counts = np.asarray([len(group) for group in usable], dtype=np.int16)
    rssi = np.asarray([group[0]["rssi"] for group in usable], dtype=np.int16)
    snr = np.asarray([group[0]["snr"] for group in usable], dtype=np.int16)
    duration = float(elapsed[-1] - elapsed[0])

    group_record_counts = collections.Counter(len(group) for group in groups)
    usable_chain_counts = collections.Counter(int(value) for value in chain_counts)
    chain_counts_by_subcarriers = collections.Counter(
        (group[0]["subcarriers"], len(group)) for group in groups
    )
    thresholds = [0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0]

    summary = {
        "source_file": args.capture.name,
        "source_size_bytes": len(parsed["data"]),
        "source_sha256": hashlib.sha256(parsed["data"]).hexdigest(),
        "method": {
            "group_boundary": "chain_info bit 15",
            "group_timestamp": "outer timestamp of the first chain record in the group",
            "usable_group": "all records have 256 subcarriers and at least one I/Q block is non-zero",
            "long_gap_thresholds_s": thresholds,
        },
        "outer_records": parsed["outer_records"],
        "closed_groups": len(groups),
        "usable_groups": len(usable),
        "quarantined_groups": len(groups) - len(usable),
        "incomplete_tail_records": parsed["incomplete_tail_records"],
        "malformed_records": parsed["malformed_records"],
        "duration_s": rounded(duration),
        "rates_hz": {
            "chain_records_per_duration": rounded(parsed["outer_records"] / duration),
            "closed_groups_per_duration": rounded(len(groups) / duration),
            "mean_interval_rate": rounded(1.0 / float(np.mean(intervals))),
        },
        "group_record_counts": {str(k): v for k, v in sorted(group_record_counts.items())},
        "usable_chain_counts": {str(k): v for k, v in sorted(usable_chain_counts.items())},
        "chain_counts_by_subcarriers": {
            f"{subcarriers}sc_x{chains}": count
            for (subcarriers, chains), count in sorted(chain_counts_by_subcarriers.items())
        },
        "subcarrier_record_counts": {
            str(k): v for k, v in sorted(parsed["subcarrier_counts"].items())
        },
        "bandwidth_record_counts": {
            str(k): v for k, v in sorted(parsed["bandwidth_counts"].items())
        },
        "stations": dict(parsed["station_counts"]),
        "intervals_s": {
            "count": int(intervals.size),
            **percentile_dict(intervals),
        },
        "long_gaps": {
            f">{threshold:g}s": int(np.sum(intervals > threshold))
            for threshold in thresholds
        },
        "rssi_dbm": {
            "min": int(rssi.min()),
            "mean": rounded(rssi.mean(), 3),
            "max": int(rssi.max()),
        },
        "snr_db": {
            "min": int(snr.min()),
            "mean": rounded(snr.mean(), 3),
            "max": int(snr.max()),
        },
        "interpretation_limits": [
            "Average rate does not establish a uniform time grid.",
            "The capture contains no human-event ground truth.",
            "A long gap can create spectral leakage or false peaks if resampled without a mask.",
        ],
    }

    args.summary.parent.mkdir(parents=True, exist_ok=True)
    args.summary.write_text(json.dumps(summary, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")

    args.img_dir.mkdir(parents=True, exist_ok=True)
    configure_plotting()
    plot_distribution(intervals, args.img_dir / "csi-quality-distribution.png")
    plot_timeline(
        elapsed,
        intervals,
        chain_counts,
        rssi,
        snr,
        args.img_dir / "csi-quality-timeline.png",
    )

    print(json.dumps(summary, indent=2, ensure_ascii=False))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())