#!/usr/bin/env python3
"""basic-benchmark: turn powerlog CSVs into graphs + an efficiency table.

    python3 power-report.py --tags before-headless
    python3 power-report.py --tags before-headless before-headed   # comparison

Reads  results/power-<tag>.csv  (+ optional results/phases-<tag>.csv and
results/scores-<tag>.csv).  Writes graphs to results/graphs/ and prints a
per-phase power/thermal summary.

scores-<tag>.csv format (one benchmark per line, name must match a phase):
    PassMark-cpu,24500
    PassMark-mem,2500
    libreoffice,3.9
"""
import argparse
import csv
import os
import statistics as st

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

HERE = os.path.dirname(os.path.abspath(__file__))
RES = os.path.join(HERE, "results")
GRAPHS = os.path.join(RES, "graphs")


def fnum(v):
    try:
        return float(v)
    except (TypeError, ValueError):
        return None


def load_power(tag):
    path = os.path.join(RES, f"power-{tag}.csv")
    if not os.path.exists(path):
        raise SystemExit(f"missing {path}")
    with open(path) as fh:
        rows = list(csv.DictReader(fh))
    return [r for r in rows]


def load_series(rows, key):
    return [(int(r["epoch"]), fnum(r.get(key))) for r in rows]


def load_phases(tag):
    path = os.path.join(RES, f"phases-{tag}.csv")
    out = []
    if os.path.exists(path):
        with open(path) as fh:
            for line in fh:
                a, _, b = line.strip().partition(",")
                try:
                    out.append((int(float(a)), b))
                except ValueError:
                    pass
    return sorted(out)


def load_scores(tag):
    path = os.path.join(RES, f"scores-{tag}.csv")
    out = {}
    if os.path.exists(path):
        with open(path) as fh:
            for line in fh:
                parts = [p.strip() for p in line.strip().split(",")]
                if len(parts) >= 2 and fnum(parts[1]) is not None:
                    out[parts[0]] = fnum(parts[1])
    return out


def load_cpu(tag):
    path = os.path.join(RES, f"cpu-{tag}.txt")
    if os.path.exists(path):
        with open(path) as fh:
            name = fh.read().strip()
            if name:
                return name
    return None


def desc(vals):
    vals = [v for v in vals if v is not None]
    if not vals:
        return None
    srt = sorted(vals)
    return {
        "n": len(vals),
        "mean": st.mean(vals),
        "p95": srt[min(len(srt) - 1, int(len(srt) * 0.95))],
        "max": srt[-1],
    }


def window(series, lo, hi):
    return [v for t, v in series if v is not None and lo <= t < hi]


def main():
    global RES, GRAPHS
    ap = argparse.ArgumentParser()
    ap.add_argument("--tags", nargs="+", required=True)
    ap.add_argument("--results", default=RES, help="results directory (default ./results)")
    ap.add_argument("--label", action="append", default=[], metavar="TAG=NAME",
                    help="friendly name for a tag in the comparison graph (repeatable)")
    args = ap.parse_args()

    RES = os.path.abspath(args.results)
    GRAPHS = os.path.join(RES, "graphs")
    labels = {}
    for spec in args.label:
        if "=" in spec:
            k, v = spec.split("=", 1)
            labels[k] = v

    def label(t):
        return labels.get(t, t)

    os.makedirs(GRAPHS, exist_ok=True)

    data = {}
    for tag in args.tags:
        rows = load_power(tag)
        data[tag] = {
            "rows": rows,
            "pkg": load_series(rows, "pkg_w"),
            "core": load_series(rows, "core_w"),
            "gpu": load_series(rows, "gpu_w"),
            "temp": load_series(rows, "cpu_temp_c"),
            "phases": load_phases(tag),
            "scores": load_scores(tag),
            "cpu": load_cpu(tag),
        }

    for tag in args.tags:
        d = data[tag]
        pkg = d["pkg"]
        if not pkg:
            continue
        t0 = pkg[0][0]
        cpu = d["cpu"] or "unknown CPU"
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(11, 7), sharex=True)
        for series_label, series, color in (("package W", d["pkg"], "#c0392b"),
                                            ("core W", d["core"], "#e67e22"),
                                            ("GPU W", d["gpu"], "#2980b9")):
            pts = [((t - t0) / 60, v) for t, v in series if v is not None]
            if pts:
                ax1.plot([p[0] for p in pts], [p[1] for p in pts], label=series_label, color=color, lw=1)
        ax1.set_ylabel("power (W)")
        ax1.legend(loc="upper right", fontsize=8)
        ax1.set_title(f"power & thermal profile — {tag}\nCPU: {cpu}")
        ax1.grid(alpha=0.3)

        temps = [((t - t0) / 60, v) for t, v in d["temp"] if v is not None]
        if temps:
            ax2.plot([p[0] for p in temps], [p[1] for p in temps], color="#16a085", lw=1)
        ax2.set_ylabel("CPU temp (°C)")
        ax2.set_xlabel("minutes from start")
        ax2.grid(alpha=0.3)

        for t, name in d["phases"]:
            for ax in (ax1, ax2):
                ax.axvline((t - t0) / 60, color="0.6", ls="--", lw=0.8)
            ax1.text((t - t0) / 60, ax1.get_ylim()[1], name, rotation=90,
                     va="top", ha="right", fontsize=7, color="0.35")
        out = os.path.join(GRAPHS, f"power-{tag}.png")
        fig.tight_layout()
        fig.savefig(out, dpi=110)
        plt.close(fig)
        print(f"graph: {out}")

        print(f"\n== {tag} — per phase  (CPU: {cpu}) ==")
        print(f"{'phase':<22}{'dur s':>7}{'avg W':>7}{'p95 W':>7}{'peak W':>7}"
              f"{'avg °C':>7}{'peak °C':>8}{'energy kJ':>11}")
        ph = d["phases"]
        spans = []
        if ph:
            for i, (t, name) in enumerate(ph):
                hi = ph[i + 1][0] if i + 1 < len(ph) else pkg[-1][0] + 1
                spans.append((name, t, hi))
        else:
            spans.append(("whole run", pkg[0][0], pkg[-1][0] + 1))
        for name, lo, hi in spans:
            pw = window(d["pkg"], lo, hi)
            tp = window(d["temp"], lo, hi)
            s = desc(pw)
            ts = desc(tp)
            if not s:
                continue
            dur = hi - lo
            print(f"{name:<22}{dur:>7}{s['mean']:>7.1f}{s['p95']:>7.1f}{s['max']:>7.1f}"
                  f"{(ts['mean'] if ts else 0):>7.1f}{(ts['max'] if ts else 0):>8.1f}"
                  f"{s['mean'] * dur / 1000:>11.1f}")

        if d["scores"]:
            def find_span(key):
                low = key.lower()
                for s in spans:
                    if s[0] == key:
                        return s
                for s in spans:
                    if low in s[0].lower() or s[0].lower() in low:
                        return s
                return None

            print(f"\n== {tag} — efficiency (score / avg package W) ==")
            print(f"{'benchmark':<22}{'score':>12}{'avg W':>8}{'score/W':>10}")
            for name, score in d["scores"].items():
                span = find_span(name)
                if not span:
                    print(f"{name:<22}{score:>12.1f}{'no phase match':>18}")
                    continue
                s = desc(window(d["pkg"], span[1], span[2]))
                if not s:
                    continue
                print(f"{name:<22}{score:>12.1f}{s['mean']:>8.1f}{score / s['mean']:>10.2f}")

    if len(args.tags) > 1:
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(11, 7), sharex=True)
        palette = ("#c0392b", "#2980b9", "#27ae60", "#8e44ad", "#d35400", "#16a085")
        # Keep legend/labels short: shorten the on-plot names and drop the
        # per-machine CPU names (they go in the printed summary, not the chart).
        def short(t):
            s = labels.get(t) or t
            s = s.replace("-headless", "").replace("_", " ")
            return s if len(s) <= 16 else s[:15] + "…"
        for i, tag in enumerate(args.tags):
            color = palette[i % len(palette)]
            d = data[tag]
            if not d["pkg"]:
                continue
            t0 = d["pkg"][0][0]
            pts = [((t - t0) / 60, v) for t, v in d["pkg"] if v is not None]
            ax1.plot([p[0] for p in pts], [p[1] for p in pts], label=short(tag), color=color, lw=1)
            tt = [((t - t0) / 60, v) for t, v in d["temp"] if v is not None]
            if tt:
                ax2.plot([p[0] for p in tt], [p[1] for p in tt], label=short(tag), color=color, lw=1)
        ax1.set_ylabel("package W")
        ax1.legend(fontsize=8, title="machine", title_fontsize=8)
        ax1.set_title(f"package power & CPU temp — {len(args.tags)} machines")
        ax1.grid(alpha=0.3)
        ax2.set_ylabel("CPU temp (°C)")
        ax2.set_xlabel("minutes from start")
        ax2.legend(fontsize=8)
        ax2.grid(alpha=0.3)
        out = os.path.join(GRAPHS, "power-thermal-compare.png")
        fig.tight_layout()
        fig.savefig(out, dpi=110)
        plt.close(fig)
        print(f"\ngraph: {out}")


if __name__ == "__main__":
    main()