"""
Supportive Periodontal Therapy Clinical Examination Data
Reproduction of Figure 2 (PPD stability thresholds)

Python equivalent of 03_reproduce_figure_2.R. Reproduces the empirically
determined thresholds of no change of residual probing depth reported in
Figure 2 of Ramseier et al. 2019 (J Clin Periodontol 46:218-230), and reports a
sensitivity analysis excluding visits with no affected sites.

These twenty thresholds are the empirical basis of the SPT interval algorithm
implemented in 01_compute_spt_algorithm.py.

Requires: pandas, numpy. Run from the folder holding the data files:
    python 03_reproduce_figure_2.py
"""

import os
import sys

import numpy as np
import pandas as pd

if len(sys.argv) > 1:
    os.chdir(sys.argv[1])

FILE = "02_supportive_periodontal_therapy_data.csv"
if not os.path.exists(FILE):
    sys.exit(f"'{FILE}' not found in the working directory ({os.getcwd()}).\n"
             "  Run this script from the folder holding the data files.")

spt = pd.read_csv(FILE).sort_values(["pat_id", "spt_id"]).reset_index(drop=True)


# ---- 1. Cumulative percentages at each visit --------------------------------

n_sites = spt["n_teeth_spt"] * 6
PPD = ["n_4mm_spt", "n_5mm_spt", "n_6mm_spt", "n_from7mm_spt"]

spt["pct4"] = spt[PPD].sum(axis=1, min_count=4) * 100 / n_sites
spt["pct5"] = spt[PPD[1:]].sum(axis=1, min_count=3) * 100 / n_sites
spt["pct6"] = spt[PPD[2:]].sum(axis=1, min_count=2) * 100 / n_sites
spt["pct7"] = spt["n_from7mm_spt"] * 100 / n_sites


# ---- 2. Change relative to the preceding visit ------------------------------
# NOTE on the grouping variable.
# The figure caption states that visits are grouped by the residual PPD category
# "recorded at the previous SPT visit". The published values are reproduced only
# when visits are grouped by the percentage recorded at the RESPECTIVE (current)
# visit, with the change measured backwards to the preceding visit. The code
# below follows the computation that reproduces the published figure.

for k in ["4", "5", "6", "7"]:
    current = spt[f"pct{k}"]
    spt[f"chg{k}"] = current - current.groupby(spt["pat_id"]).shift(1)


# ---- 3. Interval categories -------------------------------------------------
# Real time between visits, categorised into 3, 4, 6, 9 and 12+ months.
# true_interval_days is 0 at each patient's first visit (no predecessor).

months = np.round(spt["true_interval_days"] / 30)
spt["ivl_cat"] = pd.cut(months, bins=[0, 3, 5, 8, 11, np.inf],
                        labels=["3", "4", "6", "9", "12+"], right=True)
spt.loc[spt["true_interval_days"] == 0, "ivl_cat"] = np.nan


# ---- 4. Percentage categories per panel -------------------------------------
# Upper bounds of the bins as printed on the x axis of Figure 2.

BINS = {"4": [10, 20, 30, 40, 100],
        "5": [2, 4, 6, 8, 10, 20, 30, 100],
        "6": [1, 2, 3, 4, 10, 100],
        "7": [1, 2, 3, 4, 10, 100]}

INTERVALS = ["3", "4", "6", "9", "12+"]


# ---- 5. Cell means and threshold extraction ---------------------------------
# The threshold is the largest bin whose mean change is still <= 0, i.e. the
# highest percentage of affected sites at which no increase of residual PPD is
# observed for that interval length.

def panel_table(k, exclude_zero=False):
    edges = BINS[k]
    current = spt[f"pct{k}"]
    change = spt[f"chg{k}"]
    keep = change.notna() & spt["ivl_cat"].notna()
    if exclude_zero:
        keep = keep & (current > 0)

    mean_rows, count_rows = [], []
    for j, edge in enumerate(edges):
        low = -1 if j == 0 else edges[j - 1]
        mean_row, count_row = [], []
        for interval in INTERVALS:
            sel = keep & (current > low) & (current <= edge) & (spt["ivl_cat"] == interval)
            mean_row.append(change[sel].mean() if sel.any() else np.nan)
            count_row.append(int(sel.sum()))
        mean_rows.append(mean_row)
        count_rows.append(count_row)

    index = [f"<={e}" for e in edges]
    means = pd.DataFrame(mean_rows, index=index, columns=INTERVALS, dtype=float)
    counts = pd.DataFrame(count_rows, index=index, columns=INTERVALS, dtype=int)
    return means, counts, edges


def thresholds(means, edges):
    out = {}
    for interval in INTERVALS:
        last_safe = np.nan
        for j in range(len(edges)):
            value = means.iloc[j][interval]
            if pd.isna(value):
                continue
            if value <= 0:
                last_safe = edges[j]
            else:
                break
        out[interval] = last_safe
    return out


# ---- 6. Verification against the published thresholds -----------------------

PUBLISHED = {"4": [30, 20, 20, 10, 10],
             "5": [20, 10, 6, 4, 2],
             "6": [4, 3, 2, 1, 1],
             "7": [2, 1, 1, 1, 1]}

print("Thresholds of no change of residual PPD (percentage of affected sites)")
print("=====================================================================\n")
all_match = True
for k in ["4", "5", "6", "7"]:
    means, _, edges = panel_table(k)
    reconstructed = thresholds(means, edges)
    values = [reconstructed[i] for i in INTERVALS]
    match = values == PUBLISHED[k]
    all_match &= match
    print(f"PPD >= {k} mm")
    print("  published    :", "".join(f"{v:>6}" for v in PUBLISHED[k]))
    print("  reconstructed:", "".join(f"{v:>6.0f}" for v in values))
    print("  match        :", match, "\n")
print("All twenty thresholds reproduce exactly.\n" if all_match
      else "Some thresholds did not reproduce.\n")


# ---- 7. Sensitivity: excluding visits with zero affected sites --------------
# Visits at which a patient had no residual pocket in a given category
# contribute a change of approximately zero and dominate the lowest bin, in
# particular for the deeper categories. Re-deriving the thresholds without them
# tests whether the algorithm depends on those zeros.

print("Sensitivity: thresholds with and without visits at 0% affected sites")
print("====================================================================\n")
for k in ["4", "5", "6", "7"]:
    share = (spt[f"pct{k}"] == 0).mean() * 100
    m1, _, e1 = panel_table(k, exclude_zero=False)
    m0, _, e0 = panel_table(k, exclude_zero=True)
    t1, t0 = thresholds(m1, e1), thresholds(m0, e0)
    fmt = lambda t: "".join(f"{t[i]:>6.0f}" if pd.notna(t[i]) else f"{'--':>6}"
                            for i in INTERVALS)
    print(f"PPD >= {k} mm   ({share:.1f}% of visits had no affected site)")
    print("  with zeros   :", fmt(t1))
    print("  without zeros:", fmt(t0), "\n")


# ---- 8. Cell means, for inspection or plotting ------------------------------

for k in ["4", "5", "6", "7"]:
    means, counts, _ = panel_table(k)
    print(f"\nMean change, PPD >= {k} mm (n in brackets)")
    display = means.copy().astype(object)
    for row in means.index:
        for col in INTERVALS:
            v, n = means.loc[row, col], counts.loc[row, col]
            display.loc[row, col] = "-" if pd.isna(v) else f"{v:+.2f} ({int(n)})"
    print(display.to_string())
    means.to_csv(f"figure2_panel_ppd{k}.csv")

print(f"\npandas {pd.__version__}, numpy {np.__version__}, "
      f"Python {sys.version.split()[0]}")