"""
§3.3 — ROC-AUC of severity vs exploitation-likelihood as predictors of
realised exploitation (CISA KEV membership).

- EPSS-AUC: rank-AUC of the continuous EPSS score over every CVE that has an
  EPSS score, labelled KEV / not-KEV. (Caveat: current EPSS is *informed by*
  observed exploitation, so this is closer to detection than time-of-disclosure
  prediction — stated honestly in the paper.)
- CVSS-AUC: severity-banded AUC computed analytically from the §3.1 4x2 table
  (we do not have per-CVE continuous base scores at population scale).
- Censoring sensitivity: repeat EPSS-AUC excluding CVE-2025/2026 IDs (too recent
  to have been exploited-and-listed), using CVE-ID year as a transparent proxy.
"""
import gzip, json, csv

# ---------- load EPSS ----------
epss = {}
score_date = None
with gzip.open("epss.csv.gz", "rt") as f:
    first = f.readline()
    if "score_date" in first:
        for tok in first.strip("#\n").split(","):
            if "score_date" in tok:
                score_date = tok.split(":", 1)[1]
    else:
        f.seek(0)
    rdr = csv.DictReader(f)
    for row in rdr:
        try:
            epss[row["cve"]] = float(row["epss"])
        except (KeyError, ValueError):
            pass

# ---------- load KEV ----------
kev = set()
with open("kev.json", "r", encoding="utf-8") as f:
    kj = json.load(f)
for v in kj.get("vulnerabilities", []):
    cid = v.get("cveID")
    if cid:
        kev.add(cid)

print(f"EPSS score_date: {score_date}")
print(f"EPSS CVEs: {len(epss):,}   KEV CVEs (total): {len(kev):,}")
kev_in_epss = sum(1 for c in kev if c in epss)
print(f"KEV CVEs covered by EPSS: {kev_in_epss:,}")

def rank_auc(pairs):
    """pairs: list of (score, label). Returns AUC via mean-rank Mann-Whitney."""
    pairs = sorted(pairs, key=lambda x: x[0])
    n = len(pairs)
    # assign average ranks (1-based) handling ties
    ranks = [0.0] * n
    i = 0
    while i < n:
        j = i
        while j + 1 < n and pairs[j + 1][0] == pairs[i][0]:
            j += 1
        avg = (i + 1 + j + 1) / 2.0
        for k in range(i, j + 1):
            ranks[k] = avg
        i = j + 1
    nP = sum(1 for _, l in pairs if l == 1)
    nN = n - nP
    if nP == 0 or nN == 0:
        return None, nP, nN
    sum_ranks_pos = sum(r for r, (_, l) in zip(ranks, pairs) if l == 1)
    auc = (sum_ranks_pos - nP * (nP + 1) / 2.0) / (nP * nN)
    return auc, nP, nN

# ---------- EPSS-AUC (full) ----------
data = [(s, 1 if c in kev else 0) for c, s in epss.items()]
auc_full, nP, nN = rank_auc(data)
print(f"\n[Full] EPSS-AUC vs KEV: {auc_full:.4f}  (positives={nP:,}, negatives={nN:,})")

# ---------- EPSS-AUC (censored: drop 2025/2026 CVE-ID years) ----------
def id_year(cve):
    try:
        return int(cve.split("-")[1])
    except Exception:
        return 0
data_c = [(s, 1 if c in kev else 0) for c, s in epss.items() if id_year(c) <= 2024]
auc_cens, nPc, nNc = rank_auc(data_c)
print(f"[Censored <=2024] EPSS-AUC vs KEV: {auc_cens:.4f}  (positives={nPc:,}, negatives={nNc:,})")

# ---------- CVSS severity-banded AUC (analytic, from §3.1) ----------
# order: Critical > High > Medium > Low
bands = ["C", "H", "M", "L"]
P = {"C": 475, "H": 654, "M": 134, "L": 3}                 # exploited (KEV)
POP = {"C": 30301, "H": 75456, "M": 82119, "L": 3203}
N = {b: POP[b] - P[b] for b in bands}
order = {b: i for i, b in enumerate(bands)}                 # 0=highest severity
nP_b = sum(P.values()); nN_b = sum(N.values())
concordant = 0.0
for bi in bands:          # exploited band
    for bj in bands:      # non-exploited band
        if order[bi] < order[bj]:      # exploited MORE severe -> concordant
            concordant += P[bi] * N[bj]
        elif order[bi] == order[bj]:   # tie
            concordant += 0.5 * P[bi] * N[bj]
auc_cvss = concordant / (nP_b * nN_b)
print(f"\nCVSS severity-banded AUC vs KEV: {auc_cvss:.4f}")

# ---------- precise lifts (rounding-nit fix) ----------
base = nP_b / (nP_b + nN_b)
print("\nPrecise per-band lift (2 dp):")
for b, name in zip(bands, ["Critical", "High", "Medium", "Low"]):
    p = P[b] / POP[b]
    print(f"  {name:8} lift = {p/base:.2f}x")

json.dump({
    "epss_score_date": score_date,
    "epss_cve_count": len(epss),
    "kev_total": len(kev),
    "kev_in_epss": kev_in_epss,
    "epss_auc_full": round(auc_full, 4),
    "epss_auc_positives": nP,
    "epss_auc_censored_le2024": round(auc_cens, 4),
    "cvss_banded_auc": round(auc_cvss, 4),
    "base_rate_banded": base,
    "lifts": {name: round(P[b]/POP[b]/base, 2) for b, name in zip(bands, ["Critical","High","Medium","Low"])},
}, open("auc_results.json", "w"), indent=2)
print("\nSaved auc_results.json")