"""
One-time data prep: AMPds2 -> compact hourly parquet for the Space.

Usage
-----
    python prepare_data.py /path/to/dataverse_files.zip
    python prepare_data.py /path/to/AMPds2_folder
    python prepare_data.py            # uses the default Colab/Drive path below

Produces data/ampds2_hourly.parquet (small enough to commit to the HF Space).
Resamples the minutely AMPds2 active-power readings to hourly means.
"""
import io
import os
import sys
import zipfile
import tempfile
import numpy as np
import pandas as pd

DEFAULT_INPUT = "/content/drive/MyDrive/AMPds2/dataverse_files.zip"
OUT = "data/ampds2_hourly.parquet"

METER_RE = __import__("re").compile(r"^[A-Z][A-Z0-9]{1,2}E$")   # WHE, FGE, HPE, B1E, B2E, ...


def _read_table(buf_or_path, name):
    sep = "\t" if name.lower().endswith((".tab", ".tsv")) else ","
    return pd.read_csv(buf_or_path, sep=sep, low_memory=False)


def _score(df):
    """How likely this table is the wide active-power file (meter-code columns incl. WHE)."""
    cols = [str(c).strip().upper() for c in df.columns]
    meters = [c for c in cols if METER_RE.match(c)]
    return (len(meters), "WHE" in cols)


def _iter_tables(path):
    """Yield (name, dataframe) for every csv/tab/tsv found in a zip (incl. nested) or folder."""
    if zipfile.is_zipfile(path):
        with zipfile.ZipFile(path) as z:
            for m in z.namelist():
                low = m.lower()
                if low.endswith(".zip"):
                    with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tf:
                        tf.write(z.read(m)); nested = tf.name
                    try:
                        yield from _iter_tables(nested)
                    finally:
                        os.unlink(nested)
                elif low.endswith((".csv", ".tab", ".tsv")):
                    try:
                        yield m, _read_table(io.BytesIO(z.read(m)), m)
                    except Exception as e:
                        print("  skip", m, "->", e)
    elif os.path.isdir(path):
        for root, _, files in os.walk(path):
            for f in files:
                if f.lower().endswith((".csv", ".tab", ".tsv")):
                    fp = os.path.join(root, f)
                    try:
                        yield f, _read_table(fp, f)
                    except Exception as e:
                        print("  skip", f, "->", e)
    else:  # single file
        name = os.path.basename(path)
        yield name, _read_table(path, name)


def find_power_table(path):
    best, best_score, best_name = None, (-1, False), None
    for name, df in _iter_tables(path):
        # an active-power file is preferred (…_P… in AMPds2); compute a score regardless
        s = _score(df)
        bonus = (1, s[1]) if ("_p" in name.lower() or name.lower().startswith("electricity")) else (0, s[1])
        score = (s[0] + bonus[0] * 100, s[1])
        print(f"  candidate {name:40s} meters={s[0]:2d} has_WHE={s[1]}")
        if score > best_score:
            best, best_score, best_name = df, score, name
    if best is None or best_score[0] < 1:
        raise SystemExit("No AMPds2 power table found (need a CSV with meter-code columns like WHE, FGE).")
    print("  -> using:", best_name)
    return best


def to_hourly(df):
    df = df.copy()
    df.columns = [str(c).strip() for c in df.columns]
    upper = {c: c.upper() for c in df.columns}

    # timestamp: prefer an explicit UNIX seconds column, else first datetime-like column
    ts_col = next((c for c in df.columns if upper[c] in ("UNIX_TS", "TS", "TIMESTAMP", "TIME")), None)
    if ts_col is not None and np.issubdtype(df[ts_col].dropna().dtype, np.number):
        idx = pd.to_datetime(df[ts_col], unit="s")
    elif ts_col is not None:
        idx = pd.to_datetime(df[ts_col], errors="coerce")
    else:
        ts_col = df.columns[0]
        idx = pd.to_datetime(df[ts_col], errors="coerce")
        if idx.isna().mean() > 0.5:                     # maybe it's unix seconds in col 0
            idx = pd.to_datetime(pd.to_numeric(df[ts_col], errors="coerce"), unit="s")
    df.index = idx

    meters = [c for c in df.columns if METER_RE.match(upper[c]) and c != ts_col]
    out = df[meters].apply(pd.to_numeric, errors="coerce")
    out.columns = [upper[c] for c in meters]
    out = out[~out.index.isna()].sort_index()
    hourly = out.resample("h").mean()
    return hourly


def main():
    inp = sys.argv[1] if len(sys.argv) > 1 else DEFAULT_INPUT
    if not os.path.exists(inp):
        raise SystemExit(f"Input not found: {inp}\nPass the path to dataverse_files.zip or the AMPds2 folder.")
    print("Reading AMPds2 from:", inp)
    raw = find_power_table(inp)
    print(f"  raw shape: {raw.shape}")
    hourly = to_hourly(raw)
    os.makedirs(os.path.dirname(OUT), exist_ok=True)
    hourly.index.name = "ts"
    hourly.to_parquet(OUT)
    mb = os.path.getsize(OUT) / 1e6
    print(f"\nWrote {OUT}  ({hourly.shape[0]} hours x {hourly.shape[1]} meters, {mb:.1f} MB)")
    print("Columns:", list(hourly.columns))
    print("Commit this parquet to your Space (or upload via the Files tab).")


if __name__ == "__main__":
    main()