# SPDX-License-Identifier: MIT
"""YS-EN-0006 · reduce the VTT SCOTT IAQ dataset (Zenodo 3774723, CC BY 4.0) to hourly rooms.

  python prep_vtt.py VTT_SCOTT_IAQ_dataset.zip results/vtt_hourly.csv.gz

Per room: per-minute median CO2 across the room's CO2 sensors and mean PIR activity
across its PIR nodes; then hourly means on the local (Europe/Helsinki) wall clock,
kept only where at least 30 of the hour's 60 minutes were reported.  The spring-forward
hour is missing; the doubled autumn hour is averaged.  No other cleaning or imputation.
"""
from __future__ import annotations

import io
import sys
import zipfile
from collections import defaultdict

import numpy as np
import pandas as pd

MIN_MINUTES = 30
IDX = pd.date_range("2019-01-01", "2020-01-01", freq="h", inclusive="left")


def main(zpath, out):
    z = zipfile.ZipFile(zpath)
    files = defaultdict(list)
    for n in z.namelist():
        if n.lower().endswith(".csv"):
            parts = n.split("/")
            room = next((p for p in parts if p.lower().startswith("room")), None)
            if room:
                files[room.lower()].append(n)
    frames = []
    for room in sorted(files):
        co2, pir = [], []
        for n in files[room]:
            raw = z.read(n)
            if b"," not in raw[:200]:
                continue
            d = pd.read_csv(io.BytesIO(raw))
            t = pd.to_datetime(d["devicetimestamp"], unit="s", utc=True).dt.floor("min")
            if "co2" in d:
                co2.append(pd.Series(d["co2"].to_numpy(float), index=t).groupby(level=0).mean())
            if "pir_cnt" in d:
                pir.append(pd.Series(d["pir_cnt"].to_numpy(float), index=t).groupby(level=0).mean())
        c = pd.concat(co2, axis=1).median(axis=1)
        c = c[(c > 300) & (c < 5000)]      # sensor range sanity only
        p = pd.concat(pir, axis=1).mean(axis=1) if pir else pd.Series(dtype=float)
        row = {"room": room, "co2_sensors": len(co2), "pir_nodes": len(pir)}
        h = {}
        for key, s in (("co2", c), ("pir", p)):
            local = s.tz_convert("Europe/Helsinki").tz_localize(None).index.floor("h")
            g = pd.Series(s.to_numpy(), index=local).groupby(level=0).agg(["mean", "count"])
            # the doubled autumn hour has up to 120 minutes; count threshold still applies
            v = g["mean"].where(g["count"] >= MIN_MINUTES)
            h[key] = v.reindex(IDX)
        f = pd.DataFrame({"time": IDX, "room": room, "co2": h["co2"].round(1),
                          "pir": h["pir"].round(3)})
        frames.append(f)
        print(f"{room}: {row['co2_sensors']} CO2 sensor(s), {row['pir_nodes']} PIR node(s), "
              f"hourly CO2 coverage {f['co2'].notna().mean():.1%}", flush=True)
    pd.concat(frames).to_csv(out, index=False, compression="gzip")


if __name__ == "__main__":
    main(*sys.argv[1:3])