#!/usr/bin/env python3
import re
import sys
import argparse
from pathlib import Path
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

LOG_PATTERN = re.compile(r"(\d+)\s+(TX|RX)\s+0x([0-9A-Fa-f]+)\s+\d+\s+((?:[0-9A-Fa-f]{2}\s+)+)")

def parse_trace(path: Path) -> pd.DataFrame:
    rows = []
    with open(path, "r", errors="ignore") as f:
        for line in f:
            m = LOG_PATTERN.match(line)
            if not m:
                continue
            ts = int(m.group(1))
            direction = m.group(2)
            can_id = int(m.group(3), 16)
            data = [int(x, 16) for x in m.group(4).split() if x.strip()]
            rows.append((ts, direction, can_id, data))
    df = pd.DataFrame(rows, columns=["ts","dir","id","data"])
    if df.empty:
        return df
    df["time_s"] = (df["ts"] - df["ts"].min())/1000.0
    return df

def be16(b):
    return (b[0]<<8) | b[1]

def le16(b):
    return b[0] | (b[1]<<8)

def main():
    ap = argparse.ArgumentParser(description="Per-ID explorer: generate plots for 8-bit and 16-bit combinations")
    ap.add_argument("trace", help="Single-ID .trace file (from can_split_by_id.py)")
    ap.add_argument("--outdir", default=None, help="Output directory; default: <trace>_explore")
    ap.add_argument("--prefix", default="viz", help="File prefix for exports")
    ap.add_argument("--rx-only", action="store_true", help="Use only RX frames")
    args = ap.parse_args()

    trace = Path(args.trace)
    df = parse_trace(trace)
    if df.empty:
        print("No data in trace.", file=sys.stderr)
        sys.exit(1)

    if args.rx_only:
        df = df[df["dir"]=="RX"].copy()
        if df.empty:
            print("No RX frames.", file=sys.stderr)
            sys.exit(2)

    outdir = Path(args.outdir) if args.outdir else trace.with_suffix("").parent / (trace.stem + "_explore")
    outdir.mkdir(parents=True, exist_ok=True)

    # --- 8-bit channels ---
    for idx in range(8):
        vals = [d[idx] if len(d)>idx else None for d in df["data"].tolist()]
        times = [t for t, v in zip(df["time_s"].tolist(), vals) if v is not None]
        series = [v for v in vals if v is not None]
        if not series:
            continue
        plt.figure(figsize=(10,4))
        plt.plot(times, series, marker=".", linestyle="-")
        plt.xlabel("Zeit (s)")
        plt.ylabel(f"Byte[{idx}] (8-bit)")
        plt.title(f"{trace.name} – 8-bit Byte {idx}")
        plt.grid(True)
        fn = outdir / f"{args.prefix}_byte{idx}.png"
        plt.tight_layout()
        plt.savefig(fn, dpi=150)
        plt.close()

    # --- 16-bit combos ---
    pairs = [(i,i+1) for i in range(7)]
    # LE
    for i,j in pairs:
        times, series = [], []
        for t, d in zip(df["time_s"].tolist(), df["data"].tolist()):
            if len(d) > j:
                series.append(le16([d[i], d[j]])); times.append(t)
        if not series:
            continue
        plt.figure(figsize=(10,4))
        plt.plot(times, series, marker=".", linestyle="-")
        plt.xlabel("Zeit (s)")
        plt.ylabel(f"LE16 @{i}-{j}")
        plt.title(f"{trace.name} – LE16 Bytes {i}-{j}")
        plt.grid(True)
        fn = outdir / f"{args.prefix}_le16_{i}-{j}.png"
        plt.tight_layout()
        plt.savefig(fn, dpi=150)
        plt.close()

    # BE
    for i,j in pairs:
        times, series = [], []
        for t, d in zip(df["time_s"].tolist(), df["data"].tolist()):
            if len(d) > j:
                series.append(be16([d[i], d[j]])); times.append(t)
        if not series:
            continue
        plt.figure(figsize=(10,4))
        plt.plot(times, series, marker=".", linestyle="-")
        plt.xlabel("Zeit (s)")
        plt.ylabel(f"BE16 @{i}-{j}")
        plt.title(f"{trace.name} – BE16 Bytes {i}-{j}")
        plt.grid(True)
        fn = outdir / f"{args.prefix}_be16_{i}-{j}.png"
        plt.tight_layout()
        plt.savefig(fn, dpi=150)
        plt.close()

    # Summary stats
    stats = []
    # 8-bit stats
    for idx in range(8):
        vals = [d[idx] if len(d)>idx else None for d in df["data"].tolist()]
        vals = [v for v in vals if v is not None]
        if not vals:
            continue
        arr = np.array(vals, dtype=float)
        stats.append({"type":"byte8", "slot":idx, "min":float(arr.min()), "max":float(arr.max()), "var":float(arr.var())})
    # 16-bit stats
    for i,j in pairs:
        vals = [le16([d[i],d[j]]) for d in df["data"].tolist() if len(d)>j]
        if vals:
            arr = np.array(vals, dtype=float)
            stats.append({"type":"le16", "slot":f"{i}-{j}", "min":float(arr.min()), "max":float(arr.max()), "var":float(arr.var())})
        vals = [be16([d[i],d[j]]) for d in df["data"].tolist() if len(d)>j]
        if vals:
            arr = np.array(vals, dtype=float)
            stats.append({"type":"be16", "slot":f"{i}-{j}", "min":float(arr.min()), "max":float(arr.max()), "var":float(arr.var())})

    pd.DataFrame(stats).to_csv(outdir / "summary_stats.csv", index=False)
    print(f"Exported 8-bit & 16-bit plots and summary_stats.csv to {outdir}")

if __name__ == "__main__":
    main()