"""
Supportive Periodontal Therapy Clinical Examination Data
Residual-PPD-based algorithm for computing SPT intervals

Python equivalent of 01_compute_spt_algorithm.R. Both implementations were
written independently against the same specification and both reproduce every
derived variable in 03_supplementary_data.csv exactly.

The algorithm implements the empirically determined PPD stability thresholds
reported in Ramseier et al. 2019 (J Clin Periodontol 46:218-230), Figure 2.

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

import os
import sys

import numpy as np
import pandas as pd


# ---- 1. Read the data -------------------------------------------------------

def read_data(filename):
    if not os.path.exists(filename):
        sys.exit(
            f"'{filename}' not found in the working directory ({os.getcwd()}).\n"
            "  Run this script from the folder holding the data files, or pass "
            "that folder as the first argument."
        )
    return pd.read_csv(filename)


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

spt = read_data("02_supportive_periodontal_therapy_data.csv")
spt = spt.sort_values(["pat_id", "spt_id"]).reset_index(drop=True)


# ---- 2. Cumulative site counts at each SPT visit ----------------------------
# PPD categories are recorded as counts of sites at exactly 4, 5, 6 mm and at
# >= 7 mm. The algorithm operates on cumulative counts (>= 4, >= 5, >= 6 mm).
#
# min_count keeps a sum missing when any of its parts is missing, which is what
# we want for the visits flagged by `imputed`.

PPD = ["n_4mm_spt", "n_5mm_spt", "n_6mm_spt", "n_from7mm_spt"]

spt["n_from4mm_spt"] = spt[PPD].sum(axis=1, min_count=4)
spt["n_from5mm_spt"] = spt[PPD[1:]].sum(axis=1, min_count=3)
spt["n_from6mm_spt"] = spt[PPD[2:]].sum(axis=1, min_count=2)


# ---- 3. Site percentages ----------------------------------------------------
# PPD was recorded at six sites per tooth, so the denominator is n_teeth * 6.

n_sites = spt["n_teeth_spt"] * 6

for name, source in [
    ("percent_4mm_spt", "n_4mm_spt"),
    ("percent_5mm_spt", "n_5mm_spt"),
    ("percent_6mm_spt", "n_6mm_spt"),
    ("percent_from7mm_spt", "n_from7mm_spt"),
    ("percent_from4mm_spt", "n_from4mm_spt"),
    ("percent_from5mm_spt", "n_from5mm_spt"),
    ("percent_from6mm_spt", "n_from6mm_spt"),
]:
    spt[name] = spt[source] * 100 / n_sites


# ---- 4. PPD stability thresholds -------------------------------------------
# For each residual PPD category (>= 4, >= 5, >= 6 mm) and each candidate SPT
# interval (3, 4, 6, 12 months), a threshold percentage of affected sites was
# determined below which no increase of residual PPD was observed between two
# consecutive visits. A candidate interval is admissible only when the observed
# percentage lies at or below its threshold.
#
#   PPD >= 4 mm :  3 months if > 20% ; 4 and 6 months if <= 20% ; 12 months if <= 10%
#   PPD >= 5 mm :  3 months if > 10% ; 4 months if <= 10% ; 6 months if <= 6% ; 12 months if <= 2%
#   PPD >= 6 mm :  3 months if >  3% ; 4 months if <=  3% ; 6 months if <= 2% ; 12 months if <= 1%
#
# PPD >= 7 mm is deliberately NOT used: sites of >= 8 mm were not recorded at
# the MSDH, so >= 7 mm did not qualify as a threshold variable
# (Ramseier et al. 2019, section 3.2).
#
# NaN encodes "this interval is not admissible for this PPD category".

def threshold(condition, months):
    return np.where(condition.fillna(False), float(months), np.nan)


p4 = spt["percent_from4mm_spt"]
p5 = spt["percent_from5mm_spt"]
p6 = spt["percent_from6mm_spt"]

spt["from4mm_3months"] = threshold(p4 > 20, 3)
spt["from4mm_4months"] = threshold(p4 <= 20, 4)
spt["from4mm_6months"] = threshold(p4 <= 20, 6)
spt["from4mm_12months"] = threshold(p4 <= 10, 12)

spt["from5mm_3months"] = threshold(p5 > 10, 3)
spt["from5mm_4months"] = threshold(p5 <= 10, 4)
spt["from5mm_6months"] = threshold(p5 <= 6, 6)
spt["from5mm_12months"] = threshold(p5 <= 2, 12)

spt["from6mm_3months"] = threshold(p6 > 3, 3)
spt["from6mm_4months"] = threshold(p6 <= 3, 4)
spt["from6mm_6months"] = threshold(p6 <= 2, 6)
spt["from6mm_12months"] = threshold(p6 <= 1, 12)


# ---- 5. Longest admissible interval per PPD category ------------------------
# NaN here has two distinct meanings and they must not be conflated: an interval
# that is not admissible for an observed percentage, and a percentage that was
# never measured. The first counts as 0 and competes in the maximum; the second
# makes the whole result unknown. Visits flagged by `imputed` fall in the second
# case, and every value derived from them stays missing.

def longest_admissible(percentage, *columns):
    stacked = np.vstack([np.nan_to_num(spt[c].to_numpy(dtype=float), nan=0.0)
                         for c in columns])
    result = stacked.max(axis=0)
    result[percentage.isna().to_numpy()] = np.nan
    return result


spt["max_from4mm"] = longest_admissible(
    p4, "from4mm_3months", "from4mm_4months", "from4mm_6months", "from4mm_12months")
spt["max_from5mm"] = longest_admissible(
    p5, "from5mm_3months", "from5mm_4months", "from5mm_6months", "from5mm_12months")
spt["max_from6mm"] = longest_admissible(
    p6, "from6mm_3months", "from6mm_4months", "from6mm_6months", "from6mm_12months")


# ---- 6. Computed SPT interval ----------------------------------------------
# The binding constraint is the shortest of the three category-wise intervals.
# A category contributing 0 (no admissible interval) is ignored.

def ignore_zero(x):
    return np.where(x != 0, x, 999.0)


spt["algorithm_based_interval_months"] = np.minimum(
    np.minimum(ignore_zero(spt["max_from4mm"]), ignore_zero(spt["max_from5mm"])),
    ignore_zero(spt["max_from6mm"]),
)


# ---- 7. Adherence to the computed interval ----------------------------------
# true_interval_days is 0 at each patient's first visit (no preceding visit
# exists). Months are obtained by rounding days / 30, as in the original
# analysis. R rounds halves to even, and numpy does the same, so the two
# implementations agree.

spt["true_interval_months"] = np.round(spt["true_interval_days"] / 30)

interval_dif = pd.Series(
    np.where(spt["true_interval_days"] == 0, 0.0,
             spt["true_interval_months"] - spt["algorithm_based_interval_months"]),
    index=spt.index,
)
interval_dif[spt["algorithm_based_interval_months"].isna()] = np.nan

spt["interval_shorter"] = np.where(interval_dif < 0, 1.0, np.nan)
spt["interval_exact"] = np.where(interval_dif == 0, 1.0, np.nan)
spt["interval_longer"] = np.where(interval_dif > 0, 1.0, np.nan)

# NOTE: interval_shorter / interval_exact / interval_longer compare the actual
# interval against the ALGORITHM-BASED interval, not against the interval
# assigned by the dental hygienist (assigned_spt_interval_months).


# ---- 8. Runs of consecutive visits in the same adherence category -----------

def run_length(patient_ids, flags):
    flags = np.nan_to_num(np.asarray(flags, dtype=float), nan=0.0).astype(int)
    patient_ids = np.asarray(patient_ids)
    out = np.zeros(len(flags), dtype=int)
    run = 0
    for i in range(len(flags)):
        if i > 0 and patient_ids[i] != patient_ids[i - 1]:
            run = 0
        run = run + 1 if flags[i] == 1 else 0
        out[i] = run
    return out


for target, source in [("consecutive_shorter", "interval_shorter"),
                       ("consecutive_exact", "interval_exact"),
                       ("consecutive_longer", "interval_longer")]:
    spt[target] = run_length(spt["pat_id"], spt[source])


# ---- 9. Verification against the published supplementary file ---------------

CHECK = [
    "n_from4mm_spt", "n_from5mm_spt", "n_from6mm_spt",
    "percent_4mm_spt", "percent_5mm_spt", "percent_6mm_spt",
    "percent_from7mm_spt", "percent_from4mm_spt",
    "percent_from5mm_spt", "percent_from6mm_spt",
    "from4mm_3months", "from4mm_4months", "from4mm_6months", "from4mm_12months",
    "from5mm_3months", "from5mm_4months", "from5mm_6months", "from5mm_12months",
    "from6mm_3months", "from6mm_4months", "from6mm_6months", "from6mm_12months",
    "max_from4mm", "max_from5mm", "max_from6mm",
    "algorithm_based_interval_months",
    "interval_shorter", "interval_exact", "interval_longer",
    "consecutive_shorter", "consecutive_exact", "consecutive_longer",
]

if os.path.exists("03_supplementary_data.csv"):
    reference = pd.read_csv("03_supplementary_data.csv")
    reference = reference.sort_values(["pat_id", "spt_id"]).reset_index(drop=True)

    print("\nVerification against 03_supplementary_data.csv")
    print("---------------------------------------------")
    failures = 0
    for variable in CHECK:
        a = spt[variable].astype(float)
        b = reference[variable].astype(float)
        identical = ((a.isna() & b.isna())
                     | (a.notna() & b.notna() & ((a - b).abs() < 1e-9)))
        n_ok = int(identical.sum())
        if n_ok < len(reference):
            failures += 1
        print(f"  {variable:<32} {n_ok:>6} / {len(reference)} identical"
              f"{'   <-- MISMATCH' if n_ok < len(reference) else ''}")

    print()
    if failures == 0:
        print("All derived variables reproduce exactly.")
    else:
        print(f"{failures} variable(s) did not reproduce.")

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