#!/usr/bin/env python3
"""Audit per-chain presence and amplitude stability in a router CSI capture.

Input container:
    <float64 timestamp> <uint16 payload-length> <payload>

24-byte CSI payload header:
    <I timestamp_us 6s transmitter_mac b rssi B snr B bandwidth
     B primary_channel B receive_mode B sequence/padding
     H subcarrier_count H tx_chain H rx_chain H chain_info>

The 8-chain mode in this capture uses every tx/rx pair once. `chain_info` bit 15
marks the final record in one closed measurement group.
"""

from __future__ import annotations

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

import numpy as np


RECORD = struct.Struct("<dH")
HEADER = struct.Struct("<I6sbBBBBBHHHH")
EXPECTED_CHAINS = [(tx, rx) for tx in range(2) for rx in range(4)]


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


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


def parse_capture(path: Path) -> dict:
    data = path.read_bytes()
    offset = 0
    outer_records = 0
    malformed_records = 0
    pending: list[dict] = []
    groups: list[dict] = []

    while offset + RECORD.size <= len(data):
        outer_ts, 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_ts,
            station,
            rssi,
            snr,
            bandwidth,
            channel,
            receive_mode,
            sequence,
            subcarrier_count,
            tx,
            rx,
            chain_info,
        ) = header

        iq_offset = payload_start + HEADER.size
        available_pairs = max(0, (payload_end - iq_offset) // 4)
        pairs = min(subcarrier_count, available_pairs)
        iq = np.frombuffer(
            data,
            dtype="<i2",
            count=pairs * 2,
            offset=iq_offset,
        ).astype(np.float64)
        amplitude = np.hypot(iq[0::2], iq[1::2]) if iq.size else np.empty(0)
        complete = pairs == 256 and amplitude.size == 256 and np.any(amplitude)

        pending.append(
            {
                "outer_ts": outer_ts,
                "payload_ts": payload_ts,
                "station": station,
                "rssi": rssi,
                "snr": snr,
                "bandwidth": bandwidth,
                "channel": channel,
                "receive_mode": receive_mode,
                "sequence": sequence,
                "subcarrier_count": subcarrier_count,
                "tx": tx,
                "rx": rx,
                "chain_info": chain_info,
                "pairs": pairs,
                "complete": complete,
                "amplitude": amplitude,
            }
        )
        outer_records += 1

        if chain_info & 0x8000:
            groups.append({"records": pending})
            pending = []
        offset = payload_end

    return {
        "outer_records": outer_records,
        "malformed_records": malformed_records,
        "incomplete_tail_records": len(pending),
        "groups": groups,
    }


def summarize(parsed: dict) -> tuple[dict, list[dict]]:
    groups = parsed["groups"]
    closed_groups = len(groups)
    group_sizes = collections.Counter()
    complete_groups = 0
    same_rssi_groups = 0
    same_snr_groups = 0
    chain_records: dict[tuple[int, int], list[dict]] = {
        key: [] for key in EXPECTED_CHAINS
    }
    chain_presence = collections.Counter()
    group_mean_amplitude: dict[tuple[int, int], list[float]] = {
        key: [] for key in EXPECTED_CHAINS
    }

    for group in groups:
        records = group["records"]
        group_sizes[len(records)] += 1
        if records and all(record["complete"] for record in records):
            complete_groups += 1
        if len({record["rssi"] for record in records}) == 1:
            same_rssi_groups += 1
        if len({record["snr"] for record in records}) == 1:
            same_snr_groups += 1

        seen = set()
        for record in records:
            key = (record["tx"], record["rx"])
            if key not in EXPECTED_CHAINS:
                continue
            if key not in seen:
                chain_presence[key] += 1
                seen.add(key)
            if not record["complete"]:
                continue
            amplitude = record["amplitude"]
            mean_amplitude = float(amplitude.mean())
            within_record_cv = float(amplitude.std() / (mean_amplitude + 1e-12))
            chain_records[key].append(
                {
                    "mean_amplitude": mean_amplitude,
                    "within_record_cv": within_record_cv,
                    "rssi": record["rssi"],
                    "snr": record["snr"],
                }
            )
            group_mean_amplitude[key].append(mean_amplitude)

    chain_rows = []
    summary_chains = []
    for tx, rx in EXPECTED_CHAINS:
        key = (tx, rx)
        records = chain_records[key]
        means = np.asarray(group_mean_amplitude[key], dtype=np.float64)
        mean_amplitude = np.asarray(
            [record["mean_amplitude"] for record in records], dtype=np.float64
        )
        within_cv = np.asarray(
            [record["within_record_cv"] for record in records], dtype=np.float64
        )
        temporal_cv = float(means.std() / (means.mean() + 1e-12))
        row = {
            "chain": f"tx{tx}/rx{rx}",
            "tx": tx,
            "rx": rx,
            "group_presence": chain_presence[key],
            "group_presence_fraction": rounded(chain_presence[key] / closed_groups),
            "complete_record_count": len(records),
            "median_mean_amplitude": rounded(np.median(mean_amplitude)),
            "median_within_record_cv": rounded(np.median(within_cv)),
            "temporal_cv": rounded(temporal_cv),
            "group_mean_amplitude_p05": rounded(np.percentile(means, 5)),
            "group_mean_amplitude_p50": rounded(np.percentile(means, 50)),
            "group_mean_amplitude_p95": rounded(np.percentile(means, 95)),
        }
        chain_rows.append(row)
        summary_chains.append(
            {
                **row,
                "median_rssi_dbm": int(np.median([r["rssi"] for r in records])),
                "median_snr_db": int(np.median([r["snr"] for r in records])),
            }
        )

    summary = {
        "source": {
            "file": "",
            "size_bytes": 0,
            "sha256": "",
        },
        "capture": {
            "outer_records": parsed["outer_records"],
            "closed_groups": closed_groups,
            "complete_groups": complete_groups,
            "incomplete_tail_records": parsed["incomplete_tail_records"],
            "malformed_records": parsed["malformed_records"],
            "duration_s": rounded(
                groups[-1]["records"][-1]["outer_ts"] - groups[0]["records"][0]["outer_ts"]
            )
            if groups
            else 0,
        },
        "groups": {
            "eight_chain": group_sizes[8],
            "four_chain": group_sizes[4],
            "eight_chain_fraction": rounded(group_sizes[8] / closed_groups),
            "four_chain_fraction": rounded(group_sizes[4] / closed_groups),
        },
        "metadata": {
            "same_rssi_groups": same_rssi_groups,
            "same_snr_groups": same_snr_groups,
            "selects_chain": False,
            "explanation": (
                "RSSI and SNR were identical across all chains inside every closed group. "
                "They describe the reception context and cannot rank individual tx/rx chains."
            ),
        },
        "chains": summary_chains,
        "method": {
            "record_layout": "float64 outer timestamp + uint16 length + 24-byte CSI header + int16 I/Q",
            "group_boundary": "chain_info bit 15",
            "complete_record": "256 subcarrier pairs and non-zero I/Q",
            "amplitude": "sqrt(I^2 + Q^2) for each subcarrier pair",
            "within_record_cv": "standard deviation divided by mean across 256 subcarriers",
            "temporal_cv": "standard deviation divided by mean of each group's mean amplitude",
        },
        "limitations": [
            "Neither capture has a controlled noise-only baseline or human-event ground truth.",
            "Low temporal variance is not automatically better: it may also indicate weak motion sensitivity.",
            "Median amplitude depends on radio gain and changed between the two captures from the same device.",
            "Chain ordering is hardware- and firmware-specific; tx/rx labels cannot be assumed to represent physical antenna positions.",
            "Two captures from one device and one home do not establish behaviour across rooms, postures, sessions or traffic profiles.",
            "This audit supports chain health gating and experiment design, not breathing, fall or identity inference.",
        ],
    }
    return summary, chain_rows


def write_csv(path: Path, rows: list[dict]) -> None:
    fields = [
        "chain",
        "tx",
        "rx",
        "group_presence",
        "group_presence_fraction",
        "complete_record_count",
        "median_mean_amplitude",
        "median_within_record_cv",
        "temporal_cv",
        "group_mean_amplitude_p05",
        "group_mean_amplitude_p50",
        "group_mean_amplitude_p95",
    ]
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=fields)
        writer.writeheader()
        writer.writerows(rows)


def pearson(values_a: list[float], values_b: list[float]) -> float:
    a = np.asarray(values_a, dtype=np.float64)
    b = np.asarray(values_b, dtype=np.float64)
    return float(np.corrcoef(a, b)[0, 1])


def ranks(values: dict[str, float], descending: bool) -> dict[str, int]:
    ordered = sorted(values.items(), key=lambda item: item[1], reverse=descending)
    return {chain: index + 1 for index, (chain, _) in enumerate(ordered)}


def build_comparison(primary: dict, comparison: dict) -> tuple[dict, list[dict]]:
    primary_by_chain = {row["chain"]: row for row in primary["chains"]}
    comparison_by_chain = {row["chain"]: row for row in comparison["chains"]}
    chains = [row["chain"] for row in primary["chains"]]

    amplitude_primary = {
        chain: primary_by_chain[chain]["median_mean_amplitude"] for chain in chains
    }
    amplitude_comparison = {
        chain: comparison_by_chain[chain]["median_mean_amplitude"] for chain in chains
    }
    cv_primary = {chain: primary_by_chain[chain]["temporal_cv"] for chain in chains}
    cv_comparison = {chain: comparison_by_chain[chain]["temporal_cv"] for chain in chains}
    amplitude_rank_primary = ranks(amplitude_primary, descending=True)
    amplitude_rank_comparison = ranks(amplitude_comparison, descending=True)
    cv_rank_primary = ranks(cv_primary, descending=False)
    cv_rank_comparison = ranks(cv_comparison, descending=False)

    rows = []
    for chain in chains:
        rows.append(
            {
                "chain": chain,
                "presence_10m": primary_by_chain[chain]["group_presence"],
                "presence_60s": comparison_by_chain[chain]["group_presence"],
                "median_mean_amplitude_10m": amplitude_primary[chain],
                "median_mean_amplitude_60s": amplitude_comparison[chain],
                "amplitude_rank_10m": amplitude_rank_primary[chain],
                "amplitude_rank_60s": amplitude_rank_comparison[chain],
                "temporal_cv_10m": cv_primary[chain],
                "temporal_cv_60s": cv_comparison[chain],
                "temporal_cv_rank_10m": cv_rank_primary[chain],
                "temporal_cv_rank_60s": cv_rank_comparison[chain],
                "delta_amplitude": rounded(amplitude_comparison[chain] - amplitude_primary[chain]),
                "delta_temporal_cv": rounded(cv_comparison[chain] - cv_primary[chain]),
            }
        )

    top_amplitude_primary = min(amplitude_rank_primary, key=amplitude_rank_primary.get)
    top_amplitude_comparison = min(amplitude_rank_comparison, key=amplitude_rank_comparison.get)
    lowest_cv_primary = min(cv_rank_primary, key=cv_rank_primary.get)
    lowest_cv_comparison = min(cv_rank_comparison, key=cv_rank_comparison.get)
    result = {
        "source": comparison["source"],
        "capture": comparison["capture"],
        "groups": comparison["groups"],
        "metadata": comparison["metadata"],
        "top_amplitude_chain_10m": top_amplitude_primary,
        "top_amplitude_chain_60s": top_amplitude_comparison,
        "lowest_temporal_cv_chain_10m": lowest_cv_primary,
        "lowest_temporal_cv_chain_60s": lowest_cv_comparison,
        "pearson_amplitude": rounded(
            pearson(
                [amplitude_primary[chain] for chain in chains],
                [amplitude_comparison[chain] for chain in chains],
            )
        ),
        "pearson_temporal_cv": rounded(
            pearson(
                [cv_primary[chain] for chain in chains],
                [cv_comparison[chain] for chain in chains],
            )
        ),
    }
    return result, rows


def write_comparison_csv(path: Path, rows: list[dict]) -> None:
    fields = [
        "chain",
        "presence_10m",
        "presence_60s",
        "median_mean_amplitude_10m",
        "median_mean_amplitude_60s",
        "amplitude_rank_10m",
        "amplitude_rank_60s",
        "temporal_cv_10m",
        "temporal_cv_60s",
        "temporal_cv_rank_10m",
        "temporal_cv_rank_60s",
        "delta_amplitude",
        "delta_temporal_cv",
    ]
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=fields)
        writer.writeheader()
        writer.writerows(rows)


def write_svg(path: Path, summary: dict) -> None:
    width = 1200
    height = 650
    left_x = 90
    left_y = 155
    left_w = 470
    left_h = 375
    right_x = 720
    right_y = 155
    right_w = 390
    right_h = 375
    chains = summary["chains"]
    max_presence = summary["capture"]["closed_groups"]
    amplitude_values = [row["median_mean_amplitude"] for row in chains]
    amplitude_max = max(amplitude_values) * 1.18
    temporal_values = [row["temporal_cv"] for row in chains]
    temporal_max = max(temporal_values) * 1.18
    colors = ["#22c55e" if row["group_presence"] == max_presence else "#f59e0b" for row in chains]

    parts = [
        f'<svg xmlns="http://www.w3.org/2000/svg" width="{width}" height="{height}" '
        f'viewBox="0 0 {width} {height}" role="img" aria-labelledby="title desc">',
        "<title id=\"title\">GL-MT6000 CSI chain presence and temporal variability</title>",
        "<desc id=\"desc\">Bars show group presence for eight tx/rx chains. "
        "The scatter plot compares median mean amplitude with temporal coefficient of variation. "
        "The plot is diagnostic and does not identify which chain is physically best.</desc>",
        '<rect width="1200" height="650" fill="#0b0f17"/>',
        '<text x="60" y="54" fill="#f8fafc" font-family="Arial, sans-serif" font-size="28" '
        'font-weight="700">CSI chain presence and temporal variability</text>',
        '<text x="60" y="84" fill="#94a3b8" font-family="Arial, sans-serif" font-size="16">'
        "10-minute GL-MT6000 capture, 7,609 closed groups; low variance alone is not a quality verdict.</text>",
        f'<rect x="{left_x}" y="{left_y}" width="{left_w}" height="{left_h}" fill="#111827" '
        'stroke="#334155"/>',
        f'<rect x="{right_x}" y="{right_y}" width="{right_w}" height="{right_h}" fill="#111827" '
        'stroke="#334155"/>',
        f'<text x="{left_x}" y="128" fill="#e2e8f0" font-family="Arial, sans-serif" '
        'font-size="18">Chain presence in closed groups</text>',
        f'<text x="{right_x}" y="128" fill="#e2e8f0" font-family="Arial, sans-serif" '
        'font-size="18">Amplitude vs temporal CV</text>',
    ]

    bar_gap = 18
    bar_width = 42
    for index, row in enumerate(chains):
        bar_h = left_h * row["group_presence"] / max_presence
        x = left_x + 22 + index * ((left_w - 44) / 8)
        y = left_y + left_h - bar_h
        parts.append(
            f'<rect x="{x:.1f}" y="{y:.1f}" width="{bar_width}" height="{bar_h:.1f}" '
            f'fill="{colors[index]}" rx="4"/>'
        )
        parts.append(
            f'<text x="{x + bar_width / 2:.1f}" y="{y - 8:.1f}" text-anchor="middle" '
            f'fill="#e2e8f0" font-family="Arial, sans-serif" font-size="11">'
            f'{row["group_presence"]:,}</text>'
        )
        parts.append(
            f'<text x="{x + bar_width / 2:.1f}" y="{left_y + left_h + 23}" text-anchor="end" '
            f'transform="rotate(-35 {x + bar_width / 2:.1f} {left_y + left_h + 23})" '
            f'fill="#cbd5e1" font-family="Arial, sans-serif" font-size="11">{row["chain"]}</text>'
        )

    parts.extend(
        [
            f'<line x1="{right_x + 55}" y1="{right_y + right_h - 45}" '
            f'x2="{right_x + right_w - 25}" y2="{right_y + right_h - 45}" stroke="#64748b"/>',
            f'<line x1="{right_x + 55}" y1="{right_y + 28}" '
            f'x2="{right_x + 55}" y2="{right_y + right_h - 45}" stroke="#64748b"/>',
            f'<text x="{right_x + right_w / 2}" y="{right_y + right_h + 35}" text-anchor="middle" '
            'fill="#cbd5e1" font-family="Arial, sans-serif" font-size="13">'
            "Median mean amplitude (arbitrary CSI scale)</text>",
            f'<text x="{right_x - 20}" y="{right_y + right_h / 2}" text-anchor="middle" '
            f'transform="rotate(-90 {right_x - 20} {right_y + right_h / 2})" '
            'fill="#cbd5e1" font-family="Arial, sans-serif" font-size="13">Temporal CV</text>',
            f'<text x="{right_x + 60}" y="{right_y + right_h - 25}" fill="#94a3b8" '
            'font-family="Arial, sans-serif" font-size="11">0</text>',
            f'<text x="{right_x + right_w - 75}" y="{right_y + right_h - 25}" fill="#94a3b8" '
            f'font-family="Arial, sans-serif" font-size="11">{amplitude_max:.0f}</text>',
            f'<text x="{right_x + 58}" y="{right_y + 35}" fill="#94a3b8" '
            f'font-family="Arial, sans-serif" font-size="11">{temporal_max:.2f}</text>',
        ]
    )

    label_offsets = {
        "tx0/rx0": (-66, 4, "end"),
        "tx0/rx1": (12, 4, "start"),
        "tx0/rx2": (-64, -8, "end"),
        "tx0/rx3": (12, 4, "start"),
        "tx1/rx0": (12, 4, "start"),
        "tx1/rx1": (12, 4, "start"),
        "tx1/rx2": (12, -8, "start"),
        "tx1/rx3": (12, 18, "start"),
    }
    for index, row in enumerate(chains):
        x = right_x + 55 + (right_w - 80) * row["median_mean_amplitude"] / amplitude_max
        y = right_y + right_h - 45 - (right_h - 73) * row["temporal_cv"] / temporal_max
        parts.append(
            f'<circle cx="{x:.1f}" cy="{y:.1f}" r="8" fill="{colors[index]}" '
            'stroke="#f8fafc" stroke-width="1.5"/>'
        )
        dx, dy, anchor = label_offsets.get(row["chain"], (12, 4, "start"))
        parts.append(
            f'<text x="{x + dx:.1f}" y="{y + dy:.1f}" text-anchor="{anchor}" fill="#e2e8f0" '
            f'font-family="Arial, sans-serif" font-size="10.5">{row["chain"]}</text>'
        )

    parts.extend(
        [
            '<circle cx="90" cy="605" r="6" fill="#22c55e"/>',
            '<text x="104" y="610" fill="#cbd5e1" font-family="Arial, sans-serif" font-size="12">'
            "Present in every closed group</text>",
            '<circle cx="330" cy="605" r="6" fill="#f59e0b"/>',
            '<text x="344" y="610" fill="#cbd5e1" font-family="Arial, sans-serif" font-size="12">'
            "Intermittent four-chain groups</text>",
            "</svg>",
        ]
    )
    path.write_text("\n".join(parts) + "\n", encoding="utf-8")


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("capture", type=Path)
    parser.add_argument("--summary", type=Path, required=True)
    parser.add_argument("--csv", type=Path, required=True)
    parser.add_argument("--img", type=Path, required=True)
    parser.add_argument("--compare-capture", type=Path)
    parser.add_argument("--comparison-csv", type=Path)
    args = parser.parse_args()
    if bool(args.compare_capture) != bool(args.comparison_csv):
        raise SystemExit("--compare-capture and --comparison-csv must be supplied together")

    parsed = parse_capture(args.capture)
    if not parsed["groups"]:
        raise SystemExit("no closed measurement groups found")
    summary, rows = summarize(parsed)
    summary["source"] = {
        "file": args.capture.name,
        "size_bytes": args.capture.stat().st_size,
        "sha256": sha256_file(args.capture),
    }
    comparison_rows = None
    if args.compare_capture:
        compared = parse_capture(args.compare_capture)
        if not compared["groups"]:
            raise SystemExit("no closed measurement groups found in comparison capture")
        comparison_summary, _ = summarize(compared)
        comparison_summary["source"] = {
            "file": args.compare_capture.name,
            "size_bytes": args.compare_capture.stat().st_size,
            "sha256": sha256_file(args.compare_capture),
        }
        summary["comparison"], comparison_rows = build_comparison(summary, comparison_summary)

    args.summary.parent.mkdir(parents=True, exist_ok=True)
    args.csv.parent.mkdir(parents=True, exist_ok=True)
    args.img.parent.mkdir(parents=True, exist_ok=True)
    args.summary.write_text(json.dumps(summary, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
    write_csv(args.csv, rows)
    if comparison_rows is not None:
        write_comparison_csv(args.comparison_csv, comparison_rows)
    write_svg(args.img, summary)

    print(f"closed_groups: {summary['capture']['closed_groups']}")
    print(f"complete_groups: {summary['capture']['complete_groups']}")
    print(f"eight_chain_groups: {summary['groups']['eight_chain']}")
    print(f"four_chain_groups: {summary['groups']['four_chain']}")
    for row in summary["chains"]:
        print(
            f"{row['chain']}: presence={row['group_presence']} "
            f"median_amplitude={row['median_mean_amplitude']:.3f} "
            f"temporal_cv={row['temporal_cv']:.6f}"
        )
    if "comparison" in summary:
        comparison = summary["comparison"]
        print(
            "comparison: "
            f"groups={comparison['capture']['closed_groups']} "
            f"eight_chain={comparison['groups']['eight_chain']} "
            f"four_chain={comparison['groups']['four_chain']} "
            f"amplitude_r={comparison['pearson_amplitude']:.6f} "
            f"cv_r={comparison['pearson_temporal_cv']:.6f}"
        )
    return 0


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