from __future__ import annotations

import json
from pathlib import Path

import joblib
from sklearn.datasets import load_iris
from sklearn.model_selection import GridSearchCV
from sklearn.neighbors import KNeighborsClassifier
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler


ARTIFACT_PATH = Path(__file__).with_name("model.joblib")
METRICS_PATH = Path(__file__).with_name("metrics.json")


def train_and_save(artifact_path: Path = ARTIFACT_PATH) -> dict:
    iris = load_iris(as_frame=True)
    X = iris.data
    y = iris.target

    target_names = [str(name) for name in iris.target_names]
    feature_names = [str(name) for name in iris.feature_names]

    pipeline = Pipeline(
        steps=[
            ("scaler", StandardScaler()),
            ("knn", KNeighborsClassifier()),
        ]
    )

    param_grid = {
        "knn__n_neighbors": list(range(1, 21)),
        "knn__weights": ["uniform", "distance"],
        "knn__p": [1, 2],
    }

    search = GridSearchCV(
        estimator=pipeline,
        param_grid=param_grid,
        cv=5,
        scoring="accuracy",
        n_jobs=-1,
        refit=True,
    )
    search.fit(X, y)

    best_model = search.best_estimator_
    joblib.dump(
        {
            "model": best_model,
            "target_names": target_names,
            "feature_names": feature_names,
        },
        artifact_path,
    )

    metrics = {
        "cv_best_accuracy": float(search.best_score_),
        "best_params": search.best_params_,
    }
    METRICS_PATH.write_text(json.dumps(metrics, indent=2), encoding="utf-8")
    return metrics


if __name__ == "__main__":
    metrics = train_and_save()
    print(f"Saved model to: {ARTIFACT_PATH}")
    print(json.dumps(metrics, indent=2))