"""Plot three hypothetical multitask contracts from the generated Week 4 CSV.

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

The script preserves the exact lab CSV under assets/data. All observations,
reward slopes, cost parameters and gross-value weights are invented.
"""

from pathlib import Path
import csv
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

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

PROVENANCE = "SYNTHETIC_TEACHING_NOT_PUBLISHED_ESTIMATES"
LAB = BOOK / "labs/04-pay-incentives-multitasking-and-relational-contracts"


def main() -> None:
    source = LAB / "output/transfer/multitask.csv"
    with source.open(newline="") as stream:
        rows = list(csv.DictReader(stream))
    assert len(rows) == 3 and all(r["provenance"] == PROVENANCE for r in rows)
    expected = [(0.8, 0.8, 0.5, 0.5, 1.5),
                (1.2, 0.8, 1.125, 0.125, 1.375),
                (1.0, 1.0, 0.625, 0.625, 1.875)]
    by_slopes = {(float(r["b1"]), float(r["b2"])): r for r in rows}
    assert set(by_slopes) == {(v[0], v[1]) for v in expected}
    data = []
    for b1, b2, e1, e2, value in expected:
        row = by_slopes[b1, b2]
        actual = [float(row[k]) for k in ("task1_effort", "task2_effort", "gross_value")]
        for observed, target in zip(actual, (e1, e2, value)):
            assert math.isclose(observed, target, rel_tol=0, abs_tol=1e-8)
        assert math.isclose(e1 + 0.6 * e2, b1) and math.isclose(e2 + 0.6 * e1, b2)
        assert e1 >= 0 and e2 >= 0
        assert math.isclose(float(row["variable_pay"]), b1 * e1 + b2 * e2, abs_tol=1e-8)
        assert math.isclose(float(row["effort_cost"]), 0.5*(e1**2 + e2**2 + 1.2*e1*e2), abs_tol=1e-8)
        data.append(actual)
    apply_style()
    fig, axes = plt.subplots(1, 2, figsize=(10.5, 5.1), gridspec_kw={"width_ratios": [1.4, 1]})
    labels = ["Baseline\n(0.8, 0.8)", "Quantity only\n(1.2, 0.8)", "Balanced\n(1.0, 1.0)"]
    x = list(range(3))
    for offset, column, color, name in [(-0.18, 0, COLORS["navy"], "Task 1: quantity"),
                                       (0.18, 1, COLORS["teal"], "Task 2: quality")]:
        bars = axes[0].bar([i + offset for i in x], [r[column] for r in data], width=0.34,
                           color=color, label=name)
        axes[0].bar_label(bars, labels=[f"{r[column]:.3f}".rstrip("0").rstrip(".") for r in data],
                          padding=3, fontsize=9)
    axes[0].set(title="Effort moves toward the rewarded task", ylabel="Effort (hypothetical units)",
                xticks=x, xticklabels=labels, ylim=(0, 1.45))
    axes[0].legend(loc="upper right", fontsize=9)
    bars = axes[1].bar(x, [r[2] for r in data], color=[COLORS["navy"], COLORS["rust"], COLORS["teal"]], width=0.58)
    axes[1].bar_label(bars, labels=[f"{r[2]:.3f}".rstrip("0").rstrip(".") for r in data], padding=4, fontsize=10)
    axes[1].set(title="Gross value can fall", ylabel="Gross value: task 1 + 2 × task 2",
                xticks=x, xticklabels=labels, ylim=(0, 2.22))
    for ax in axes:
        ax.tick_params(axis="x", labelsize=9)
        style_axis(ax)
    fig.text(0.5, 0.105, "Two-task quadratic costs; κ = 0.6; interior nonnegative effort; labels show reward slopes.",
             ha="center", fontsize=9, color=COLORS["muted"])
    fig.text(0.5, 0.062, "Gross value excludes compensation, effort costs and measurement costs; it is not profit or worker welfare.",
             ha="center", fontsize=8.5, color=COLORS["muted"])
    fig.text(0.5, 0.018, PROVENANCE, ha="center", fontsize=8, color=COLORS["muted"])
    fig.tight_layout(rect=(0, 0.16, 1, 1))
    target = BOOK / "assets/data/04-multitask.csv"
    target.parent.mkdir(parents=True, exist_ok=True)
    target.write_bytes(source.read_bytes())
    output = BOOK / "assets/figures/04-multitask-contracts"
    save_figure(fig, output.with_suffix(".png"))
    fig.savefig(output.with_suffix(".pdf"), bbox_inches="tight", facecolor="white",
                metadata={"CreationDate": None, "ModDate": None})
    plt.close(fig)
    print("PASS: three generated contracts verified; exact CSV copy and PNG/PDF saved.")


if __name__ == "__main__":
    main()