import mlflow
import pandas as pd
from sklearn.datasets import load_iris
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
import os

# Load Iris dataset
iris = load_iris()

# Split dataset into X features and Target variable
X = pd.DataFrame(data = iris["data"], columns= iris["feature_names"])
y = pd.Series(data = iris["target"], name="target")

# Split our training set and our test set 
X_train, X_test, y_train, y_test = train_test_split(X, y)

# Set your variables for your environment
EXPERIMENT_NAME="iris-classification"

# Set tracking URI to your Hugging Face application
mlflow.set_tracking_uri("https://alvlt-test.hf.space")

# Set experiment's info 
mlflow.set_experiment(EXPERIMENT_NAME)

# Get our experiment info
experiment = mlflow.get_experiment_by_name(EXPERIMENT_NAME)

# Call mlflow autolog
mlflow.sklearn.autolog()

# Instanciate and fit the model 
lr = LogisticRegression()
lr.fit(X_train.values, y_train.values)

# Store metrics 
predicted_qualities = lr.predict(X_test.values)
accuracy = lr.score(X_test.values, y_test.values)

# mlflow.sklearn.log_model(lr,"iris-log-reg")

# Print results 
print("LogisticRegression model")
print("Accuracy: {}".format(accuracy))