#!/usr/bin/env python3
"""Train/evaluate meeting-minutes evidence classifiers.

Outputs one model bundle containing:
- overall_status classifier, used here as evidence_type (single-label)
- action_state classifier, used here as commitment_state (single-label)
- signals classifier (multi-label evidence cues)
"""
from __future__ import annotations

import argparse
import json
import warnings
from pathlib import Path

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, f1_score, precision_recall_fscore_support
from sklearn.model_selection import StratifiedKFold, cross_val_predict
from sklearn.multiclass import OneVsRestClassifier
from sklearn.preprocessing import LabelEncoder, MultiLabelBinarizer

warnings.filterwarnings("ignore")


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


def safe_n_splits(y: np.ndarray, requested: int) -> int:
    _, counts = np.unique(y, return_counts=True)
    if len(counts) < 2:
        return 0
    return int(max(0, min(requested, counts.min())))


def train_single_label(name: str, labels: pd.Series, X: np.ndarray, out: Path, requested_splits: int, metrics: dict):
    le = LabelEncoder()
    y = le.fit_transform(labels.astype(str))
    clf = LogisticRegression(max_iter=2000, class_weight="balanced")
    n_splits = safe_n_splits(y, requested_splits)
    if n_splits >= 2:
        cv = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=42)
        pred = cross_val_predict(clf, X, y, cv=cv)
        display_name = {"overall_status": "evidence_type", "action_state": "commitment_state"}.get(name, name)
        print(f"\n{display_name} CV report:")
        report = classification_report(y, pred, target_names=le.classes_, zero_division=0)
        print(report)
        metrics[f"{name}_cv_macro_f1"] = float(f1_score(y, pred, average="macro", zero_division=0))
        metrics[f"{name}_cv_micro_f1"] = float(f1_score(y, pred, average="micro", zero_division=0))
    else:
        print(f"\n{name} CV skipped: not enough examples per class yet.")
        metrics[f"{name}_cv_skipped"] = True
    clf.fit(X, y)
    return le, clf


def train_signals(signals: list[list[str]], schema: dict, X: np.ndarray, df_len: int, out: Path, requested_splits: int, metrics: dict):
    all_signal_labels = list(schema["signals"]["labels"].keys())
    full_mlb = MultiLabelBinarizer(classes=all_signal_labels)
    Y_full = full_mlb.fit_transform(signals)
    support = Y_full.sum(axis=0)
    active_idx = np.where(support > 0)[0]
    active_labels = [all_signal_labels[i] for i in active_idx]
    if not active_labels:
        print("\nSignal training skipped: no signal labels present yet.")
        metrics["signals_training_skipped"] = True
        return full_mlb, None

    mlb = MultiLabelBinarizer(classes=active_labels)
    Y = mlb.fit_transform([[s for s in row if s in active_labels] for row in signals])
    clf = OneVsRestClassifier(LogisticRegression(max_iter=2000, class_weight="balanced"))

    if df_len >= 10 and int(Y.sum(axis=0).max()) > 1:
        rng = np.random.default_rng(42)
        idx = np.arange(df_len)
        rng.shuffle(idx)
        folds = np.array_split(idx, min(requested_splits, df_len))
        oof = np.zeros_like(Y, dtype=float)
        evaluated_any = False
        for test_idx in folds:
            train_idx = np.setdiff1d(idx, test_idx)
            train_support = Y[train_idx].sum(axis=0)
            fold_active = np.where((train_support > 0) & (train_support < len(train_idx)))[0]
            if len(fold_active) == 0:
                continue
            fold_clf = OneVsRestClassifier(LogisticRegression(max_iter=2000, class_weight="balanced"))
            fold_clf.fit(X[train_idx], Y[train_idx][:, fold_active])
            try:
                probs = np.asarray(fold_clf.predict_proba(X[test_idx]))
            except Exception:
                probs = np.asarray(fold_clf.predict(X[test_idx]), dtype=float)
            if probs.ndim == 1:
                probs = probs.reshape(-1, 1)
            oof[np.ix_(test_idx, fold_active)] = probs
            evaluated_any = True
        if evaluated_any:
            pred = (oof >= 0.5).astype(int)
            micro = f1_score(Y, pred, average="micro", zero_division=0)
            macro = f1_score(Y, pred, average="macro", zero_division=0)
            print("\nSignals CV @0.5, active labels only:")
            print(f"  micro_f1: {micro:.3f}")
            print(f"  macro_f1: {macro:.3f}")
            rows=[]
            pr, rc, f1, sup = precision_recall_fscore_support(Y, pred, average=None, zero_division=0)
            for label, p, r, f, s in zip(mlb.classes_, pr, rc, f1, sup):
                rows.append({"signal": label, "support": int(s), "precision": float(p), "recall": float(r), "f1": float(f)})
            per_signal = pd.DataFrame(rows)
            print(per_signal.to_string(index=False, float_format=lambda x: f"{x:.3f}"))
            per_signal.to_csv(out / "per_signal_cv_metrics.csv", index=False)
            metrics["signals_cv_micro_f1"] = float(micro)
            metrics["signals_cv_macro_f1"] = float(macro)
        else:
            print("\nSignal CV skipped: each fold was too sparse for learnable labels.")
            metrics["signals_cv_skipped"] = True
    else:
        print("\nSignal CV skipped: not enough rows/signal support yet.")
        metrics["signals_cv_skipped"] = True

    learnable = np.where((Y.sum(axis=0) > 0) & (Y.sum(axis=0) < df_len))[0]
    if len(learnable) == 0:
        print("Signal final training skipped: no learnable active labels.")
        return mlb, None
    if len(learnable) < Y.shape[1]:
        kept = [active_labels[i] for i in learnable]
        print("Signal final training uses learnable labels only:", ", ".join(kept))
        mlb = MultiLabelBinarizer(classes=kept)
        Y = mlb.fit_transform([[s for s in row if s in kept] for row in signals])
    clf.fit(X, Y)
    return mlb, clf


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", default="data/training_data.csv")
    ap.add_argument("--schema", default="labels/status_schema.json")
    ap.add_argument("--output", default="model_bundle")
    ap.add_argument("--embedding-model", default="sentence-transformers/all-MiniLM-L6-v2")
    ap.add_argument("--splits", type=int, default=5)
    ap.add_argument("--reviewed-only", action="store_true")
    args = ap.parse_args()

    out = Path(args.output); out.mkdir(parents=True, exist_ok=True)
    schema = json.loads(Path(args.schema).read_text())
    df = pd.read_csv(args.data).fillna("")
    if "action_state" not in df.columns:
        df["action_state"] = "no_action"
    if args.reviewed_only:
        df = df[df["reviewed"].astype(str).str.lower().isin(["true", "1", "yes", "y"])]
    df = df[df["text"].astype(str).str.strip() != ""].copy()
    if len(df) < 5:
        print("Not enough rows to train/evaluate yet. Add more labelled examples first.")
        return 2

    print(f"Rows: {len(df)}")
    print(f"Embedding model: {args.embedding_model}")
    embedder = SentenceTransformer(args.embedding_model)
    X = embedder.encode(df["text"].astype(str).tolist(), normalize_embeddings=True, convert_to_numpy=True, batch_size=32, show_progress_bar=True)

    metrics = {"rows": int(len(df)), "embedding_model": args.embedding_model}
    status_le, status_clf = train_single_label("overall_status", df["overall_status"], X, out, args.splits, metrics)
    action_le = action_clf = None
    if schema.get("action_state"):
        action_le, action_clf = train_single_label("action_state", df["action_state"], X, out, args.splits, metrics)
    signal_mlb, signal_clf = train_signals(df["signals"].apply(split_signals).tolist(), schema, X, len(df), out, args.splits, metrics)

    bundle = {
        "schema": schema,
        "embedding_model": args.embedding_model,
        "status_label_encoder": status_le,
        "status_classifier": status_clf,
        "action_state_label_encoder": action_le,
        "action_state_classifier": action_clf,
        "signal_binarizer": signal_mlb,
        "all_signal_labels": list(schema["signals"]["labels"].keys()),
        "signal_classifier": signal_clf,
        "signal_threshold": 0.5,
    }
    joblib.dump(bundle, out / "classifier.joblib")
    (out / "metrics.json").write_text(json.dumps(metrics, indent=2))
    df.to_csv(out / "training_data_used.csv", index=False)
    print(f"\nSaved model bundle to: {out}")
    return 0

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