"""Train 23 binary XGBoost classifiers for page structure labels.

Usage:
    agentic/.venv/bin/python -m classifier.train
    agentic/.venv/bin/python -m classifier.train --eval-only
"""

from __future__ import annotations

import argparse
import json
import logging
import sys
from pathlib import Path

import joblib
import numpy as np
import pandas as pd
from sklearn.metrics import (
    accuracy_score,
    f1_score,
    precision_score,
    recall_score,
    roc_auc_score,
)
from sklearn.model_selection import train_test_split

sys.path.insert(0, str(Path(__file__).parent.parent))

from classifier.features import FEATURE_NAMES

logger = logging.getLogger(__name__)

DATA_DIR = Path(__file__).parent / "data"
BUNDLE_PATH = Path(__file__).parent / "model_bundle.joblib"

LABEL_COLS = [
    "has_handwriting",
    "has_hand_filled_form",
    "has_rotated_90",
    "has_low_quality_scan",
    "has_boxed_fields",
    "has_organizing_boxes",
    "has_checkboxes",
    "has_table",
    "has_multiple_tables",
    "has_column_headers",
    "table_has_h_ruling_lines",
    "table_has_h_zebra_stripes",
    "table_has_h_whitespace",
    "table_has_v_ruling_lines",
    "table_has_v_whitespace",
    "has_key_value_pairs",
    "has_section_headers",
    "has_narrative",
    "has_multi_column_text",
    "has_empty_form",
    "has_fill_in_blanks",
    "has_visual_media",
    "has_multiple_records",
]

# XGBoost common hyperparameters
XGB_PARAMS = {
    "n_estimators": 200,
    "max_depth": 4,
    "learning_rate": 0.1,
    "min_child_weight": 5,
    "subsample": 0.8,
    "colsample_bytree": 0.8,
    "objective": "binary:logistic",
    "eval_metric": "logloss",
    "random_state": 42,
    "n_jobs": -1,
    "verbosity": 0,
}


def load_data() -> pd.DataFrame:
    """Load and merge features + labels."""
    features_df = pd.read_parquet(DATA_DIR / "features.parquet")
    labels_df = pd.read_parquet(DATA_DIR / "labels.parquet")

    df = features_df.merge(labels_df, on=["pdf_path", "page_number"], how="inner")
    logger.info("Loaded %d rows (%d features, %d labels)", len(df), len(FEATURE_NAMES), len(LABEL_COLS))
    return df


def find_best_threshold(y_true, y_prob):
    """Sweep thresholds to find F1-maximizing value."""
    best_f1 = 0.0
    best_t = 0.5
    for t in np.arange(0.10, 0.91, 0.05):
        preds = (y_prob >= t).astype(int)
        score = f1_score(y_true, preds, zero_division=0)
        if score > best_f1:
            best_f1 = score
            best_t = t
    return round(float(best_t), 2)


def train_all(df: pd.DataFrame, eval_only: bool = False):
    """Train 23 binary XGBoost classifiers."""
    import xgboost as xgb

    X = df[FEATURE_NAMES].values
    idx = df.index.values

    X_train, X_test, idx_train, idx_test = train_test_split(
        X, idx, test_size=0.2, random_state=42
    )

    logger.info("Train: %d, Test: %d", len(idx_train), len(idx_test))

    models = {}
    thresholds = {}
    metrics = {}

    # Header
    print()
    print(f"{'Label':35s} {'Acc':>5s} {'F1':>5s} {'Prec':>5s} {'Rec':>5s} {'AUC':>5s} {'Pos':>5s} {'Thr':>5s}")
    print("-" * 80)

    for label in LABEL_COLS:
        y = df[label].values.astype(int)
        y_train = y[idx_train]
        y_test = y[idx_test]

        n_pos = int(y_train.sum())
        n_neg = int((y_train == 0).sum())

        if n_pos < 3:
            logger.warning("Skipping %s — only %d positives in train", label, n_pos)
            continue

        scale_pos_weight = n_neg / max(n_pos, 1)

        clf = xgb.XGBClassifier(**XGB_PARAMS, scale_pos_weight=scale_pos_weight)
        clf.fit(
            X_train,
            y_train,
            eval_set=[(X_test, y_test)],
            verbose=False,
        )

        y_prob = clf.predict_proba(X_test)[:, 1]
        best_t = find_best_threshold(y_test, y_prob)
        y_pred = (y_prob >= best_t).astype(int)

        test_pos = int(y_test.sum())
        acc = accuracy_score(y_test, y_pred)
        f1 = f1_score(y_test, y_pred, zero_division=0)
        prec = precision_score(y_test, y_pred, zero_division=0)
        rec = recall_score(y_test, y_pred, zero_division=0)
        try:
            auc = roc_auc_score(y_test, y_prob)
        except ValueError:
            auc = 0.0

        print(
            f"{label:35s} {acc:5.2f} {f1:5.2f} {prec:5.2f} {rec:5.2f} {auc:5.2f} {test_pos:5d} {best_t:5.2f}"
        )

        models[label] = clf
        thresholds[label] = best_t
        metrics[label] = {
            "accuracy": round(acc, 4),
            "f1": round(f1, 4),
            "precision": round(prec, 4),
            "recall": round(rec, 4),
            "roc_auc": round(auc, 4),
            "test_pos": test_pos,
            "train_pos": n_pos,
            "threshold": best_t,
        }

    print("-" * 80)

    # Summary stats
    f1_values = [m["f1"] for m in metrics.values()]
    auc_values = [m["roc_auc"] for m in metrics.values()]
    print(f"Mean F1: {np.mean(f1_values):.3f}, Mean AUC: {np.mean(auc_values):.3f}")
    print(f"Models trained: {len(models)}/{len(LABEL_COLS)}")

    # Feature importance (average across models)
    if models:
        importances = np.zeros(len(FEATURE_NAMES))
        for clf in models.values():
            importances += clf.feature_importances_
        importances /= len(models)
        top_idx = np.argsort(importances)[::-1][:15]
        print("\nTop 15 features (avg importance):")
        for i in top_idx:
            print(f"  {FEATURE_NAMES[i]:40s} {importances[i]:.4f}")

    if not eval_only:
        bundle = {
            "models": models,
            "thresholds": thresholds,
            "feature_names": FEATURE_NAMES,
            "label_names": LABEL_COLS,
            "metrics": metrics,
            "train_size": len(idx_train),
            "test_size": len(idx_test),
        }
        joblib.dump(bundle, BUNDLE_PATH)
        logger.info("\nSaved model bundle to %s", BUNDLE_PATH)

        # Also save metrics as JSON for easy inspection
        metrics_path = DATA_DIR / "eval_report.json"
        with open(metrics_path, "w") as fp:
            json.dump(metrics, fp, indent=2)
        logger.info("Saved eval report to %s", metrics_path)


def main():
    parser = argparse.ArgumentParser(description="Train page structure classifiers")
    parser.add_argument("--eval-only", action="store_true", help="Evaluate without saving")
    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO, format="%(message)s")

    df = load_data()
    train_all(df, eval_only=args.eval_only)


if __name__ == "__main__":
    main()
