# SPDX-License-Identifier: MIT
"""Merge the per-file drive summaries from extract.py into one row per drive.

    python merge.py out/ drives.csv.gz

For each (serial, model) across every file: the first and last day it was
seen, power-on hours on those days, days observed, the first day it reported
failure=1 and how many such days there were.
"""
import json
import sys
from pathlib import Path

import numpy as np
import pandas as pd


def main(src, dst):
    parts = []
    for f in sorted(Path(src).glob("*.csv.gz")):
        d = pd.read_csv(f, dtype={"serial": str, "model": str})
        parts.append(d)
    d = pd.concat(parts, ignore_index=True)
    # some files write "WDC  WUH721414ALE6L4" with a double space
    d["model"] = d["model"].str.split().str.join(" ")
    d["fail_day"] = d["fail_day"].where(d["fail_day"] >= 0, np.nan)

    d = d.sort_values("first_day")
    g = d.groupby(["serial", "model"], sort=False)
    first = g.first()[["capacity", "first_day", "first_poh"]]
    cap_unknown_first_seen = int((first.capacity <= 0).sum())
    d = d.sort_values("last_day")
    g = d.groupby(["serial", "model"], sort=False)
    last = g.last()[["last_day", "last_poh"]]
    sums = g.agg(n_obs=("n_obs", "sum"), n_fail=("n_fail", "sum"), fail_day=("fail_day", "min"),
                 capacity=("capacity", "max"))
    out = first.drop(columns="capacity").join(last).join(sums).reset_index()
    # capacity_bytes is sometimes logged as -1; fall back to the model's usual value
    known = out[out.capacity > 0].groupby("model").capacity.median()
    unk = out.capacity <= 0
    out.loc[unk, "capacity"] = out.loc[unk, "model"].map(known).fillna(-1)
    print(f"capacity unknown on every day for {int(unk.sum())} drives; "
          f"{int((out.capacity <= 0).sum())} still unknown after the model fallback")
    out["fail_day"] = out["fail_day"].fillna(-1).astype(int)
    report = {"drives": int(len(out)), "capacity_unknown_when_first_seen": cap_unknown_first_seen,
              "capacity_unknown_every_day": int(unk.sum())}
    Path(dst).with_name("merge_report.json").write_text(json.dumps(report, indent=1))
    out.to_csv(dst, index=False, compression="gzip")
    print(f"{len(out)} drives, {int((out.fail_day >= 0).sum())} with a failure record")


if __name__ == "__main__":
    main(sys.argv[1], sys.argv[2])