"""Plot efficiency vs input voltage vs current from bench CSV logs.

Auto-detects the three CSV formats produced by the tooling:
  - GUI data log      (data_*.csv: instrument columns + merged stm_* snapshot)
  - GUI telemetry log (*_telem.csv: full-rate 100 Hz board broadcast)
  - CLI sweep         (sweep_vi_*.csv: voltage_set/load_setpoint grid)

Left panel: operating-point scatter (x = Vin, y = current, color = efficiency).
Right panel: efficiency vs current, one curve per Vin bin.

Multiple files are pooled into one dataset. When several files contain the
same Vin bin, a checkbox picker opens (one row per bin, one column per file)
so you can choose which file's data to use per bin; force it with --pick or
suppress it with --no-pick.

Usage:
    bench-plot data_20260703_120000.csv
    bench-plot run_telem.csv another_telem.csv --current iin --save eff.png
    bench-plot sweep_a.csv sweep_b.csv --pick
"""

from __future__ import annotations

import argparse
import csv
import math
import os
import sys
from collections import Counter

import matplotlib.pyplot as plt
import numpy as np

MIN_P_IN_W = 0.1     # same gate as the firmware/GUI efficiency calc


def _f(row: dict, key: str) -> float:
    """Float cell value; blank/missing/garbage -> NaN."""
    v = row.get(key, "")
    if v is None or v == "":
        return math.nan
    try:
        return float(v)
    except ValueError:
        return math.nan


def _detect_format(header: list[str]) -> str:
    cols = set(header)
    if "stm_eff_net_pct" in cols:
        return "datalog"
    if "iout_slow" in cols and "p_in_W" in cols:
        return "telem"
    if "voltage_set" in cols and "efficiency" in cols:
        return "sweep"
    if not header:
        raise ValueError("empty file (a log that is still being written?)")
    raise ValueError(f"unrecognized CSV header: {header[:6]}...")


def _fatal(msg: str) -> None:
    """Print the error AND show it in a messagebox: when spawned from the
    GUI button nobody sees stderr, so exiting silently looks like a no-op."""
    print(msg, file=sys.stderr)
    if not os.environ.get("BENCH_PLOT_HEADLESS"):
        try:
            import tkinter as tk
            from tkinter import messagebox
            r = tk.Tk()
            r.withdraw()
            r.attributes("-topmost", True)
            messagebox.showerror("bench-plot", msg, parent=r)
            r.destroy()
        except Exception:
            pass
    sys.exit(1)


def _extract(path: str, source: str) -> tuple[dict, str]:
    """Read one CSV -> dict of float lists (vin_V, iin_A, iout_A, eff_pct,
    p_out_W) plus the efficiency-source label actually used."""
    with open(path, newline="") as fh:
        reader = csv.DictReader(fh)
        header = reader.fieldnames or []
        rows = list(reader)
    fmt = _detect_format(header)

    out: dict[str, list] = {k: [] for k in ("vin_V", "iin_A", "iout_A", "eff_pct", "p_out_W")}

    if fmt == "sweep":
        for r in rows:
            out["vin_V"].append(_f(r, "supply_V"))
            out["iin_A"].append(_f(r, "supply_I"))
            out["iout_A"].append(_f(r, "load_I"))
            out["eff_pct"].append(_f(r, "efficiency"))
            out["p_out_W"].append(_f(r, "output_power"))
        return out, "sweep efficiency (HIOKI)"

    if fmt == "telem":
        # Wire units: vin/vout in mV, currents in mA (iin negative into the
        # converter); eff net = (P_out - P_sys) / P_in, same as the GUI panel.
        for r in rows:
            p_in = _f(r, "p_in_W")
            p_out = _f(r, "p_out_W")
            p_sys = _f(r, "vout") * _f(r, "sys_current_ma") / 1e6
            eff = (p_out - p_sys) / p_in * 100.0 if p_in > MIN_P_IN_W else math.nan
            out["vin_V"].append(_f(r, "vin") / 1000.0)
            out["iin_A"].append(-_f(r, "iin") / 1000.0)
            out["iout_A"].append(_f(r, "iout_slow") / 1000.0)
            out["eff_pct"].append(eff)
            out["p_out_W"].append(p_out)
        return out, "board eff net (iout_slow, -P_sys)"

    # datalog: three consistent (vin, current, eff) triples to choose from
    if source == "auto":
        med_eff1 = np.nanmedian([_f(r, "meter_EFF1") for r in rows]) if rows else math.nan
        med_psup = np.nanmedian([_f(r, "supply_P") for r in rows]) if rows else math.nan
        if med_eff1 > 1.0:
            source = "hioki"
        elif med_psup > MIN_P_IN_W:
            source = "instr"
        else:
            source = "stm"

    for r in rows:
        if source == "hioki":
            out["vin_V"].append(_f(r, "meter_U5"))
            out["iin_A"].append(_f(r, "meter_I5"))
            out["iout_A"].append(_f(r, "meter_I6"))
            out["eff_pct"].append(_f(r, "meter_EFF1"))
            out["p_out_W"].append(_f(r, "meter_P6"))
        elif source == "instr":
            p_sup = _f(r, "supply_P")
            eff = _f(r, "load_P") / p_sup * 100.0 if p_sup > MIN_P_IN_W else math.nan
            out["vin_V"].append(_f(r, "supply_V"))
            out["iin_A"].append(_f(r, "supply_I"))
            out["iout_A"].append(_f(r, "load_I"))
            out["eff_pct"].append(eff)
            out["p_out_W"].append(_f(r, "load_P"))
        else:  # stm
            out["vin_V"].append(_f(r, "stm_vin_mV") / 1000.0)
            out["iin_A"].append(-_f(r, "stm_iin_mA") / 1000.0)
            out["iout_A"].append(_f(r, "stm_iout_slow_mA") / 1000.0)
            out["eff_pct"].append(_f(r, "stm_eff_net_pct"))
            out["p_out_W"].append(_f(r, "stm_p_out_W"))
    labels = {"hioki": "HIOKI EFF1", "instr": "supply/load power",
              "stm": "board eff net"}
    return out, labels[source]


class _BinPicker:
    """Checkbox grid: one row per Vin bin, one column per file.

    result is a list of selected-center sets (aligned with bin_table)
    after OK, or None if the window was cancelled/closed.
    """

    ROW_H = 26

    def __init__(self, bin_table: list[tuple[str, dict]], w: float):
        import tkinter as tk
        from tkinter import ttk
        self._tk = tk
        self.root = tk.Tk()
        self.root.title(f"Select Vin bins per file ({w:g} V bins)")
        # Spawned from the GUI's Plot Eff button, Windows denies foreground
        # to the new process (the user is mid-click in the GUI) and the
        # picker opens hidden BEHIND it -- keep this short-lived dialog
        # on top and grab focus.
        self.root.lift()
        self.root.attributes("-topmost", True)
        self.root.focus_force()
        self._n_files = len(bin_table)
        self._vars: dict[tuple[int, float], object] = {}
        self.result: list[set] | None = None

        centers = sorted({c for _, bins in bin_table for c in bins})
        shared = {c for c in centers
                  if sum(c in bins for _, bins in bin_table) > 1}

        top = ttk.Frame(self.root, padding=6)
        top.pack(fill="both", expand=True)
        canvas = tk.Canvas(
            top, highlightthickness=0,
            height=min(self.ROW_H * (len(centers) + 2) + 8, 560))
        vsb = ttk.Scrollbar(top, orient="vertical", command=canvas.yview)
        grid = ttk.Frame(canvas)
        grid.bind("<Configure>",
                  lambda e: canvas.configure(scrollregion=canvas.bbox("all"),
                                             width=grid.winfo_reqwidth()))
        canvas.create_window((0, 0), window=grid, anchor="nw")
        canvas.configure(yscrollcommand=vsb.set)
        canvas.pack(side="left", fill="both", expand=True)
        vsb.pack(side="right", fill="y")
        canvas.bind_all("<MouseWheel>", lambda e: canvas.yview_scroll(
            -1 if e.delta > 0 else 1, "units"))

        ttk.Label(grid, text="Vin bin").grid(row=0, column=0, padx=6, sticky="w")
        for fi, (name, _bins) in enumerate(bin_table):
            short = name if len(name) <= 26 else "..." + name[-23:]
            ttk.Label(grid, text=short).grid(row=0, column=1 + fi, padx=8)
            ttk.Button(grid, text="all/none", width=8,
                       command=lambda fi=fi: self._toggle_col(fi)
                       ).grid(row=1, column=1 + fi, padx=8)
        for ri, c in enumerate(centers):
            lbl = ttk.Label(grid, text=f"{c:g} V")
            if c in shared:
                lbl.configure(foreground="#cc6600")
            lbl.grid(row=2 + ri, column=0, padx=6, sticky="w")
            for fi, (_name, bins) in enumerate(bin_table):
                if c not in bins:
                    continue
                var = tk.BooleanVar(self.root, value=True)
                self._vars[(fi, c)] = var
                ttk.Checkbutton(grid, text=str(bins[c]), variable=var
                                ).grid(row=2 + ri, column=1 + fi,
                                       padx=8, sticky="w")

        bar = ttk.Frame(self.root, padding=(6, 0, 6, 6))
        bar.pack(fill="x")
        ttk.Label(bar, text="numbers = points per bin, "
                            "orange = bin present in several files"
                  ).pack(side="left")
        ttk.Button(bar, text="Cancel",
                   command=self.root.destroy).pack(side="right", padx=4)
        ttk.Button(bar, text="Plot", command=self._ok).pack(side="right")

    def _toggle_col(self, fi: int) -> None:
        cells = [v for (f, _c), v in self._vars.items() if f == fi]
        state = not all(v.get() for v in cells)
        for v in cells:
            v.set(state)

    def _ok(self) -> None:
        self.result = [
            {c for (f, c), v in self._vars.items() if f == fi and v.get()}
            for fi in range(self._n_files)]
        self.root.destroy()

    def run(self) -> list[set] | None:
        self.root.mainloop()
        return self.result


def _pick_bins(bin_table: list[tuple[str, dict]], w: float) -> list[set] | None:
    """Show the picker; on a headless/Tk failure fall back to everything."""
    try:
        picker = _BinPicker(bin_table, w)
    except Exception as e:
        print(f"bin picker unavailable ({e}) - using all bins")
        return [set(bins) for _, bins in bin_table]
    return picker.run()


def main() -> None:
    ap = argparse.ArgumentParser(
        description="Plot efficiency vs input voltage vs current from bench CSVs.")
    ap.add_argument("csv", nargs="*",
                    help="logged CSV file(s); a file dialog opens if omitted")
    ap.add_argument("--current", choices=("iout", "iin"), default="iout",
                    help="current axis: output (default) or input current")
    ap.add_argument("--source", choices=("auto", "hioki", "instr", "stm"),
                    default="auto",
                    help="efficiency source for GUI data logs (default auto: "
                         "HIOKI if present, else supply/load, else board)")
    ap.add_argument("--vin-bin", type=float, default=1.0, metavar="V",
                    help="Vin bin width for the per-voltage curves (default 1.0)")
    ap.add_argument("--min-pout", type=float, default=5.0, metavar="W",
                    help="drop points below this output power (default 5.0)")
    ap.add_argument("--save", metavar="PNG", help="write the figure instead of showing it")
    ap.add_argument("--title", default=None, help="figure title override")
    g = ap.add_mutually_exclusive_group()
    g.add_argument("--pick", action="store_true",
                   help="always show the per-file Vin-bin checkbox picker")
    g.add_argument("--no-pick", action="store_true",
                   help="never show the picker (default: it opens when "
                        "several files contain the same Vin bin)")
    args = ap.parse_args()

    if not args.csv:
        import tkinter as tk
        from tkinter import filedialog
        root = tk.Tk()
        root.withdraw()
        root.attributes("-topmost", True)
        args.csv = list(filedialog.askopenfilenames(
            parent=root,
            title="Select logged CSV(s) to plot",
            filetypes=[("CSV files", "*.csv"), ("All files", "*.*")]))
        root.destroy()
        if not args.csv:
            sys.exit("no file selected")

    # Load each file separately (kept apart for the per-file bin picker)
    w = args.vin_bin
    parts, labels, n_total, bad = [], set(), 0, []
    for path in args.csv:
        try:
            part, label = _extract(path, args.source)
        except Exception as e:
            bad.append(f"{os.path.basename(path)}: {e}")
            print(f"  skipping {os.path.basename(path)}: {e}")
            continue
        labels.add(label)
        fvin = np.asarray(part["vin_V"])
        fcur = np.asarray(part["iin_A" if args.current == "iin" else "iout_A"])
        feff = np.asarray(part["eff_pct"])
        fpout = np.asarray(part["p_out_W"])
        n_total += len(fvin)
        keep = (np.isfinite(fvin) & np.isfinite(fcur) & np.isfinite(feff)
                & (feff > 0.0) & (feff <= 105.0) & (fpout >= args.min_pout))
        p = {"name": os.path.basename(path), "vin": fvin[keep],
             "cur": fcur[keep], "eff": feff[keep]}
        p["bins"] = np.round(p["vin"] / w) * w
        parts.append(p)
        if len(args.csv) > 1:
            rng = (f"Vin bins {p['bins'].min():g}..{p['bins'].max():g} V"
                   if len(p["vin"]) else "no usable points")
            print(f"  {p['name']}: {len(p['vin'])} pts, {rng}")

    if not parts:
        _fatal("none of the selected files were readable:\n" + "\n".join(bad))

    # Per-file Vin-bin selection: opens automatically when files overlap
    bin_table = [(p["name"], Counter(p["bins"].tolist())) for p in parts]
    files_per_bin = Counter(c for _, bins in bin_table for c in bins)
    overlap = any(n > 1 for n in files_per_bin.values())
    if not args.no_pick and (args.pick or (overlap and len(parts) > 1)):
        sels = _pick_bins(bin_table, w)
        if sels is None:
            sys.exit("cancelled")
        for p, sel in zip(parts, sels):
            m = (np.isin(p["bins"], sorted(sel)) if sel
                 else np.zeros(len(p["bins"]), dtype=bool))
            for k in ("vin", "cur", "eff", "bins"):
                p[k] = p[k][m]

    vin = np.concatenate([p["vin"] for p in parts])
    cur = np.concatenate([p["cur"] for p in parts])
    eff = np.concatenate([p["eff"] for p in parts])
    if len(vin) == 0:
        _fatal(f"no usable points ({n_total} rows read; all filtered or "
               f"deselected — check --min-pout / --source / bin picker)")

    ipk = int(np.argmax(eff))
    print(f"{len(vin)} points ({n_total - len(vin)} filtered) | "
          f"Vin {vin.min():.1f}..{vin.max():.1f} V | "
          f"I {cur.min():.2f}..{cur.max():.2f} A | "
          f"peak eff {eff[ipk]:.2f} % @ {vin[ipk]:.1f} V, {cur[ipk]:.2f} A")

    cur_name = "Input current (A)" if args.current == "iin" else "Output current (A)"
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13.5, 5.8))
    fig.suptitle(args.title or f"Efficiency map — {', '.join(sorted(labels))}")

    # Left: operating-point scatter, color = efficiency
    vmin = np.percentile(eff, 5)
    sc = ax1.scatter(vin, cur, c=eff, s=14, cmap="viridis",
                     vmin=vmin, vmax=eff.max(), rasterized=True)
    fig.colorbar(sc, ax=ax1, label="Efficiency (%)")
    ax1.plot(vin[ipk], cur[ipk], "r*", ms=14, mec="k",
             label=f"peak {eff[ipk]:.2f} %")
    ax1.set_xlabel("Input voltage (V)")
    ax1.set_ylabel(cur_name)
    ax1.legend(loc="best", fontsize=8)
    ax1.grid(alpha=0.3)

    # Right: efficiency vs current, one mean curve per Vin bin
    centers = np.unique(np.round(vin / w) * w)
    cmap = plt.cm.plasma(np.linspace(0.0, 0.9, len(centers)))
    for color, c0 in zip(cmap, centers):
        m = np.abs(vin - c0) <= w / 2
        if m.sum() < 2:
            ax2.plot(cur[m], eff[m], "o", color=color, ms=4,
                     label=f"{c0:g} V")
            continue
        edges = np.linspace(cur[m].min(), cur[m].max() + 1e-9, 41)
        idx = np.digitize(cur[m], edges)
        xs, ys = [], []
        for b in np.unique(idx):
            bm = idx == b
            xs.append(cur[m][bm].mean())
            ys.append(eff[m][bm].mean())
        ax2.plot(xs, ys, "-o", color=color, ms=3, lw=1.2, label=f"{c0:g} V")
    ax2.set_xlabel(cur_name)
    ax2.set_ylabel("Efficiency (%)")
    ax2.grid(alpha=0.3)
    if len(centers) <= 14:
        ax2.legend(title="Vin bin", fontsize=8, ncols=1 + len(centers) // 8)
    else:
        norm = plt.Normalize(centers.min(), centers.max())
        fig.colorbar(plt.cm.ScalarMappable(norm=norm, cmap="plasma"),
                     ax=ax2, label="Vin bin (V)")

    fig.tight_layout()
    if args.save:
        fig.savefig(args.save, dpi=140)
        print(f"saved: {args.save}")
    else:
        plt.show()


if __name__ == "__main__":
    main()