# SPDX-License-Identifier: MIT
"""YS-EN-0005: what a 40-year MTBF promises, checked against Backblaze Drive Stats.

    python analyse.py drives.csv.gz        # writes results.json next to this file

Input is the one-row-per-drive table from extract.py + merge.py. Everything
the note quotes is computed here and written to results.json.

Method
------
Each drive is followed from the first day it appears to the day it first
reports failure=1 (an event) or the last day it appears (censored). Drives
that were already running when the records begin (2013-04-10) enter at the
age given by SMART attribute 9, power-on hours, on that day, and contribute
exposure from that age onwards only (left truncation). Later drives enter at
age zero on their first day, which is Backblaze's own convention: SMART 9 can
be corrupted or non-zero on new drives. AGE_MODE = "smart9" uses SMART 9 for
every drive instead, as a sensitivity check. Hazard in an age bin is failures divided by
drive-years spent in that bin, i.e. Backblaze's annualized failure rate
computed per age instead of per calendar period.
"""
from __future__ import annotations

import json
import math
import sys
from pathlib import Path

import numpy as np
import pandas as pd
from scipy import optimize, stats
from scipy.special import gamma as G

HERE = Path(__file__).resolve().parent
YEAR = 365.25

# ---------------------------------------------------------------- inputs
MTBF_YR = 40.0          # Model 4077 datasheet
FLEET = 225             # cores aboard the SS Mariposa
MIN_DRIVES = 2000       # per-model fits: models with at least this many drives
MIN_FAILS = 50          # ... and at least this many failures
BIN_MONTHS = 3          # width of the age bins for the fleet hazard curve
MAX_AGE_YR = 12.0       # fleet curve stops where exposure gets thin
# models shown one by one: one near-constant, one rising, one falling, and the
# Exos X24 whose datasheet we quote
EXEMPLARS = ["ST4000DM000", "HGST HUH721212ALN604", "HGST HMS5C4040BLE640", "ST24000NM002H"]
AGE_MODE = "backblaze"  # or "smart9"


# ---------------------------------------------------------------- MTBF arithmetic
def mtbf_arithmetic():
    lam = 1 / MTBF_YR
    out = {
        "mtbf_yr": MTBF_YR,
        "mtbf_hours": MTBF_YR * 8760,   # vendors' 8,760 power-on hours a year
        "afr_exponential": 1 - math.exp(-lam),
        "afr_rate": lam,
        "survive_to_mtbf": math.exp(-1),
        "unrepaired_failures_40yr": FLEET * (1 - math.exp(-1)),
        "unrepaired_failures_10yr": FLEET * (1 - math.exp(-10 * lam)),
        "repaired_failures_40yr": FLEET * 40 * lam,
        "median_life_yr": MTBF_YR * math.log(2),
    }
    # same mean life, different Weibull shapes: R(t) = exp(-(t/eta)^beta), eta = MTBF / Gamma(1 + 1/beta)
    shapes = {}
    for b in (0.66, 1.0, 1.45, 2.18, 4.0):   # shapes fitted to real drive models below
        eta = MTBF_YR / G(1 + 1 / b)
        R = lambda t: math.exp(-((t / eta) ** b))
        H = lambda t: (t / eta) ** b
        h = lambda t: b / eta * (t / eta) ** (b - 1)
        # mean residual life at age 30: integral_30^inf R(t) dt / R(30)
        from scipy.integrate import quad
        mrl30 = quad(R, 30, np.inf, limit=200)[0] / R(30)
        shapes[f"{b:g}"] = {
            "eta_yr": eta,
            "fail_by_10": FLEET * (1 - R(10)),
            "fail_by_40": FLEET * (1 - R(40)),
            "survive_to_mtbf": R(MTBF_YR),
            "afr_year1": 1 - R(1),
            "afr_year30": 1 - R(31) / R(30),
            "mrl_at_30": mrl30,
            "hazard_ratio_30_vs_1": h(30) / h(1) if b != 1 else 1.0,
        }
    out["weibull_same_mtbf"] = shapes
    return out


def afr_to_mtbf_hours(afr):
    """Backblaze-style AFR (failures per drive-year) to MTBF in hours."""
    return YEAR * 24 / afr


# ---------------------------------------------------------------- data
def load(path):
    d = pd.read_csv(path, dtype={"serial": str, "model": str})
    n0 = len(d)
    rep = {"rows_in": n0}

    # hard drives only: boot SSDs carry "SSD" or are small flash models
    ssd = d.model.str.contains(r"SSD|DELLBOSS|Micron|MTFDD|WDS\d|^CT\d|SATA SSD|ZA\d{3}", case=False, regex=True)
    small = d.capacity < 1e12
    d = d[~ssd & ~small].copy()
    rep["dropped_ssd_or_small"] = int(n0 - len(d))

    d["end_day"] = np.where(d.fail_day >= 0, d.fail_day, d.last_day)
    d["event"] = (d.fail_day >= 0).astype(int)
    rep["failed_then_seen_again"] = int(((d.fail_day >= 0) & (d.last_day > d.fail_day)).sum())
    rep["multiple_failure_days"] = int((d.n_fail > 1).sum())

    # SMART 9 is only trusted where it is needed: for drives present on day 0
    # (or for every drive in "smart9" mode). Power-on hours must then be
    # plausible and advance about 24 h per calendar day; this rejects drives
    # whose SMART 9 is in other units, wraps, or is missing.
    need = (d.first_day == 0) if AGE_MODE == "backblaze" else pd.Series(True, index=d.index)
    span = d.last_day - d.first_day
    ratio = (d.last_poh - d.first_poh) / (24 * span.where(span > 0))
    ok_poh = d.first_poh.notna() & (d.first_poh >= 0) & (d.first_poh < 15 * YEAR * 24)
    ok_rate = (span < 30) | ratio.between(0.8, 1.1)
    bad = need & ~(ok_poh & ok_rate)
    rep["age_mode"] = AGE_MODE
    rep["present_on_day0"] = int((d.first_day == 0).sum())
    rep["dropped_bad_poh"] = int(bad.sum())
    rep["dropped_bad_poh_failures"] = int(d.event[bad].sum())
    d = d[~bad].copy()
    need = need[~bad]

    late = d.first_day > 0
    late_days = d.first_poh[late] / 24
    rep["later_drives_poh_over_30d_share"] = float((late_days > 30).mean())
    rep["later_drives_poh_over_30d_median_days"] = float(late_days[late_days > 30].median())
    rep["later_drives_poh_over_1yr_share"] = float((late_days > YEAR).mean())
    rep["drives_missing_over_10pct_of_days_share"] = float((d.n_obs / (d.last_day - d.first_day + 1) < 0.9).mean())
    mr = Path(path).with_name("merge_report.json")
    if mr.exists():
        rep["merge"] = json.loads(mr.read_text())

    d["entry"] = np.where(need, d.first_poh / 24 / YEAR, 0.0)
    d["exit"] = d.entry + (d.end_day - d.first_day + 1) / YEAR
    d["manufacturer"] = d.model.map(maker)
    rep["drives"] = int(len(d))
    rep["failures"] = int(d.event.sum())
    rep["drive_years"] = float((d.exit - d.entry).sum())
    rep["models"] = int(d.model.nunique())
    rep["first_day"] = str(pd.Timestamp("2013-04-10") + pd.Timedelta(days=int(d.first_day.min())))
    rep["last_day"] = str(pd.Timestamp("2013-04-10") + pd.Timedelta(days=int(d.last_day.max())))
    rep["entered_at_age_over_30d"] = int((d.entry > 30 / YEAR).sum())
    return d, rep


def maker(m):
    u = m.upper()
    if u.startswith(("ST", "SEAGATE")):
        return "Seagate"
    if u.startswith(("HGST", "HITACHI")):
        return "HGST"
    if u.startswith(("WDC", "WD")):
        return "WDC"
    if u.startswith("TOSHIBA") or u.startswith(("MG0", "MD0")):
        return "Toshiba"
    return "other"


# ---------------------------------------------------------------- hazard by age
def exposure_by_bin(entry, exit_, event, edges):
    expo = np.empty(len(edges) - 1)
    fails = np.empty(len(edges) - 1)
    for i, (a, b) in enumerate(zip(edges[:-1], edges[1:])):
        expo[i] = np.clip(np.minimum(exit_, b) - np.maximum(entry, a), 0, None).sum()
        fails[i] = (event & (exit_ > a) & (exit_ <= b)).sum()
    return expo, fails


def poisson_ci(k, expo, conf=0.95):
    a = (1 - conf) / 2
    lo = np.where(k > 0, stats.chi2.ppf(a, 2 * k) / 2, 0.0) / expo
    hi = stats.chi2.ppf(1 - a, 2 * k + 2) / 2 / expo
    return lo, hi


def hazard_curve(d, bin_months=BIN_MONTHS, max_age=MAX_AGE_YR):
    edges = np.arange(0, max_age + 1e-9, bin_months / 12)
    expo, fails = exposure_by_bin(d.entry.values, d.exit.values, d.event.values.astype(bool), edges)
    with np.errstate(divide="ignore", invalid="ignore"):
        h = fails / expo
    lo, hi = poisson_ci(fails, np.where(expo > 0, expo, np.nan))
    return {"edges": edges.tolist(), "drive_years": expo.tolist(), "failures": fails.astype(int).tolist(),
            "afr": h.tolist(), "lo": lo.tolist(), "hi": hi.tolist()}


def first_month_curve(d):
    edges = np.arange(0, 1 + 1e-9, 1 / 12)
    expo, fails = exposure_by_bin(d.entry.values, d.exit.values, d.event.values.astype(bool), edges)
    return {"edges": edges.tolist(), "drive_years": expo.tolist(), "failures": fails.astype(int).tolist(),
            "afr": (fails / expo).tolist()}


# ---------------------------------------------------------------- Weibull vs exponential
@np.errstate(over="ignore", invalid="ignore")
def weibull_fit(entry, exit_, event):
    """MLE for a Weibull with left truncation and right censoring.
    log L = sum delta*log h(x) - sum [H(x) - H(e)],  H(t) = (t/eta)^beta."""
    e, x, dl = entry, exit_, event.astype(float)
    nz = e > 0
    lx = np.log(x)
    le = np.where(nz, np.log(np.where(nz, e, 1.0)), 0.0)

    def nll(p):
        lb, leta = p
        b = math.exp(lb)
        z_x = np.exp(b * (lx - leta))
        z_e = np.where(nz, np.exp(b * (le - leta)), 0.0)
        ll = (dl * (lb - leta + (b - 1) * (lx - leta))).sum() - (z_x - z_e).sum()
        return -ll if np.isfinite(ll) else 1e300

    lam = dl.sum() / (x - e).sum()
    r = optimize.minimize(nll, [0.0, math.log(1 / lam)], method="Nelder-Mead",
                          options={"xatol": 1e-7, "fatol": 1e-7, "maxiter": 4000})
    r = optimize.minimize(nll, r.x, method="BFGS")
    ll_w = -r.fun
    ll_e = (dl * math.log(lam)).sum() - lam * (x - e).sum()
    lb, leta = r.x
    # profile-likelihood 95% CI on beta
    def prof(lb_):
        return optimize.minimize_scalar(lambda le_: nll([lb_, le_]), bracket=(leta - 1, leta + 1)).fun
    crit = stats.chi2.ppf(0.95, 1) / 2
    f = lambda lb_: prof(lb_) - (-ll_w) - crit
    lo = optimize.brentq(f, lb - 3, lb) if f(lb - 3) > 0 else float("nan")
    hi = optimize.brentq(f, lb, lb + 3) if f(lb + 3) > 0 else float("nan")
    lrt = 2 * (ll_w - ll_e)
    return {"beta": math.exp(lb), "beta_lo": math.exp(lo), "beta_hi": math.exp(hi),
            "eta_yr": math.exp(leta), "lambda_exp": lam, "lrt": lrt,
            "p_exponential": float(stats.chi2.sf(max(lrt, 0), 1))}


def per_model(d):
    rows = []
    for m, g in d.groupby("model"):
        if len(g) < MIN_DRIVES or g.event.sum() < MIN_FAILS:
            continue
        fit = weibull_fit(g.entry.values, g.exit.values, g.event.values)
        dy = float((g.exit - g.entry).sum())
        # first-year rate, extrapolated as a constant, vs observed survival later on
        hz = hazard_curve(g, bin_months=3, max_age=MAX_AGE_YR)
        edges = np.array(hz["edges"])
        expo = np.array(hz["drive_years"])
        fails = np.array(hz["failures"])
        y1 = edges[1:] <= 1 + 1e-9
        rate_y1 = fails[y1].sum() / expo[y1].sum() if expo[y1].sum() > 0 else float("nan")
        # piecewise-constant hazard -> cumulative hazard; only where the bin has real exposure
        enough = expo >= 200
        reach = edges[1:][enough].max() if enough.any() else 0.0
        H = np.cumsum(np.where(enough, fails / np.where(expo > 0, expo, 1) * (edges[1:] - edges[:-1]), np.nan))
        surv = {}
        for A in (3, 5, 7):
            i = np.searchsorted(edges[1:], A - 1e-9)
            if i < len(H) and np.isfinite(H[: i + 1]).all() and A <= reach:
                surv[str(A)] = {"observed_failed": float(1 - math.exp(-H[i])),
                                "first_year_rate_predicts": float(1 - math.exp(-rate_y1 * A))}
        rows.append({
            "model": m, "manufacturer": maker(m), "capacity_tb": round(float(g.capacity.median()) / 1e12, 1),
            "drives": int(len(g)), "failures": int(g.event.sum()), "drive_years": dy,
            "afr": float(g.event.sum() / dy),
            "max_age_yr": float(g.exit.max()), "median_exit_age_yr": float(g.exit.median()),
            "first_year_afr": float(rate_y1), "reach_yr": float(reach), "cum_failed": surv,
            **fit,
        })
    rows.sort(key=lambda r: -r["drive_years"])
    return rows


def main(path, mode=None):
    global AGE_MODE
    if mode:
        AGE_MODE = mode
    d, rep = load(path)
    res = {"inputs": {"mtbf_yr": MTBF_YR, "fleet": FLEET, "min_drives": MIN_DRIVES, "min_fails": MIN_FAILS,
                      "bin_months": BIN_MONTHS, "max_age_yr": MAX_AGE_YR},
           "mtbf": mtbf_arithmetic(), "data": rep}
    res["fleet_afr"] = rep["failures"] / rep["drive_years"]
    res["fleet_afr_as_mtbf_hours"] = afr_to_mtbf_hours(res["fleet_afr"])
    res["hazard_all"] = hazard_curve(d)
    res["hazard_first_year"] = first_month_curve(d)
    res["hazard_by_maker"] = {m: hazard_curve(g) for m, g in d.groupby("manufacturer") if m != "other"}
    h = res["hazard_all"]
    H = np.cumsum(np.nan_to_num(np.array(h["afr"])) * np.diff(h["edges"]))
    res["pooled_survival"] = {str(a): float(math.exp(-H[int(round(a * 12 / BIN_MONTHS)) - 1])) for a in (1, 2, 4, 6, 8, 10)}
    res["models"] = per_model(d)
    res["hazard_by_model"] = {m: hazard_curve(d[d.model == m], bin_months=12) for m in EXEMPLARS}
    for m in res["models"]:
        m["mtbf_yr_observed"] = 1 / m["afr"]
    ms = res["models"]
    res["model_summary"] = {
        "n": len(ms),
        "beta_gt1_sig": sum(1 for m in ms if m["beta_lo"] > 1),
        "beta_lt1_sig": sum(1 for m in ms if m["beta_hi"] < 1),
        "beta_ci_covers_1": sum(1 for m in ms if m["beta_lo"] <= 1 <= m["beta_hi"]),
        "drives_covered": sum(m["drives"] for m in ms),
        "failures_covered": sum(m["failures"] for m in ms),
    }
    out = "results.json" if AGE_MODE == "backblaze" else f"results_{AGE_MODE}.json"
    (HERE / out).write_text(json.dumps(res, indent=1, default=float))
    print(json.dumps({k: res[k] for k in ("data", "fleet_afr", "model_summary")}, indent=1))


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