"""Convert the neural models' scaled-unit metrics to kWh and assemble the
final five-model comparison table."""
import json
from pathlib import Path

import pandas as pd

from huggingface_hub import hf_hub_download

import sys
sys.path.insert(0, str(Path.cwd()))
from train_seq import daily_origins, download_all, prepare, REPO_CODE  # noqa

files = download_all(Path.cwd())
D = prepare(files)
sd = D["sd"]  # per-consumer kWh std from the fit window

index = pd.date_range("2024-12-31 23:00", periods=len(D["Y"]), freq="15min", tz="UTC")
origins_test = daily_origins(index, D["test"])
M = D["M"]
cnt = [int(sum(M[o + 96:o + 192, c].sum() for o in origins_test)) for c in range(27)]
total = sum(cnt)
print(f"observed test pairs per consumer: {cnt[0]:d} ... (total {total})")

out = {}
for model in ["gnn", "lstm"]:
    p = hf_hub_download(REPO_CODE, f"results/{model}.json", repo_type="model",
                        local_dir=Path.cwd())
    r = json.load(open(p))
    per_c = r["test"]["pinball_per_consumer"]
    kwh_c = {k: per_c[k] * float(sd[i]) for i, k in enumerate(per_c)}
    agg = sum(cnt[i] * v for i, k in enumerate(sorted(kwh_c)) for kk, v in
              [(k, kwh_c[k])]) / total
    cov = {k: round(v, 3) for k, v in r["test"].items() if k.startswith("coverage")}
    out[model] = {
        "pinball_kwh": round(agg, 3),
        "pinball_kwh_per_consumer": {k: round(v, 3) for k, v in kwh_c.items()},
        "coverage": cov,
        "note": "converted from z-scaled units via per-consumer fit-window sd",
    }
    print(f"\n{model}: pinball {agg:.3f} kWh | coverage@0.90 "
          f"{cov['coverage_0.90']:.3f}")
    print("  worst 5 consumers (kWh):",
          dict(sorted(kwh_c.items(), key=lambda kv: -kv[1])[:5]))

for f in ["baselines.json", "xgboost_residual.json"]:
    p = hf_hub_download(REPO_CODE, f"results/{f}", repo_type="model", local_dir=Path.cwd())
    r = json.load(open(p))
    for k, v in r.items():
        if isinstance(v, dict) and "pinball" in v:
            out.setdefault(k, {})["pinball_kwh"] = round(v["pinball"], 3)
            out[k]["coverage"] = {kk: round(vv, 3) for kk, vv in v.items()
                                  if kk.startswith("coverage")}
            print(f"{k}: pinball {v['pinball']:.3f} kWh | coverage@0.90 "
                  f"{out[k]['coverage'].get('coverage_0.90', float('nan')):.3f}")

Path("results").mkdir(exist_ok=True)
Path("results/comparison.json").write_text(json.dumps(out, indent=2))
print("\nwrote results/comparison.json")