# SPDX-License-Identifier: MIT
"""Reduce Backblaze Drive Stats daily snapshots to one row per drive.

    python extract.py data_Q1_2024.zip out/Q1_2024.csv.gz

Reads each daily CSV inside a quarterly (or yearly) zip, keeping only date,
serial_number, model, capacity_bytes, failure and smart_9_raw (power-on
hours), and writes one row per (serial, model) seen in that file:

    serial, model, capacity, first_day, first_poh, last_day, last_poh,
    n_obs, fail_day, n_fail

Days are counted from 2013-04-10, the first day in the data set. fail_day is
-1 if the drive never reported failure=1. Memory stays small because each day
is folded into the running summary before the next is read.
"""
import sys, zipfile, io
import numpy as np
import pandas as pd

EPOCH = pd.Timestamp("2013-04-10")
KEEP = {"date", "serial_number", "model", "capacity_bytes", "failure", "smart_9_raw"}


def day_frames(path):
    with zipfile.ZipFile(path) as z:
        names = sorted(n for n in z.namelist()
                       if n.endswith(".csv") and "__MACOSX" not in n)
        for n in names:
            with z.open(n) as fh:
                df = pd.read_csv(fh, usecols=lambda c: c in KEEP,
                                 dtype={"serial_number": str, "model": str},
                                 low_memory=False)
            if df.empty:
                continue
            day = (pd.Timestamp(df["date"].iloc[0]) - EPOCH).days
            df["model"] = df["model"].str.split().str.join(" ")
            df["serial_number"] = df["serial_number"].str.strip()
            df = df.drop_duplicates(["serial_number", "model"])
            df = df.set_index(["serial_number", "model"])
            yield day, df


def reduce_zip(path):
    agg = None
    for day, df in day_frames(path):
        poh = pd.to_numeric(df["smart_9_raw"], errors="coerce") if "smart_9_raw" in df else np.nan
        cur = pd.DataFrame({
            "capacity": df["capacity_bytes"].values,
            "first_day": day, "first_poh": poh,
            "last_day": day, "last_poh": poh,
            "n_obs": 1,
            "fail_day": np.where(df["failure"].values == 1, day, -1),
            "n_fail": (df["failure"].values == 1).astype(int),
        }, index=df.index)
        if agg is None:
            agg = cur
            continue
        new = cur.index.difference(agg.index)
        old = cur.index.intersection(agg.index)
        c = cur.loc[old]
        # capacity_bytes is sometimes -1 on a given day; keep the largest seen
        agg.loc[old, "capacity"] = np.maximum(agg.loc[old, "capacity"].values, c["capacity"].values)
        agg.loc[old, "last_day"] = day
        agg.loc[old, "last_poh"] = c["last_poh"].values
        agg.loc[old, "n_obs"] += 1
        hit = c.index[c["n_fail"].values == 1]
        if len(hit):
            agg.loc[hit, "n_fail"] += 1
            first = agg.loc[hit, "fail_day"] < 0
            agg.loc[hit[first.values], "fail_day"] = day
        if len(new):
            agg = pd.concat([agg, cur.loc[new]])
    agg.index.names = ["serial", "model"]
    return agg


if __name__ == "__main__":
    src, dst = sys.argv[1], sys.argv[2]
    out = reduce_zip(src)
    out.to_csv(dst, compression="gzip")
    print(f"{src}: {len(out)} drives, {int((out.n_fail > 0).sum())} failed")