"""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.

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

from __future__ import annotations

import argparse
import csv
import math
import sys

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"
    raise ValueError(f"unrecognized CSV header: {header[:6]}...")


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]


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")
    args = ap.parse_args()

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

    data: dict[str, list] = {k: [] for k in ("vin_V", "iin_A", "iout_A", "eff_pct", "p_out_W")}
    labels = set()
    for path in args.csv:
        part, label = _extract(path, args.source)
        for k in data:
            data[k].extend(part[k])
        labels.add(label)

    vin = np.asarray(data["vin_V"])
    cur = np.asarray(data["iin_A" if args.current == "iin" else "iout_A"])
    eff = np.asarray(data["eff_pct"])
    pout = np.asarray(data["p_out_W"])

    keep = (np.isfinite(vin) & np.isfinite(cur) & np.isfinite(eff)
            & (eff > 0.0) & (eff <= 105.0) & (pout >= args.min_pout))
    n_total = len(vin)
    vin, cur, eff = vin[keep], cur[keep], eff[keep]
    if len(vin) == 0:
        sys.exit(f"no usable points ({n_total} rows read; all filtered — "
                 f"check --min-pout / --source)")

    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
    w = args.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()
