#!/usr/bin/env python3
"""Controlled multitask training run for the meeting-minutes classifier rework.

This script is deliberately separate from the earlier cross-validation trainer.
It honours per-head masks, recommended splits and sparse labels from the
reformatted 667-row master CSV.
"""
from __future__ import annotations

import argparse
import json
import math
import shutil
from collections import Counter
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import joblib
import numpy as np
import pandas as pd
from sentence_transformers import SentenceTransformer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import (
    classification_report,
    confusion_matrix,
    f1_score,
    precision_recall_fscore_support,
)
from sklearn.multiclass import OneVsRestClassifier
from sklearn.preprocessing import LabelEncoder, MultiLabelBinarizer


EXPECTED = {
    "master_rows": 667,
    "evidence_rows": 547,
    "minutes_quality_rows": 120,
    "quality_ready": 63,
    "quality_intervention": 57,
    "split_train": 449,
    "split_validation": 90,
    "split_test": 128,
}


def split_labels(value: Any) -> list[str]:
    if pd.isna(value) or str(value).strip() == "":
        return []
    return [part.strip() for part in str(value).split("|") if part.strip()]


def now_slug() -> str:
    return datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")


def as_bool_mask(series: pd.Series) -> pd.Series:
    return series.fillna(0).astype(int).eq(1)


def fail(message: str) -> None:
    raise SystemExit(f"VALIDATION FAILED: {message}")


def validate_master(df: pd.DataFrame, expected: dict[str, int] | None = None) -> dict[str, Any]:
    expected = expected or EXPECTED
    checks: dict[str, Any] = {}
    checks["master_rows"] = int(len(df))
    if len(df) != expected["master_rows"]:
        fail(f"expected {expected['master_rows']} master rows, found {len(df)}")

    if not df["row_id"].is_unique:
        dupes = df.loc[df["row_id"].duplicated(), "row_id"].tolist()[:10]
        fail(f"row_id values are not unique: {dupes}")
    checks["unique_row_id"] = True

    evidence_rows = int(as_bool_mask(df["train_evidence_head"]).sum())
    checks["evidence_rows"] = evidence_rows
    if evidence_rows != expected["evidence_rows"]:
        fail(f"expected {expected['evidence_rows']} evidence-family rows, found {evidence_rows}")

    quality_rows = int(as_bool_mask(df["train_quality_gate"]).sum())
    checks["minutes_quality_rows"] = quality_rows
    if quality_rows != expected["minutes_quality_rows"]:
        fail(f"expected {expected['minutes_quality_rows']} minutes-quality rows, found {quality_rows}")

    quality_counts = df.loc[as_bool_mask(df["train_quality_gate"]), "quality_gate"].value_counts().to_dict()
    checks["quality_gate_counts"] = {str(k): int(v) for k, v in quality_counts.items()}
    if quality_counts.get("ready", 0) != expected["quality_ready"]:
        fail(f"expected {expected['quality_ready']} ready quality rows, found {quality_counts.get('ready', 0)}")
    if quality_counts.get("intervention", 0) != expected["quality_intervention"]:
        fail(f"expected {expected['quality_intervention']} intervention quality rows, found {quality_counts.get('intervention', 0)}")

    split_counts = df["recommended_split"].value_counts().to_dict()
    checks["recommended_split_counts"] = {str(k): int(v) for k, v in split_counts.items()}
    for split in ["train", "validation", "test"]:
        expected_split = expected[f"split_{split}"]
        if split_counts.get(split, 0) != expected_split:
            fail(f"expected {expected_split} {split} rows, found {split_counts.get(split, 0)}")

    group_splits = df.groupby("group_id")["recommended_split"].nunique(dropna=False)
    leaking_groups = group_splits[group_splits > 1]
    checks["group_id_single_split"] = leaking_groups.empty
    if not leaking_groups.empty:
        fail(f"group_id crosses splits: {leaking_groups.index.tolist()[:10]}")

    for mask_col, target_col in [
        ("train_evidence_head", "evidence_type"),
        ("train_action_head", "action_state"),
        ("train_signal_head", "signals"),
        ("train_quality_gate", "quality_gate"),
        ("train_quality_defect_head", "quality_defects"),
    ]:
        enabled = as_bool_mask(df[mask_col])
        if target_col == "signals":
            # Empty signal sets are valid for explicitly enabled rows.
            continue
        missing = df.loc[enabled, target_col].fillna("").astype(str).str.strip().eq("")
        if missing.any():
            fail(f"{target_col} has empty values where {mask_col}=1: {df.loc[enabled & missing, 'row_id'].head(10).tolist()}")
    checks["enabled_heads_have_targets"] = True

    disabled_excluded = {}
    for mask_col in [
        "train_evidence_head",
        "train_action_head",
        "train_signal_head",
        "train_quality_gate",
        "train_quality_defect_head",
        "train_disposition_head",
    ]:
        disabled_excluded[mask_col] = int((~as_bool_mask(df[mask_col])).sum())
    checks["disabled_rows_available_for_exclusion"] = disabled_excluded

    duplicates = df[df["model_text"].fillna("").astype(str).str.strip() != ""].copy()
    cross_split_dupes = (
        duplicates.groupby("model_text")["recommended_split"]
        .nunique(dropna=False)
        .loc[lambda s: s > 1]
    )
    checks["no_exact_duplicate_model_text_crosses_splits"] = cross_split_dupes.empty
    if not cross_split_dupes.empty:
        fail(f"exact duplicate model_text crosses splits; first example: {cross_split_dupes.index[0][:120]}")

    return checks


def encode_texts(embedder: SentenceTransformer, texts: list[str]) -> np.ndarray:
    return embedder.encode(
        texts,
        normalize_embeddings=True,
        convert_to_numpy=True,
        batch_size=32,
        show_progress_bar=True,
    )


def split_counts(df: pd.DataFrame, label_col: str) -> dict[str, dict[str, int]]:
    out: dict[str, dict[str, int]] = {}
    for split in ["train", "validation", "test"]:
        out[split] = {
            str(k): int(v)
            for k, v in df.loc[df["recommended_split"].eq(split), label_col].value_counts().sort_index().items()
        }
    return out


def abstain_mask(probs: np.ndarray, min_confidence: float, min_margin: float) -> np.ndarray:
    sorted_probs = np.sort(probs, axis=1)
    top = sorted_probs[:, -1]
    second = sorted_probs[:, -2] if probs.shape[1] > 1 else np.zeros_like(top)
    return (top < min_confidence) | ((top - second) < min_margin)


def tune_abstention(y_true: np.ndarray, probs: np.ndarray, pred: np.ndarray) -> dict[str, float]:
    best = {"min_confidence": 0.0, "min_margin": 0.0, "macro_f1_non_abstained": 0.0, "abstention_rate": 1.0}
    for conf in np.linspace(0.0, 0.8, 17):
        for margin in np.linspace(0.0, 0.3, 7):
            abstain = abstain_mask(probs, float(conf), float(margin))
            keep = ~abstain
            if keep.sum() == 0:
                continue
            score = f1_score(y_true[keep], pred[keep], average="macro", zero_division=0)
            penalty = 0.05 * abstain.mean()
            ranked = score - penalty
            current = best["macro_f1_non_abstained"] - 0.05 * best["abstention_rate"]
            if ranked > current:
                best = {
                    "min_confidence": float(conf),
                    "min_margin": float(margin),
                    "macro_f1_non_abstained": float(score),
                    "abstention_rate": float(abstain.mean()),
                }
    return best


def single_metrics(y_true: np.ndarray, pred: np.ndarray, labels: list[str]) -> dict[str, Any]:
    label_ids = list(range(len(labels)))
    report = classification_report(
        y_true,
        pred,
        labels=label_ids,
        target_names=labels,
        output_dict=True,
        zero_division=0,
    )
    return {
        "macro_f1": float(f1_score(y_true, pred, average="macro", zero_division=0)),
        "micro_f1": float(f1_score(y_true, pred, average="micro", zero_division=0)),
        "per_class": report,
        "confusion_matrix": confusion_matrix(y_true, pred, labels=label_ids).astype(int).tolist(),
        "labels": labels,
    }


def evaluate_single_split(clf, le: LabelEncoder, X: np.ndarray, y: np.ndarray, abstention: dict[str, float]) -> dict[str, Any]:
    probs = clf.predict_proba(X)
    pred = probs.argmax(axis=1)
    metrics = single_metrics(y, pred, list(le.classes_))
    abstain = abstain_mask(probs, abstention["min_confidence"], abstention["min_margin"])
    keep = ~abstain
    metrics["abstention"] = {
        "rate": float(abstain.mean()) if len(abstain) else 0.0,
        "count": int(abstain.sum()),
        "non_abstained_count": int(keep.sum()),
    }
    if keep.sum():
        metrics["non_abstained"] = single_metrics(y[keep], pred[keep], list(le.classes_))
    else:
        metrics["non_abstained"] = None
    return metrics


def train_single_head(
    name: str,
    df: pd.DataFrame,
    mask_col: str,
    target_col: str,
    X_all: np.ndarray,
    out: Path,
    c_value: float = 1.0,
) -> dict[str, Any]:
    head_df = df.loc[as_bool_mask(df[mask_col])].copy()
    idx = head_df.index.to_numpy()
    le = LabelEncoder()
    y_all = le.fit_transform(head_df[target_col].astype(str))
    X = X_all[idx]
    train = head_df["recommended_split"].eq("train").to_numpy()
    validation = head_df["recommended_split"].eq("validation").to_numpy()
    test = head_df["recommended_split"].eq("test").to_numpy()

    clf = LogisticRegression(max_iter=3000, class_weight="balanced", C=c_value)
    clf.fit(X[train], y_all[train])

    val_probs = clf.predict_proba(X[validation])
    val_pred = val_probs.argmax(axis=1)
    abstention = tune_abstention(y_all[validation], val_probs, val_pred)

    split_metrics = {
        "train": evaluate_single_split(clf, le, X[train], y_all[train], abstention),
        "validation": evaluate_single_split(clf, le, X[validation], y_all[validation], abstention),
        "test": evaluate_single_split(clf, le, X[test], y_all[test], abstention),
    }

    predictions = []
    probs_all = clf.predict_proba(X)
    pred_all = probs_all.argmax(axis=1)
    abstain_all = abstain_mask(probs_all, abstention["min_confidence"], abstention["min_margin"])
    for pos, row in enumerate(head_df.itertuples()):
        predictions.append(
            {
                "row_id": row.row_id,
                "split": row.recommended_split,
                "true_label": getattr(row, target_col),
                "predicted_label": str(le.classes_[pred_all[pos]]),
                "abstained": bool(abstain_all[pos]),
                **{f"prob_{cls}": float(probs_all[pos, i]) for i, cls in enumerate(le.classes_)},
            }
        )
    pd.DataFrame(predictions).to_csv(out / f"{name}_predictions.csv", index=False)
    pd.DataFrame({"row_id": head_df["row_id"], "split": head_df["recommended_split"]}).to_csv(out / f"{name}_row_ids.csv", index=False)

    joblib.dump({"label_encoder": le, "classifier": clf, "abstention": abstention}, out / f"{name}_classifier.joblib")
    return {
        "row_count": int(len(head_df)),
        "train_count": int(train.sum()),
        "validation_count": int(validation.sum()),
        "test_count": int(test.sum()),
        "class_distribution": split_counts(head_df, target_col),
        "abstention_policy": abstention,
        "metrics": split_metrics,
    }


def tune_multilabel_thresholds(y_true: np.ndarray, probs: np.ndarray) -> list[float]:
    thresholds = []
    for i in range(y_true.shape[1]):
        best_t = 0.5
        best_f1 = -1.0
        for threshold in np.linspace(0.1, 0.9, 17):
            pred = (probs[:, i] >= threshold).astype(int)
            score = f1_score(y_true[:, i], pred, zero_division=0)
            if score > best_f1:
                best_f1 = score
                best_t = float(threshold)
        thresholds.append(best_t)
    return thresholds


def multilabel_metrics(y_true: np.ndarray, probs: np.ndarray, thresholds: list[float], labels: list[str]) -> dict[str, Any]:
    pred = (probs >= np.asarray(thresholds).reshape(1, -1)).astype(int)
    pr, rc, f1, sup = precision_recall_fscore_support(y_true, pred, average=None, zero_division=0)
    per_label = {}
    for i, label in enumerate(labels):
        yt = y_true[:, i]
        yp = pred[:, i]
        tp = int(((yt == 1) & (yp == 1)).sum())
        fp = int(((yt == 0) & (yp == 1)).sum())
        fn = int(((yt == 1) & (yp == 0)).sum())
        tn = int(((yt == 0) & (yp == 0)).sum())
        per_label[label] = {
            "support": int(sup[i]),
            "precision": float(pr[i]),
            "recall": float(rc[i]),
            "f1": float(f1[i]),
            "tp": tp,
            "fp": fp,
            "fn": fn,
            "tn": tn,
            "threshold": float(thresholds[i]),
        }
    margin_to_threshold = np.max(probs - np.asarray(thresholds).reshape(1, -1), axis=1)
    abstain = margin_to_threshold < 0.05
    keep = ~abstain
    result: dict[str, Any] = {
        "macro_f1": float(f1_score(y_true, pred, average="macro", zero_division=0)),
        "micro_f1": float(f1_score(y_true, pred, average="micro", zero_division=0)),
        "per_label": per_label,
        "abstention": {
            "rate": float(abstain.mean()) if len(abstain) else 0.0,
            "count": int(abstain.sum()),
            "non_abstained_count": int(keep.sum()),
        },
    }
    if keep.sum():
        result["non_abstained"] = {
            "macro_f1": float(f1_score(y_true[keep], pred[keep], average="macro", zero_division=0)),
            "micro_f1": float(f1_score(y_true[keep], pred[keep], average="micro", zero_division=0)),
        }
    else:
        result["non_abstained"] = None
    return result


def train_multilabel_head(
    name: str,
    df: pd.DataFrame,
    mask_col: str,
    target_col: str,
    X_all: np.ndarray,
    out: Path,
    fixed_labels: list[str] | None = None,
    c_value: float = 1.0,
) -> dict[str, Any]:
    head_df = df.loc[as_bool_mask(df[mask_col])].copy()
    idx = head_df.index.to_numpy()
    labels_per_row = head_df[target_col].apply(split_labels).tolist()
    if fixed_labels is None:
        labels = sorted({label for row in labels_per_row for label in row})
    else:
        labels = fixed_labels
    mlb = MultiLabelBinarizer(classes=labels)
    y_all = mlb.fit_transform(labels_per_row)
    X = X_all[idx]
    train = head_df["recommended_split"].eq("train").to_numpy()
    validation = head_df["recommended_split"].eq("validation").to_numpy()
    test = head_df["recommended_split"].eq("test").to_numpy()

    train_support = y_all[train].sum(axis=0)
    learnable = np.where((train_support > 0) & (train_support < train.sum()))[0]
    if len(learnable) < len(labels):
        labels = [labels[i] for i in learnable]
        mlb = MultiLabelBinarizer(classes=labels)
        y_all = mlb.fit_transform([[label for label in row if label in labels] for row in labels_per_row])

    clf = OneVsRestClassifier(LogisticRegression(max_iter=3000, class_weight="balanced", C=c_value))
    clf.fit(X[train], y_all[train])
    val_probs = np.asarray(clf.predict_proba(X[validation]))
    thresholds = tune_multilabel_thresholds(y_all[validation], val_probs)

    split_metrics = {}
    for split_name, mask in [("train", train), ("validation", validation), ("test", test)]:
        probs = np.asarray(clf.predict_proba(X[mask]))
        split_metrics[split_name] = multilabel_metrics(y_all[mask], probs, thresholds, labels)

    probs_all = np.asarray(clf.predict_proba(X))
    pred_all = (probs_all >= np.asarray(thresholds).reshape(1, -1)).astype(int)
    predictions = []
    for pos, row in enumerate(head_df.itertuples()):
        predicted = [labels[i] for i, val in enumerate(pred_all[pos]) if val]
        truth = [label for label in split_labels(getattr(row, target_col)) if label in labels]
        item = {
            "row_id": row.row_id,
            "split": row.recommended_split,
            "true_labels": "|".join(truth),
            "predicted_labels": "|".join(predicted),
        }
        item.update({f"prob_{label}": float(probs_all[pos, i]) for i, label in enumerate(labels)})
        predictions.append(item)
    pd.DataFrame(predictions).to_csv(out / f"{name}_predictions.csv", index=False)
    pd.DataFrame({"row_id": head_df["row_id"], "split": head_df["recommended_split"]}).to_csv(out / f"{name}_row_ids.csv", index=False)
    pd.DataFrame({"label": labels, "threshold": thresholds}).to_csv(out / f"{name}_thresholds.csv", index=False)

    joblib.dump({"binarizer": mlb, "classifier": clf, "thresholds": thresholds}, out / f"{name}_classifier.joblib")
    return {
        "row_count": int(len(head_df)),
        "train_count": int(train.sum()),
        "validation_count": int(validation.sum()),
        "test_count": int(test.sum()),
        "label_distribution": {
            split_name: Counter(
                label
                for row in head_df.loc[head_df["recommended_split"].eq(split_name), target_col].apply(split_labels)
                for label in row
                if label in labels
            )
            for split_name in ["train", "validation", "test"]
        },
        "labels": labels,
        "thresholds": {label: float(thresholds[i]) for i, label in enumerate(labels)},
        "metrics": split_metrics,
    }


def reliability_summary(probs: np.ndarray, correct: np.ndarray) -> list[dict[str, Any]]:
    confidence = probs.max(axis=1)
    rows = []
    for lo, hi in [(0.0, 0.5), (0.5, 0.6), (0.6, 0.7), (0.7, 0.8), (0.8, 0.9), (0.9, 1.01)]:
        mask = (confidence >= lo) & (confidence < hi)
        rows.append(
            {
                "confidence_bin": f"{lo:.1f}-{min(hi, 1.0):.1f}",
                "count": int(mask.sum()),
                "accuracy": float(correct[mask].mean()) if mask.any() else None,
                "mean_confidence": float(confidence[mask].mean()) if mask.any() else None,
            }
        )
    return rows


def add_reliability(head_result: dict[str, Any], pred_csv: Path) -> None:
    pred_df = pd.read_csv(pred_csv)
    for split in ["train", "validation", "test"]:
        sdf = pred_df[pred_df["split"].eq(split)].copy()
        prob_cols = [c for c in sdf.columns if c.startswith("prob_")]
        if sdf.empty or not prob_cols:
            continue
        probs = sdf[prob_cols].to_numpy(float)
        correct = sdf["true_label"].astype(str).eq(sdf["predicted_label"].astype(str)).to_numpy()
        head_result["metrics"][split]["reliability"] = reliability_summary(probs, correct)


def quality_breakdowns(df: pd.DataFrame, predictions_path: Path) -> dict[str, Any]:
    pred = pd.read_csv(predictions_path)
    merged = pred.merge(df[["row_id", "input_stage", "source_origin", "dataset_family"]], on="row_id", how="left")
    out: dict[str, Any] = {}
    for field, groups in {
        "input_stage": ["raw_transcript", "generated_minutes", "human_minutes"],
        "source_family": ["synthetic_seed", "transcript_derived", "human_minutes_derived"],
    }.items():
        out[field] = {}
        if field == "source_family":
            def family(value: Any) -> str:
                text = str(value)
                if "human_minutes" in text:
                    return "human_minutes_derived"
                if "transcript" in text or "jacqui3" in text:
                    return "transcript_derived"
                if "synthetic" in text:
                    return "synthetic_seed"
                return "other"
            merged["_source_family"] = merged["source_origin"].apply(family)
            series = merged["_source_family"]
        else:
            series = merged[field]
        for group in groups:
            sdf = merged[series.eq(group)]
            if sdf.empty:
                continue
            labels = sorted(set(sdf["true_label"].astype(str)) | set(sdf["predicted_label"].astype(str)))
            le = LabelEncoder().fit(labels)
            y_true = le.transform(sdf["true_label"].astype(str))
            y_pred = le.transform(sdf["predicted_label"].astype(str))
            out[field][group] = {
                "rows": int(len(sdf)),
                "macro_f1": float(f1_score(y_true, y_pred, average="macro", zero_division=0)),
                "micro_f1": float(f1_score(y_true, y_pred, average="micro", zero_division=0)),
            }
    return out


def create_faithfulness_review(pair_path: Path, out: Path, embedder: SentenceTransformer) -> dict[str, Any]:
    pairs = pd.read_csv(pair_path).fillna("")
    review = pairs[pairs["train_faithfulness_head"].astype(int).eq(1)].copy()
    review["review_flag"] = np.where(
        review["pair_review_status"].astype(str).str.contains("provisional", case=False, na=False),
        "provisional_review_required",
        "ambiguous_or_mismatched_review_required",
    )
    review[
        ["pair_id", "source_text", "candidate_text", "faithfulness_label", "rationale", "review_flag", "recommended_split", "candidate_row_id"]
    ].to_csv(out / "faithfulness_provisional_review_table.csv", index=False)

    result: dict[str, Any] = {
        "rows": int(len(pairs)),
        "trainable_rows": int(len(review)),
        "class_distribution": split_counts(review, "faithfulness_label") if len(review) else {},
        "warning": "Provisional paired diagnostic only; not promoted to production behaviour.",
    }
    if len(review) >= 10 and review["faithfulness_label"].nunique() > 1:
        texts = ("[SOURCE] " + review["source_text"].astype(str) + " [CANDIDATE] " + review["candidate_text"].astype(str)).tolist()
        X = encode_texts(embedder, texts)
        le = LabelEncoder()
        y = le.fit_transform(review["faithfulness_label"].astype(str))
        train = review["recommended_split"].eq("train").to_numpy()
        validation = review["recommended_split"].eq("validation").to_numpy()
        test = review["recommended_split"].eq("test").to_numpy()
        if train.sum() and validation.sum() and test.sum() and len(set(y[train])) > 1:
            clf = LogisticRegression(max_iter=3000, class_weight="balanced")
            clf.fit(X[train], y[train])
            result["diagnostic_baseline"] = {}
            for split_name, mask in [("train", train), ("validation", validation), ("test", test)]:
                if mask.sum() == 0:
                    continue
                pred = clf.predict(X[mask])
                result["diagnostic_baseline"][split_name] = single_metrics(y[mask], pred, list(le.classes_))
            joblib.dump({"label_encoder": le, "classifier": clf}, out / "faithfulness_diagnostic_classifier.joblib")
        else:
            result["diagnostic_baseline_skipped"] = "Not enough labelled split coverage."
    return result


def fmt_pct(value: float | None) -> str:
    if value is None or (isinstance(value, float) and math.isnan(value)):
        return "n/a"
    return f"{value * 100:.1f}%"


def write_report(out: Path, metrics: dict[str, Any], baseline: dict[str, Any] | None) -> None:
    lines = [
        "# Meeting minutes classifier repaired masking run",
        "",
        f"Run directory: `{out}`",
        "",
        "## Plain-English summary",
        "",
    ]

    comparisons = []
    if baseline:
        comparisons = [
            ("Evidence type", baseline.get("overall_status_cv_macro_f1"), metrics["heads"]["evidence_type"]["metrics"]["test"]["macro_f1"]),
            ("Action state", baseline.get("action_state_cv_macro_f1"), metrics["heads"]["action_state"]["metrics"]["test"]["macro_f1"]),
            ("Signals", baseline.get("signals_cv_macro_f1"), metrics["heads"]["signals"]["metrics"]["test"]["macro_f1"]),
        ]
        for label, before, after in comparisons:
            delta = None if before is None else after - before
            direction = "better" if delta and delta > 0 else "worse" if delta and delta < 0 else "similar"
            lines.append(f"- {label}: {fmt_pct(before)} before versus {fmt_pct(after)} on repaired test; {direction} by {fmt_pct(abs(delta)) if delta is not None else 'n/a'}.")
    lines.extend(
        [
            f"- Quality gate: {fmt_pct(metrics['heads']['quality_gate']['metrics']['test']['micro_f1'])} test accuracy/micro-F1 and {fmt_pct(metrics['heads']['quality_gate']['metrics']['test']['macro_f1'])} macro-F1.",
            f"- Quality defects: {fmt_pct(metrics['heads']['quality_defects']['metrics']['test']['micro_f1'])} test micro-F1 and {fmt_pct(metrics['heads']['quality_defects']['metrics']['test']['macro_f1'])} macro-F1.",
            "",
            "## Before and after metrics",
            "",
            "| Head | Before | Repaired test | Change |",
            "| --- | ---: | ---: | ---: |",
        ]
    )
    if baseline:
        for label, before, after in comparisons:
            delta = None if before is None else after - before
            lines.append(f"| {label} macro-F1 | {fmt_pct(before)} | {fmt_pct(after)} | {fmt_pct(delta) if delta is not None else 'n/a'} |")
    lines.extend(
        [
            f"| Quality gate macro-F1 | n/a | {fmt_pct(metrics['heads']['quality_gate']['metrics']['test']['macro_f1'])} | new |",
            f"| Quality gate micro-F1 | n/a | {fmt_pct(metrics['heads']['quality_gate']['metrics']['test']['micro_f1'])} | new |",
            f"| Quality defects macro-F1 | n/a | {fmt_pct(metrics['heads']['quality_defects']['metrics']['test']['macro_f1'])} | new |",
            f"| Quality defects micro-F1 | n/a | {fmt_pct(metrics['heads']['quality_defects']['metrics']['test']['micro_f1'])} | new |",
            "",
            "## Quality-gate confusion matrix",
            "",
        ]
    )
    q = metrics["heads"]["quality_gate"]["metrics"]["test"]
    labels = q["labels"]
    cm = q["confusion_matrix"]
    lines.append("| Actual \\ predicted | " + " | ".join(labels) + " |")
    lines.append("| --- | " + " | ".join(["---:"] * len(labels)) + " |")
    for label, row in zip(labels, cm):
        lines.append("| " + label + " | " + " | ".join(str(x) for x in row) + " |")

    lines.extend(["", "## Weakest classes and labels", ""])
    weak = []
    for head_name in ["evidence_type", "action_state", "quality_gate"]:
        per = metrics["heads"][head_name]["metrics"]["test"]["per_class"]
        for klass, stats in per.items():
            if isinstance(stats, dict) and "f1-score" in stats:
                weak.append((stats["f1-score"], head_name, klass, stats.get("support", 0)))
    for label, stats in metrics["heads"]["signals"]["metrics"]["test"]["per_label"].items():
        weak.append((stats["f1"], "signals", label, stats["support"]))
    for label, stats in metrics["heads"]["quality_defects"]["metrics"]["test"]["per_label"].items():
        weak.append((stats["f1"], "quality_defects", label, stats["support"]))
    weak.sort(key=lambda item: item[0])
    for score, head_name, klass, support in weak[:15]:
        lines.append(f"- {head_name} / {klass}: {fmt_pct(score)} F1 on {int(support)} test examples.")

    lines.extend(
        [
            "",
            "## Suitability recommendation",
            "",
            "| Head | Diagnostics | Rewrite warnings | Human-review routing | Automatic approval | Automatic deletion |",
            "| --- | --- | --- | --- | --- | --- |",
            "| Evidence type | yes | supporting only | yes | no | no |",
            "| Action state | yes | supporting only | yes | no | no |",
            "| Signals | yes | supporting only | yes | no | no |",
            "| Quality gate | yes | yes | yes | no | no |",
            "| Quality defects | yes | yes | yes | no | no |",
            "| Faithfulness | provisional only | no | possible later | no | no |",
            "",
            "## Highest-value next examples",
            "",
            "Prioritise 50-100 examples that cover:",
            "",
            "1. Real generated minutes rows that look fluent but are vague, especially `topic_only_or_missing_substance`.",
            "2. Real generated minutes rows with reported speech: `said that`, `mentioned that`, `there was a discussion around`.",
            "3. Short false starts and chairing phrases from different meeting types.",
            "4. Good human-minute rows with first names in the body only, so the gate learns that names are not the defect.",
            "5. Negative faithfulness pairs where the candidate is plausible but loses a deadline, owner, decision or constraint.",
            "6. More decision-language examples; this remains a weak old-head signal.",
            "7. More risk/process examples, because those dipped or remain fragile.",
            "",
            "## Implementation issues and assumptions",
            "",
            "- The repaired run uses MiniLM exactly, so this isolates schema repair rather than embedding-model changes.",
            "- Faithfulness labels remain provisional; the pair test split currently lacks unsupported examples, so faithfulness metrics are not production evidence.",
            "- No live minutes-generation behaviour has been changed.",
        ]
    )
    (out / "REPORT.md").write_text("\n".join(lines) + "\n")


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--master-csv", required=True)
    ap.add_argument("--faithfulness-csv", required=True)
    ap.add_argument("--schema", default="labels/status_schema.json")
    ap.add_argument("--baseline-metrics", default="runs/owner-framing-20260810T224329Z/model_bundle/metrics.json")
    ap.add_argument("--output-root", default="runs")
    ap.add_argument("--embedding-model", default="sentence-transformers/all-MiniLM-L6-v2")
    ap.add_argument("--expected-master-rows", type=int, default=EXPECTED["master_rows"])
    ap.add_argument("--expected-evidence-rows", type=int, default=EXPECTED["evidence_rows"])
    ap.add_argument("--expected-quality-rows", type=int, default=EXPECTED["minutes_quality_rows"])
    ap.add_argument("--expected-quality-ready", type=int, default=EXPECTED["quality_ready"])
    ap.add_argument("--expected-quality-intervention", type=int, default=EXPECTED["quality_intervention"])
    ap.add_argument("--expected-split-train", type=int, default=EXPECTED["split_train"])
    ap.add_argument("--expected-split-validation", type=int, default=EXPECTED["split_validation"])
    ap.add_argument("--expected-split-test", type=int, default=EXPECTED["split_test"])
    ap.add_argument("--quality-gate-c", type=float, default=1.0)
    ap.add_argument("--quality-defect-c", type=float, default=1.0)
    args = ap.parse_args()

    out = Path(args.output_root) / f"multitask-rework-{now_slug()}" / "model_bundle"
    out.mkdir(parents=True, exist_ok=False)

    df = pd.read_csv(args.master_csv)
    expected = {
        "master_rows": args.expected_master_rows,
        "evidence_rows": args.expected_evidence_rows,
        "minutes_quality_rows": args.expected_quality_rows,
        "quality_ready": args.expected_quality_ready,
        "quality_intervention": args.expected_quality_intervention,
        "split_train": args.expected_split_train,
        "split_validation": args.expected_split_validation,
        "split_test": args.expected_split_test,
    }
    checks = validate_master(df, expected)
    schema = json.loads(Path(args.schema).read_text())
    shutil.copy2(args.master_csv, out / "training_data_used.csv")
    shutil.copy2(args.faithfulness_csv, out / "faithfulness_pairs_used.csv")
    (out / "validation_checks.json").write_text(json.dumps(checks, indent=2))

    print(f"Rows: {len(df)}")
    print(f"Output: {out}")
    print(f"Embedding model: {args.embedding_model}")
    embedder = SentenceTransformer(args.embedding_model)
    X = encode_texts(embedder, df["model_text"].fillna("").astype(str).tolist())

    signal_labels = list(schema["signals"]["labels"].keys())
    quality_defect_labels = [
        "none",
        "transcript_chatter",
        "fragment_or_false_start",
        "speaker_reported_speech",
        "topic_only_or_missing_substance",
    ]
    heads = {
        "evidence_type": train_single_head("evidence_type", df, "train_evidence_head", "evidence_type", X, out),
        "action_state": train_single_head("action_state", df, "train_action_head", "action_state", X, out),
        "signals": train_multilabel_head("signals", df, "train_signal_head", "signals", X, out, signal_labels),
        "quality_gate": train_single_head(
            "quality_gate",
            df,
            "train_quality_gate",
            "quality_gate",
            X,
            out,
            args.quality_gate_c,
        ),
        "quality_defects": train_multilabel_head(
            "quality_defects",
            df,
            "train_quality_defect_head",
            "quality_defects",
            X,
            out,
            quality_defect_labels,
            args.quality_defect_c,
        ),
    }
    for head in ["evidence_type", "action_state", "quality_gate"]:
        add_reliability(heads[head], out / f"{head}_predictions.csv")
    heads["quality_gate"]["breakdowns"] = quality_breakdowns(df, out / "quality_gate_predictions.csv")

    faithfulness = create_faithfulness_review(Path(args.faithfulness_csv), out, embedder)

    baseline = None
    if Path(args.baseline_metrics).exists():
        baseline = json.loads(Path(args.baseline_metrics).read_text())

    metrics = {
        "run_directory": str(out),
        "embedding_model": args.embedding_model,
        "rows": int(len(df)),
        "validation_checks": checks,
        "baseline_metrics": baseline,
        "heads": heads,
        "faithfulness": faithfulness,
    }
    (out / "metrics.json").write_text(json.dumps(metrics, indent=2, default=lambda o: dict(o) if isinstance(o, Counter) else str(o)))
    write_report(out, metrics, baseline)
    print(f"Saved repaired multitask run to: {out}")
    print(f"Report: {out / 'REPORT.md'}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
