"""Train and score the Bleu.js forest-cover XGBoost classifier.

Protocol, fixed before the test set is scored:
- Data: sklearn.datasets.fetch_covtype (UCI Covertype, full table)
- Labels shifted from 1..7 to 0..6
- Outer split: stratified 80/20, random_state=42. The 20% is scored once.
- Inner split: stratified 90/10 of the training rows, random_state=42,
  used only to pick n_estimators with early stopping.
- Final fit: all training rows, n_estimators = best_iteration + 1,
  sample_weight = balanced class weights from the training rows.
- Baselines, same test rows, no retuning: majority class; a depth-3
  XGBoost with 50 trees and no class weights.
"""

import json
from pathlib import Path

import numpy as np
import xgboost as xgb
from sklearn.datasets import fetch_covtype
from sklearn.metrics import (
    accuracy_score,
    classification_report,
    f1_score,
    log_loss,
)
from sklearn.model_selection import train_test_split
from sklearn.utils.class_weight import compute_sample_weight

OUT = Path(__file__).resolve().parent
RANDOM_STATE = 42
CLASS_NAMES = [
    "Spruce/Fir",
    "Lodgepole Pine",
    "Ponderosa Pine",
    "Cottonwood/Willow",
    "Aspen",
    "Douglas-fir",
    "Krummholz",
]
FEATURE_NAMES = [
    "Elevation",
    "Aspect",
    "Slope",
    "Horizontal_Distance_To_Hydrology",
    "Vertical_Distance_To_Hydrology",
    "Horizontal_Distance_To_Roadways",
    "Hillshade_9am",
    "Hillshade_Noon",
    "Hillshade_3pm",
    "Horizontal_Distance_To_Fire_Points",
    *[f"Wilderness_Area_{i}" for i in range(4)],
    *[f"Soil_Type_{i}" for i in range(40)],
]


def scores(y_true, pred, proba):
    return {
        "accuracy": float(accuracy_score(y_true, pred)),
        "macro_f1": float(f1_score(y_true, pred, average="macro")),
        "weighted_f1": float(f1_score(y_true, pred, average="weighted")),
        "log_loss": float(log_loss(y_true, proba, labels=list(range(7)))),
    }


def main():
    bunch = fetch_covtype(as_frame=False)
    X = np.asarray(bunch.data, dtype=np.float32)
    y = np.asarray(bunch.target, dtype=np.int32) - 1
    if X.shape[1] != len(FEATURE_NAMES) or y.min() != 0 or y.max() != 6:
        raise SystemExit(f"Unexpected covertype shape {X.shape} labels {y.min()}..{y.max()}")

    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=RANDOM_STATE, stratify=y
    )
    X_fit, X_val, y_fit, y_val = train_test_split(
        X_train, y_train, test_size=0.1, random_state=RANDOM_STATE, stratify=y_train
    )

    common = dict(
        max_depth=8,
        learning_rate=0.1,
        subsample=0.8,
        colsample_bytree=0.8,
        min_child_weight=5,
        reg_lambda=1.0,
        tree_method="hist",
        objective="multi:softprob",
        num_class=7,
        random_state=RANDOM_STATE,
        n_jobs=-1,
    )
    search = xgb.XGBClassifier(n_estimators=400, early_stopping_rounds=20, **common)
    search.fit(
        X_fit,
        y_fit,
        sample_weight=compute_sample_weight("balanced", y_fit),
        eval_set=[(X_val, y_val)],
        verbose=False,
    )
    n_trees = int(search.best_iteration) + 1

    model = xgb.XGBClassifier(n_estimators=n_trees, **common)
    model.fit(
        X_train,
        y_train,
        sample_weight=compute_sample_weight("balanced", y_train),
        verbose=False,
    )

    simple = xgb.XGBClassifier(
        n_estimators=50,
        max_depth=3,
        learning_rate=0.1,
        tree_method="hist",
        objective="multi:softprob",
        num_class=7,
        random_state=RANDOM_STATE,
        n_jobs=-1,
    )
    simple.fit(X_train, y_train, verbose=False)

    proba = model.predict_proba(X_test)
    pred = proba.argmax(axis=1)
    simple_proba = simple.predict_proba(X_test)
    simple_pred = simple_proba.argmax(axis=1)
    majority = int(np.bincount(y_train).argmax())
    majority_pred = np.full_like(y_test, majority)
    majority_proba = np.zeros_like(proba)
    majority_proba[:, majority] = 1.0

    report = classification_report(
        y_test, pred, target_names=CLASS_NAMES, output_dict=True, digits=4
    )
    gain = model.get_booster().get_score(importance_type="gain")
    importance = sorted(
        (
            {"feature": FEATURE_NAMES[int(k[1:])], "gain": float(v)}
            for k, v in gain.items()
        ),
        key=lambda row: row["gain"],
        reverse=True,
    )

    payload = {
        "dataset": "UCI Covertype",
        "source": "sklearn.datasets.fetch_covtype",
        "citation": "Blackard, Jock A. Covertype. UCI Machine Learning Repository.",
        "n_rows": int(X.shape[0]),
        "n_features": int(X.shape[1]),
        "class_names": CLASS_NAMES,
        "label_encoding": "model class i is UCI cover type i+1",
        "split": {
            "test_size": 0.2,
            "inner_val_size": 0.1,
            "random_state": RANDOM_STATE,
            "stratify": True,
        },
        "n_train": int(len(y_train)),
        "n_test": int(len(y_test)),
        "n_estimators_selected_on_inner_val": n_trees,
        "hyperparameters": {**common, "n_estimators": n_trees, "sample_weight": "balanced"},
        "xgboost_version": xgb.__version__,
        "trainer_note": (
            "XGBClassifier. bleujs.ml.HybridTrainer targets binary classification "
            "with binary:logistic, so this multiclass model is trained with the "
            "same library and a documented multiclass objective."
        ),
        "test_metrics": scores(y_test, pred, proba),
        "simple_xgboost_baseline": scores(y_test, simple_pred, simple_proba),
        "majority_baseline": scores(y_test, majority_pred, majority_proba),
        "per_class_f1": {name: float(report[name]["f1-score"]) for name in CLASS_NAMES},
        "per_class_support": {name: int(report[name]["support"]) for name in CLASS_NAMES},
        "feature_importance_gain": importance,
    }
    model.save_model(OUT / "model.json")
    (OUT / "metrics.json").write_text(json.dumps(payload, indent=2))
    print(json.dumps({
        "n_trees": n_trees,
        "test": payload["test_metrics"],
        "simple": payload["simple_xgboost_baseline"],
        "majority": payload["majority_baseline"],
        "per_class_f1": payload["per_class_f1"],
    }, indent=2))


if __name__ == "__main__":
    main()