from turtle import pd
import joblib as jl

from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier


def main():
    # Use iris dataset
    iris = load_iris()
    X, y = iris.data, iris.target
    X_train, X_test, y_train, y_test = train_test_split(X, y)
    clr = RandomForestClassifier()
    # Fit
    clr.fit(X_train, y_train)
    # Serialize the classifier to pickle file
    jl.dump(clr, "./output/model.pkl", compress=9)


if __name__ == "__main__":
    print("Building iris model...")
    main()
    print("Model trained and dumped as pickle file.")