"""Example 2: Upload a Scikit-learn Model.

Trains a simple sklearn classifier on the Iris dataset and uploads it
to the local MLflow tracking server using ModelUploader.

Prerequisites:
    pip install mlflow-utils[sklearn]

Run:
    python examples/02_upload_sklearn.py
"""

from dgtw_mlflow_utils import (
    ModelUploader,
    ModelVariant,
    RunAction,
    Stage,
    UploadConfig,
)
from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score
from sklearn.model_selection import train_test_split


def main():
    """Main function to execute the example."""
    # ------------------------------------------------------------------
    # 1) Train a simple model
    # ------------------------------------------------------------------
    print("Loading Iris dataset...")
    X, y = load_iris(return_X_y=True, as_frame=True)
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    print("Training RandomForestClassifier...")
    model = RandomForestClassifier(n_estimators=100, random_state=42)
    model.fit(X_train, y_train)

    y_pred = model.predict(X_test)
    accuracy = accuracy_score(y_test, y_pred)
    print(f"Test accuracy: {accuracy:.4f}")

    # ------------------------------------------------------------------
    # 2) Configure the upload
    # ------------------------------------------------------------------
    config = UploadConfig(
        architecture="RandomForest",  # custom architecture string
        variant=ModelVariant.BASELINE,
        component="Classification",
        stage=Stage.RESEARCH,
        use_case="IrisDemo",
        action=RunAction.UPLOAD,
        tracking_uri="http://127.0.0.1:5001",  # local MLflow server
        params={
            "n_estimators": 100,
            "test_accuracy": accuracy,
        },
        tags={
            "dataset": "iris",
            "developer": "ai-team",
        },
    )

    # ------------------------------------------------------------------
    # 3) Upload to MLflow
    # ------------------------------------------------------------------
    uploader = ModelUploader(config)

    # Print the resolved naming before uploading
    print()
    print(uploader.summary())
    print()

    result = uploader.upload_sklearn(
        sk_model=model,
        input_example=X_test.head(5),  # 5-row sample for documentation
    )

    print("Upload successful!")
    print(f"  Run ID             : {result['run_id']}")
    print(f"  Run Name           : {result['run_name']}")
    print(f"  Model URI          : {result['model_uri']}")
    print(f"  Registered Model   : {result['registered_model_name']}")


if __name__ == "__main__":
    main()
