"""
Lamp comparison export — one measurement, one publishable data bundle.

Takes a reading from the HPCS 6500 (or an existing pcap capture) and writes
the per-lamp files consumed by the buildfor.life comparison pages:

    <out>/<name>/spd.csv        full spectrum, wavelength_nm,value (380-1050 nm)
    <out>/<name>/tm30.csv       TM-30-18 hue-bin data for the color vector graphic
    <out>/<name>/metrics.json   photometric / colorimetric / CRI / TM-30 /
                                electrical summary

Usage:
    uv run lamp_export.py --name philips-a60-8w
    uv run lamp_export.py --name x --manufacturer Philips --model "A60 8W 927"
    uv run lamp_export.py --name x --readings 5          # average 5 readings
    uv run lamp_export.py --name x --voltage 230 --frequency 50 --settle 60
    uv run lamp_export.py --name x --passive             # vendor SW drives the instrument
    uv run lamp_export.py --name x --parse captures/run.pcap

With --voltage/--frequency/--current/--mode (or --psu), the built-in supply is
configured, switched on for the measurement, and switched off afterwards.
--settle waits after power-on so the lamp stabilizes before the first reading.
"""

import argparse
import csv
import json
import sys
import time
from datetime import datetime
from pathlib import Path

from hpcs6500 import HPCS6500, find_hpcs_port, parse_pcap_messages

SCALAR_GROUPS = {
    "photometric": ["Phi_lm", "eta_lm_W", "CCT_K", "Duv", "SDCM", "TLCI"],
    "chromaticity": ["x", "y", "u", "v", "u_prime", "v_prime", "CIE_X", "CIE_Y", "CIE_Z"],
    "radiometric": [
        "Phi_e_mW", "Phi_euv_mW", "Phi_eb_mW", "Phi_ey_mW",
        "Phi_er_mW", "Phi_efr_mW", "Phi_eir_mW", "Phi_e_total",
    ],
    "electrical": ["Voltage_V", "Current_A", "Power_W", "Freq_Hz", "PF", "UThd", "AThd"],
    "sensor": ["PeakSignal", "DarkSignal", "Compensate"],
}


def average_readings(readings):
    """Element-wise average of scalars and spectra across readings."""
    result = dict(readings[0])
    n = len(readings)
    if n == 1:
        return result
    for key, value in readings[0].items():
        if isinstance(value, (int, float)):
            result[key] = sum(r.get(key, 0.0) for r in readings) / n
        elif key == "spectrum":
            result[key] = [
                sum(r["spectrum"][i] for r in readings) / n
                for i in range(len(value))
            ]
    return result


def readings_from_pcap(path):
    """Extract parsed readings from a pcap capture."""
    messages = parse_pcap_messages(path)
    blocks = [m for m in messages if m["dir"] == "RX" and len(m["data"]) == 3904
              and m["data"][:8] == b"HPCS6500"]
    elec_blocks = [m for m in messages if m["dir"] == "RX" and len(m["data"]) == 1584]
    dev = HPCS6500.__new__(HPCS6500)
    readings = []
    for i, block in enumerate(blocks):
        r = dev._parse_measurement(block["data"])
        if i < len(elec_blocks):
            r.update(dev._parse_electrical(elec_blocks[i]["data"]))
        readings.append(r)
    return readings


def readings_from_device(port, count, passive, psu):
    dev = HPCS6500(port)
    name = dev.identify()
    if name:
        print(f"Device: {name}")
    readings = []
    supply = None
    try:
        if psu["mode"]:
            ok = dev.set_mode(psu["mode"])
            print(f"Mode {psu['mode'].upper()}: {'OK' if ok else 'FAILED'}")
        if psu["voltage"] is not None:
            if psu["mode"] == "dc":
                ok = dev.set_dc_voltage(psu["voltage"])
            else:
                ok = dev.set_ac_voltage(psu["voltage"])
            print(f"Voltage {psu['voltage']:g} V: {'OK' if ok else 'FAILED'}")
        if psu["frequency"] is not None:
            ok = dev.set_ac_frequency(psu["frequency"])
            print(f"Frequency {psu['frequency']:g} Hz: {'OK' if ok else 'FAILED'}")
        if psu["current"] is not None:
            ok = dev.set_dc_current(psu["current"])
            print(f"DC current limit {psu['current']:g} A: {'OK' if ok else 'FAILED'}")

        supply = dev.read_psu_settings()

        if psu["enable"]:
            if not dev.psu_on():
                print("ERROR: failed to turn PSU on")
                sys.exit(1)
            print("PSU on")
            if psu["settle"] > 0:
                print(f"Settling {psu['settle']:g} s ...")
                time.sleep(psu["settle"])
            else:
                time.sleep(0.2)

        for i in range(count):
            print(f"Reading {i + 1}/{count} ...")
            if passive:
                r = dev.read_current()
            else:
                # Vendor single-shot cycle (verified from USB captures):
                # test config -> trigger (8C 0E 02) -> poll -> read -> reset,
                # repeated per reading. auto_psu stays off; we hold the PSU.
                dev.send_test_config(auto_psu=False)
                r = dev.take_single_reading()
            if r is None:
                print("ERROR: failed to get reading")
                sys.exit(1)
            readings.append(r)
            if i < count - 1:
                if psu["enable"] and not passive:
                    # take_single_reading() ends with an instrument reset;
                    # make sure the lamp stays powered for the next reading.
                    dev.psu_on()
                if psu["interval"] > 0:
                    time.sleep(psu["interval"])
    finally:
        if psu["enable"]:
            dev.psu_off()
            print("PSU off")
        dev.close()
    return readings, supply


def write_bundle(reading, out_dir, meta, n_readings, supply=None):
    # Deferred: importing colour-science takes a few seconds, so it happens
    # after the readings rather than at startup.
    from tm30 import compute_tm30

    out_dir.mkdir(parents=True, exist_ok=True)

    spectrum = reading.get("spectrum") or []
    nm = reading.get("spectrum_nm") or []
    if not spectrum or max(spectrum) <= 0:
        print("ERROR: reading contains no spectrum data")
        sys.exit(1)

    # spd.csv — full measured range; consumers trim to visible as needed.
    with open(out_dir / "spd.csv", "w", newline="") as f:
        w = csv.writer(f)
        w.writerow(["wavelength_nm", "value"])
        for wl, val in zip(nm, spectrum):
            w.writerow([f"{wl:.2f}", f"{val:.6g}"])

    # TM-30 from the same spectrum.
    tm30 = compute_tm30(nm, spectrum)
    bin_fields = list(tm30["bins"][0].keys())
    with open(out_dir / "tm30.csv", "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=bin_fields)
        w.writeheader()
        for b in tm30["bins"]:
            w.writerow({k: (f"{v:.6g}" if isinstance(v, float) else v) for k, v in b.items()})

    # metrics.json — grouped scalars plus identification. ean/variant are set
    # by lamp_station.py when one barcode covers several versions of a lamp
    # (e.g. daylight and warm white sharing box art and EAN).
    metrics = {"name": meta["name"]}
    if meta.get("ean"):
        metrics["ean"] = meta["ean"]
    if meta.get("type"):
        # Lamp technology: led, halogen, cfl, incandescent, ...
        metrics["type"] = meta["type"]
    if meta.get("dimmable") is not None:
        metrics["dimmable"] = meta["dimmable"]
    if meta.get("variant"):
        metrics["variant"] = meta["variant"]
    if meta.get("rated"):
        # Advertised values from the packaging, keyed like their measured
        # counterparts so consumers can pair them (truth-in-advertising rows).
        metrics["rated"] = meta["rated"]
    metrics |= {
        "manufacturer": meta["manufacturer"],
        "model": meta["model"],
        "notes": meta["notes"],
        "instrument": reading.get("device", "HPCS6500"),
        "measured_at": datetime.now().astimezone().isoformat(timespec="seconds"),
        "instrument_date": reading.get("test_date"),
        "instrument_time": reading.get("test_time"),
        "readings_averaged": n_readings,
        "tm30": {"Rf": round(tm30["Rf"], 1), "Rg": round(tm30["Rg"], 1)},
        "cri": {"Ra": reading.get("Ra")}
        | {f"R{i}": reading.get(f"R{i}") for i in range(1, 16)},
    }
    if supply:
        metrics["supply"] = supply
    for group, keys in SCALAR_GROUPS.items():
        metrics[group] = {k: reading[k] for k in keys if k in reading}
    with open(out_dir / "metrics.json", "w") as f:
        json.dump(metrics, f, indent=2)

    return tm30


def main():
    parser = argparse.ArgumentParser(description="Export a lamp measurement bundle")
    parser.add_argument("--name", required=True,
                        help="Lamp slug, becomes the output directory name")
    parser.add_argument("--manufacturer", default="")
    parser.add_argument("--model", default="")
    parser.add_argument("--notes", default="")
    parser.add_argument("--out", default="lamps", help="Output base directory")
    parser.add_argument("--port", help="COM port (auto-detect if omitted)")
    parser.add_argument("--readings", type=int, default=1,
                        help="Number of readings to average (default 1)")
    parser.add_argument("--passive", action="store_true",
                        help="Read without controlling the instrument")
    parser.add_argument("--parse", metavar="PCAP",
                        help="Export from a pcap capture instead of the device")

    psu_group = parser.add_argument_group("power supply")
    psu_group.add_argument("--mode", choices=["ac", "dc"],
                           help="Supply mode (default ac when --voltage is given)")
    psu_group.add_argument("--voltage", type=float, help="Supply voltage (V)")
    psu_group.add_argument("--frequency", type=float, help="AC frequency (Hz)")
    psu_group.add_argument("--current", type=float, help="DC current limit (A)")
    psu_group.add_argument("--psu", action="store_true",
                           help="Power the lamp from the built-in supply "
                                "(implied by --mode/--voltage/--frequency/--current)")
    psu_group.add_argument("--settle", type=float, default=0,
                           help="Seconds to wait after PSU on before the first reading")
    parser.add_argument("--interval", type=float, default=1.0,
                        help="Seconds to wait between readings (default 1)")
    args = parser.parse_args()

    psu = {
        "mode": args.mode,
        "voltage": args.voltage,
        "frequency": args.frequency,
        "current": args.current,
        "settle": args.settle,
        "interval": args.interval,
        "enable": args.psu or args.mode is not None or args.voltage is not None
        or args.frequency is not None or args.current is not None,
    }

    supply = None
    if args.parse:
        readings = readings_from_pcap(args.parse)
        if not readings:
            print(f"ERROR: no measurement blocks in {args.parse}")
            sys.exit(1)
        print(f"Using {len(readings)} reading(s) from capture")
    else:
        port = args.port or find_hpcs_port()
        if not port:
            print("ERROR: HPCS 6500 not found. Connect the device or specify --port.")
            sys.exit(1)
        readings, supply = readings_from_device(port, args.readings, args.passive, psu)

    reading = average_readings(readings)
    out_dir = Path(args.out) / args.name
    meta = {k: getattr(args, k) for k in ("name", "manufacturer", "model", "notes")}
    tm30 = write_bundle(reading, out_dir, meta, len(readings), supply)

    print(f"\nWrote {out_dir}/spd.csv, tm30.csv, metrics.json")
    print(f"  {reading.get('Phi_lm', 0):.0f} lm   {reading.get('eta_lm_W', 0):.1f} lm/W   "
          f"{reading.get('CCT_K', 0):.0f} K   Ra {reading.get('Ra', 0):.1f}")
    print(f"  TM-30: Rf {tm30['Rf']:.1f}   Rg {tm30['Rg']:.1f}")


if __name__ == "__main__":
    main()