"""Учебное сравнение заранее заданных моделей на встроенном Wine.

Установка в своём окружении: python -m pip install scikit-learn
Запуск: python wine_baseline.py
Не используйте отложенную часть для дальнейшего подбора параметров.
"""
from sklearn.datasets import load_wine
from sklearn.dummy import DummyClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, confusion_matrix, f1_score
from sklearn.model_selection import train_test_split
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler


def compare():
    data = load_wine()
    indices = list(range(len(data.target)))
    train, test = train_test_split(
        indices, test_size=0.25, stratify=data.target, random_state=42
    )
    X_train, X_test = data.data[train], data.data[test]
    y_train, y_test = data.target[train], data.target[test]
    models = {
        "most_frequent": DummyClassifier(strategy="most_frequent"),
        "scaled_logistic": make_pipeline(
            StandardScaler(), LogisticRegression(max_iter=1000, random_state=42)
        ),
    }
    results = {}
    for name, model in models.items():
        model.fit(X_train, y_train)
        prediction = model.predict(X_test)
        results[name] = {
            "accuracy": accuracy_score(y_test, prediction),
            "macro_f1": f1_score(y_test, prediction, average="macro"),
            "confusion_matrix": confusion_matrix(y_test, prediction).tolist(),
        }
    return {"shape": data.data.shape, "train": len(train), "test": len(test),
            "overlap": len(set(train) & set(test)), "results": results}


if __name__ == "__main__":
    from pprint import pprint
    pprint(compare(), sort_dicts=False)