#!/usr/bin/env catnip
# Comparaison de classifieurs scikit-learn sur le dataset Wine.
# Catnip orchestre le cycle complet : load, split, train, predict, évaluation.
#
# DEPS: scikit-learn

import('sklearn.datasets', 'load_wine')
import('sklearn.model_selection', 'train_test_split')
import('sklearn.pipeline', 'make_pipeline')
import('sklearn.preprocessing', 'StandardScaler')
import('sklearn.linear_model', 'LogisticRegression')
import('sklearn.ensemble', 'RandomForestClassifier')
import('sklearn.svm', 'SVC')
import('sklearn.metrics', 'accuracy_score', 'f1_score', 'confusion_matrix')
import('sklearn.base', 'clone')
time = import('time')

print("⇒ scikit-learn — bake-off sur le dataset Wine")

wine = load_wine()
X = wine.data
y = wine.target
class_names = wine.target_names

print(f"  échantillons : {int(X.shape[0])}")
print(f"  features     : {int(X.shape[1])}")
print(f"  classes      : {', '.join(class_names)}")

split = train_test_split(X, y, test_size=0.3, random_state=42, stratify=y)
X_train = split[0]
X_test = split[1]
y_train = split[2]
y_test = split[3]

print(f"  train/test   : {len(y_train)} / {len(y_test)}")

struct Candidate {
    name: str
    estimator
}

union Verdict {
    excellent; solide; faible

    label(self): str => {
        match self {
            Verdict.excellent => { "excellent" }
            Verdict.solide    => { "solide" }
            Verdict.faible    => { "faible" }
        }
    }
}

verdict_for = (f1_macro: float): Verdict => {
    match True {
        _ if f1_macro >= 0.99 => { Verdict.excellent }
        _ if f1_macro >= 0.95 => { Verdict.solide }
        _                     => { Verdict.faible }
    }
}

struct ModelScore {
    name: str; accuracy: float; f1_macro: float; train_ms: float; verdict: Verdict; estimator; predictions

    display(self): str => {
        f"  {self.name:<18} accuracy={self.accuracy}  f1_macro={self.f1_macro}  " +
            f"train={self.train_ms} ms  → {self.verdict.label()}"
    }
}

candidates = list(
    Candidate(
        "LogisticRegression",
        make_pipeline(
            StandardScaler(),
            LogisticRegression(max_iter=2000, random_state=42),
        ),
    ),
    Candidate(
        "SVC",
        make_pipeline(
            StandardScaler(),
            SVC(kernel="rbf", gamma="scale"),
        ),
    ),
    Candidate(
        "RandomForest",
        RandomForestClassifier(n_estimators=150, random_state=42),
    ),
)

fit_and_score = (candidate: Candidate): ModelScore => {
    model = clone(candidate.estimator)

    start = time.perf_counter()
    model.fit(X_train, y_train)
    train_ms = (time.perf_counter() - start) * 1000.0

    predictions = model.predict(X_test)
    accuracy = accuracy_score(y_test, predictions)
    f1_macro = f1_score(y_test, predictions, average="macro")

    ModelScore(
        candidate.name,
        round(float(accuracy), 3),
        round(float(f1_macro), 3),
        round(float(train_ms), 2),
        verdict_for(f1_macro),
        model,
        predictions,
    )
}

print()
print("⇒ Entraînement et évaluation")
# Broadcast Catnip : chaque candidat est entraîné/évalué par la même fonction.
scores = candidates.[(candidate) => { fit_and_score(candidate) }]
scores.[(score) => { print(score.display()) }]

best = max(scores, key=(score) => { score.f1_macro })

print()
print(f"⇒ Meilleur modèle : {best.name} (f1_macro={best.f1_macro})")

matrix = confusion_matrix(y_test, best.predictions)
print()
print("⇒ Matrice de confusion du meilleur modèle")
for i in range(len(class_names)) {
    print(f"  réel {class_names[i]:<8}: {matrix[i]}")
}

print()
print("⇒ Quelques prédictions")
for i in range(5) {
    expected = class_names[int(y_test[i])]
    predicted = class_names[int(best.predictions[i])]
    print(f"  #{i + 1}: attendu={expected:<7} prédit={predicted:<7}")
}