"""Plot the Week 6 denominator and competing-exit teaching examples.

Run the lab first, then from the course root:
  conda run -n research python assets/scripts/build-week6-careers.py

Exact generated JSON/CSV inputs are preserved in assets/data. The 100-worker
examples and the separate 18-worker lottery are all invented, not article data.
"""
from pathlib import Path
import csv
import json
import math
import os
import sys

BOOK = Path(__file__).resolve().parents[2]
REPO = BOOK.parents[1]
os.environ.setdefault("MPLCONFIGDIR", str(BOOK / ".mpl-cache"))
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.ticker import PercentFormatter

sys.path.insert(0, str(REPO))
from shared.figure_style import COLORS, apply_style, save_figure, style_axis

MODE = "SYNTHETIC_TEACHING_CONSTRUCTION_NOT_PUBLISHED_REPLICATION"
LAB = BOOK / "labs/06-training-internal-mobility-promotions-and-career-ladders"


def preserve(source_name, target_name):
    source = LAB / source_name
    target = BOOK / "assets/data" / target_name
    target.parent.mkdir(parents=True, exist_ok=True)
    target.write_bytes(source.read_bytes())
    return target


def save(fig, name):
    path = BOOK / "assets/figures" / name
    save_figure(fig, path.with_suffix(".png"))
    fig.savefig(path.with_suffix(".pdf"), bbox_inches="tight", facecolor="white",
                metadata={"CreationDate": None, "ModDate": None})
    plt.close(fig)


def main():
    diagnostic_path = preserve("output/transfer/diagnostic_examples.json", "06-career-diagnostics.json")
    risks_path = preserve("output/reproduced/competing_risks.csv", "06-competing-risks.csv")
    examples = json.loads(diagnostic_path.read_text())
    assert examples["provenance"] == MODE and examples["empirical_replication"] is False
    d = examples["denominator_example"]
    n = d["eligible_each"]
    cohort = [d["treated_promotions"] / n, d["control_promotions"] / n]
    applicant = [d["treated_promotions"] / d["treated_applications"],
                 d["control_promotions"] / d["control_applications"]]
    assert all(math.isclose(x, y, abs_tol=1e-10) for x, y in zip(cohort + applicant, [.10, .08, .25, .40]))
    assert cohort[0] > cohort[1] and applicant[0] < applicant[1]
    with risks_path.open(newline="") as stream:
        rows = list(csv.DictReader(stream))
    assert len(rows) == 4 and [int(r["month"]) for r in rows] == [1, 2, 3, 4]
    assert all(r["provenance"] == MODE for r in rows)
    times, cif, promotion_only = [0], [0.0], [0.0]
    survival, fp, fx, km = 1.0, 0.0, 0.0, 1.0
    risk = int(rows[0]["risk_start"])
    initial = risk
    for row in rows:
        assert int(row["risk_start"]) == risk and risk > 0
        np, nx = int(row["promotions"]), int(row["exits"])
        assert 0 <= np + nx <= risk
        hp, hx = np / risk, nx / risk
        fp += survival * hp
        fx += survival * hx
        survival *= 1 - hp - hx
        km *= 1 - hp  # Deliberately censors exits; a different estimand.
        for key, expected in [("promotion_hazard", hp), ("exit_hazard", hx),
                              ("survival", survival), ("promotion_cumulative_incidence", fp),
                              ("exit_cumulative_incidence", fx)]:
            assert math.isclose(float(row[key]), expected, abs_tol=1e-8)
        assert math.isclose(survival + fp + fx, 1, abs_tol=1e-10)
        risk -= np + nx
        times.append(int(row["month"]))
        cif.append(fp)
        promotion_only.append(1 - km)
    assert initial == 18 and sum(int(r["promotions"]) for r in rows) == 3
    assert math.isclose(cif[-1], 3 / 18, abs_tol=1e-10)
    c = examples["competing_risks_example"]
    cf = c["promotion_month2"] / c["initial"]
    net = c["promotion_month2"] / (c["initial"] - c["exit_month1"])
    assert math.isclose(cf, c["promotion_CIF"]) and math.isclose(net, c["promotion_only_KM_complement"])
    assert math.isclose(cf, .30) and math.isclose(net, .50)

    apply_style()
    fig, axes = plt.subplots(1, 2, figsize=(10.5, 5.6), sharey=True)
    for ax, values, title, counts in zip(axes, [cohort, applicant],
            ["Advancement per eligible worker", "Advancement per applicant"],
            [["10 / 100", "8 / 100"], ["10 / 40", "8 / 20"]]):
        bars = ax.bar([0, 1], values, color=[COLORS["navy"], COLORS["teal"]], width=.52)
        ax.bar_label(bars, labels=[f"{v:.0%}" for v in values], padding=6, fontsize=12)
        ax.set(title=title, xticks=[0, 1], xticklabels=["Reform group\n" + counts[0], "Comparison group\n" + counts[1]], ylim=(0, .48))
        ax.yaxis.set_major_formatter(PercentFormatter(1))
        ax.tick_params(axis="x", labelsize=10)
        style_axis(ax)
    axes[0].set_ylabel("Promotion rate")
    fig.text(.5, .112, "Original diagnostic: 100 eligible per group; 40 vs 20 apply; 10 vs 8 advance.", ha="center", fontsize=9.5, color=COLORS["muted"])
    fig.text(.5, .068, "The applicant-conditioned comparison describes pools partly created by the reform.", ha="center", fontsize=9.5, color=COLORS["muted"])
    fig.text(.5, .022, MODE, ha="center", fontsize=8, color=COLORS["muted"])
    fig.tight_layout(rect=(0, .16, 1, 1))
    save(fig, "06-career-denominators")

    fig, axes = plt.subplots(1, 2, figsize=(10.5, 5.8))
    panels = [([0, 1, 2], [0, 0, cf], [0, 0, net], "Original 100-worker example", 2.7, .59),
              (times, cif, promotion_only, "Generated 18-worker cohort", 4.7, .245)]
    for ax, (x, fp_values, net_values, title, xmax, ymax) in zip(axes, panels):
        for values, color, linestyle, name in [(fp_values, COLORS["navy"], "-", "First-promotion incidence"),
                                               (net_values, COLORS["rust"], "--", "Promotion-only calculation")]:
            ax.step(x + [xmax], values + [values[-1]], where="post", color=color, linestyle=linestyle,
                    linewidth=2.2, label=name)
            ax.text(x[-1]+.1, values[-1]+.012*ymax, f"{100*values[-1]:.2f}%".replace(".00%", "%"),
                    color=color, fontsize=10, va="bottom")
        ax.set(title=title, xlabel="Period", ylabel="Cumulative probability", xticks=x, xlim=(0, xmax), ylim=(-.01, ymax))
        ax.yaxis.set_major_formatter(PercentFormatter(1))
        ax.legend(loc="upper left", fontsize=8.5, frameon=False)
        style_axis(ax)
    fig.text(.5, .12, "Exits are competing events. Promotion-only calculation = 1 − Kaplan–Meier survival censoring exits.", ha="center", fontsize=9, color=COLORS["muted"])
    fig.text(.5, .076, "Separate invented cohorts; no ordinary loss to follow-up. Neither curve makes assignment causal.", ha="center", fontsize=9, color=COLORS["muted"])
    fig.text(.5, .022, MODE, ha="center", fontsize=8, color=COLORS["muted"])
    fig.tight_layout(rect=(0, .17, 1, 1))
    save(fig, "06-career-competing-risks")
    print("PASS: denominator and competing-risk arithmetic checked; exact JSON/CSV inputs and two PNG/PDF figures saved.")


if __name__ == "__main__":
    main()