Failures after maintenance

"""Failures after maintenance: which recent work explains a failure, and who is next?

Simulates 150 days of work orders for a fleet of 900 centrifugal pumps at four
sites. Each work order records the part replaced (part number, lot, supplier),
the crew, the procedure and the site's operating conditions. A firmware update
rolls out to the controllers mid-period. Pumps fail now and then. HeapDB ranks
the recent work behind each failure (top_causes), learns across the fleet
(link_report), and flags pumps that got the suspect work and have not failed
yet (top_causes at "now").

Hidden ground truth (the seeded cause):

* lot:B-2291 of bearing part:6205-bearing from supplier:kestrel is bad. It was
  installed between day 40 and day 80. Each install fails with probability
  0.40, after a gamma-distributed delay (shape 2, scale 6 days: mode 6 days,
  mean 12 days). Kestrel's other bearing lots are fine.
* Any work order carries a small infant-mortality risk (2%), whatever was done.
* Background failures: 0.08% per pump-day, three times that at site:S3.

Distractors:

* Confounder: site:S3 runs high duty (duty:high) and fails more in general.
* Near-miss timing: firmware fw:4.2.0 rolls out from day 45 to day 75, in the
  same weeks as the bad lot. It is harmless.
* Busy crew: crew:C1 does about 45% of all work orders. It is harmless.

Learning: every failure is observe()d (no known cause). One in four failures
before "now" gets a teardown; its finding goes to feedback() (the work order
that caused it is confirmed, the other recent ones are rejected). Each week,
settle(partial=True) adds negative evidence for work without a failure since,
scaled by how much of its window (the kernel's mass) has passed.

Run:
    pip install heapdb
    python 0007-failures-after-maintenance.py
"""

from __future__ import annotations

import calendar
import random

from heapdb import GammaKernel, HeapDB, format_result

M, H, D = 60.0, 3600.0, 86400.0
SEED = 11
N_PUMPS = 900
DAYS = 150
NOW_DAY = 90                    # the analyst's "now"; later failures are hidden from HeapDB
LOOKAHEAD = 30                  # days after now used to score the proactive list
T0 = float(calendar.timegm((2026, 3, 2, 0, 0, 0)))
NOW = T0 + NOW_DAY * D

EFFECT = "unplanned_failure"
BAD_LOT = "lot:B-2291"
BAD_WINDOW = (40, 80)           # days in which the bad lot was being installed
P_BAD = 0.40
P_INFANT = 0.02
BASE_RATE = 0.0008              # per pump-day
FW = "fw:4.2.0"

SITES = {  # site -> (operating-condition features, background multiplier)
    "site:S1": ({"duty:normal"}, 1.0),
    "site:S2": ({"duty:normal", "env:dusty"}, 1.0),
    "site:S3": ({"duty:high"}, 3.0),
    "site:S4": ({"duty:normal"}, 1.0),
}
CREWS = ["crew:C1"] * 9 + ["crew:C2", "crew:C3", "crew:C4", "crew:C5",
                           "crew:C6", "crew:C7", "crew:C8", "crew:C9", "crew:C10",
                           "crew:C11", "crew:C12"]
# part -> list of (lot, supplier); BAD_LOT only enters during BAD_WINDOW
PARTS = {
    "part:6205-bearing": [("lot:B-2288", "supplier:kestrel"), ("lot:B-2290", "supplier:kestrel"),
                          ("lot:N-7714", "supplier:norrland"), ("lot:N-7731", "supplier:norrland")],
    "part:SL-40-seal": [("lot:S-0412", "supplier:sealtek"), ("lot:S-0419", "supplier:sealtek")],
    "part:IMP-9-impeller": [("lot:I-1180", "supplier:kestrel"), ("lot:I-1192", "supplier:castor")],
    "part:CPL-2-coupling": [("lot:C-3301", "supplier:castor")],
}
PROCS = {"part:6205-bearing": "proc:BRG-REPL-03", "part:SL-40-seal": "proc:SEAL-REPL-01",
         "part:IMP-9-impeller": "proc:IMP-REPL-02", "part:CPL-2-coupling": "proc:CPL-ALIGN-01"}


def gamma_lag(rng: random.Random) -> float:
    return rng.gammavariate(2.0, 6.0) * D


def simulate(rng: random.Random):
    """Return (events, failures, truth).

    events: (pump, t, type, features); failures: (t, pump, cause) where cause is
    "lot", "infant" or "background"; truth: index of the event behind a
    work-order-caused failure, keyed by (pump, failure time).
    """
    events, failures, truth = [], [], {}
    for p in range(N_PUMPS):
        pump = f"pump-{p:04d}"
        site = rng.choice(list(SITES))
        conds, mult = SITES[site]
        evs = []
        t = T0 + rng.uniform(0, 25) * D
        while t < T0 + DAYS * D:
            day = (t - T0) / D
            crew = rng.choice(CREWS)
            if rng.random() < 0.35:   # inspection only, no part
                feats = {"proc:PM-INSPECT", crew} | conds
            else:
                part = rng.choices(list(PARTS), weights=[5, 3, 1, 1])[0]
                lot, sup = rng.choice(PARTS[part])
                if part == "part:6205-bearing" and BAD_WINDOW[0] <= day < BAD_WINDOW[1] \
                        and rng.random() < 0.45:
                    lot, sup = BAD_LOT, "supplier:kestrel"
                feats = {part, lot, sup, PROCS[part], crew} | conds
            evs.append((t, "work_order", feats))
            t += rng.uniform(12, 40) * D
        if rng.random() < 0.6:        # firmware rollout (near-miss timing)
            evs.append((T0 + rng.uniform(45, 75) * D, "firmware", {FW}))
        evs.sort(key=lambda e: e[0])
        base = len(events)
        events += [(pump, t, typ, feats) for t, typ, feats in evs]
        for i, (t, typ, feats) in enumerate(evs):
            if typ != "work_order":
                continue
            if BAD_LOT in feats and rng.random() < P_BAD:
                tf = t + gamma_lag(rng)
                failures.append((tf, pump, "lot"))
                truth[(pump, tf)] = base + i
            elif rng.random() < P_INFANT:
                tf = t + gamma_lag(rng)
                failures.append((tf, pump, "infant"))
                truth[(pump, tf)] = base + i
        day = 0.0
        while True:                   # background failures (Poisson)
            day += rng.expovariate(BASE_RATE * mult)
            if day >= DAYS:
                break
            failures.append((T0 + day * D, pump, "background"))
    failures = [f for f in failures if f[0] < T0 + DAYS * D]
    failures.sort()
    return events, failures, truth


def make_db() -> HeapDB:
    db = HeapDB(bucket_seconds=2 * D)
    # Infant mortality after work: hazard rises over the first days, peaks
    # around a week (mode = (shape-1)*scale = 6 d), tail cut at mean + 4 std (~46 d).
    db.register_effect(EFFECT, GammaKernel(shape=2.0, scale=6 * D, tail_stds=4.0),
                       default_prior=(1.0, 9.0), leak=0.10)
    # An inspection with no part replaced rarely introduces a defect.
    db.add_rule(EFFECT, "inspection-only", lambda e: "proc:PM-INSPECT" in e.features, 0.3,
                types={"work_order"})
    return db


def main() -> None:
    rng = random.Random(SEED)
    events, failures, truth = simulate(rng)
    db = make_db()
    ids = {}                          # event index -> HeapDB event id
    for i, (pump, t, typ, feats) in enumerate(events):
        if t <= NOW:
            ids[i] = db.insert(pump, t, typ, feats).id
    past = [f for f in failures if f[0] <= NOW]

    teardowns = 0
    next_settle = T0 + 7 * D
    for tf, pump, cause in past:
        while next_settle < tf:
            db.settle(next_settle, partial=True)
            next_settle += 7 * D
        db.observe(pump, EFFECT, tf)
        if rng.random() < 0.25:       # teardown / root-cause finding
            teardowns += 1
            cand = db.top_causes(pump, EFFECT, tf, k=4)
            conf = ids.get(truth.get((pump, tf), -1))
            rej = [c.event.id for c in cand if c.event.id != conf]
            db.feedback(pump, EFFECT, tf, confirmed_event_id=conf, rejected_ids=rej)
    db.settle(NOW, partial=True)

    n_bad = sum(1 for _, t, _, f in events if BAD_LOT in f and t <= NOW)
    n_wo = sum(1 for _, t, typ, _ in events if typ == "work_order" and t <= NOW)
    by = {c: sum(1 for f in past if f[2] == c) for c in ("lot", "infant", "background")}
    print("SCENARIO")
    print(f"  {N_PUMPS} pumps, 4 sites, {NOW_DAY} days to now: {n_wo:,} work orders "
          f"({n_bad} installed {BAD_LOT}), firmware rollout {FW}.")
    print(f"  {len(past)} unplanned failures observed ({by['lot']} from the bad lot, "
          f"{by['infant']} other infant mortality, {by['background']} background); "
          f"{teardowns} teardowns.")

    # One failure, explained: the latest lot-caused failure with the firmware in its window.
    pick = None
    for tf, pump, cause in reversed(past):
        if cause != "lot":
            continue
        res = db.top_causes(pump, EFFECT, tf, k=3, explain=True)
        if any(FW in c.event.features for c in res):
            pick = (pump, tf, res)
            break
    if pick is None:
        raise SystemExit(1)
    pump, tf, res = pick
    print()
    print(f"ONE FAILURE, EXPLAINED ({pump}, day {(tf - T0) / D:.1f})")
    print(format_result(res))

    print()
    print("ACROSS ALL FAILURES (link_report, top 8)")
    rep = db.link_report(EFFECT, min_support=10)
    lines = str(rep).splitlines()
    print("\n".join(lines[:2 + 8]))
    order = [r.feature for r in rep]
    for f in (FW, "crew:C1", "duty:high"):
        print(f"  ... {f:<34} #{order.index(f) + 1} of {len(order)}, mean {rep.row(f).strength:.3f}")

    # Proactive: pumps that got the bad lot within the kernel window and have not failed
    # since. Risk = feature_risk of that work order x kernel mass still ahead of it.
    kern = db.effect(EFFECT).kernel
    flagged = []
    for i, (p, t, typ, feats) in enumerate(events):
        if BAD_LOT not in feats or not NOW - kern.support < t <= NOW:
            continue
        if any(p == fp and t < ft <= NOW for ft, fp, _ in past):
            continue
        risk = db.feature_risk(EFFECT, feats, entity=p, type=typ)
        flagged.append((risk * (1.0 - kern.mass_fraction(NOW - t)), p, (NOW - t) / D, i))
    flagged.sort(key=lambda r: (-r[0], r[1]))

    def fails_later(p: str) -> bool:
        return any(fp == p and NOW < ft <= NOW + LOOKAHEAD * D for ft, fp, _ in failures)

    print()
    print(f"PROACTIVE: {len(flagged)} PUMPS WITH {BAD_LOT} IN WINDOW, NOT YET FAILED, BY RISK")
    for score, p, age, _ in flagged[:5]:
        print(f"  {p}  installed {age:4.1f} d ago  risk={score:.4f}  "
              f"fails in next {LOOKAHEAD} d: {'yes' if fails_later(p) else 'no'}")
    half = len(flagged) // 2
    top_n = sum(fails_later(p) for _, p, _, _ in flagged[:half])
    bot_n = sum(fails_later(p) for _, p, _, _ in flagged[half:])
    others = sorted({e[0] for e in events} - {p for _, p, _, _ in flagged})
    oth_n = sum(fails_later(p) for p in others)
    print(f"  in the next {LOOKAHEAD} d (simulation): {top_n + bot_n} of {len(flagged)} flagged pumps fail "
          f"(top half {top_n} of {half}, bottom half {bot_n} of {len(flagged) - half});")
    print(f"  {oth_n} of {len(others)} pumps not flagged fail "
          f"({oth_n / len(others):.1%} vs {(top_n + bot_n) / max(len(flagged), 1):.1%} flagged).")

    print()
    print("CHECK")
    rank = order.index(BAD_LOT) + 1
    print(f"  {BAD_LOT} ranks #{rank} of {len(order)} features in link_report "
          f"(mean {rep.row(BAD_LOT).strength:.3f}); {FW} ranks #{order.index(FW) + 1}.")
    if rank != 1:
        raise SystemExit(1)


if __name__ == "__main__":
    main()