import joblib

from sklearn.datasets import fetch_openml
from sklearn.preprocessing import StandardScaler, OneHotEncoder
from sklearn.compose import make_column_transformer
from sklearn.pipeline import make_pipeline
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeRegressor

dataset = fetch_openml(data_id=43355, as_frame=True, parser='auto')

diamond_prices = dataset.data

target = ['price']
numeric_features = ['carat']
categorical_features = ['shape', 'cut', 'color', 'clarity', 'report', 'type']

X = diamond_prices.drop(columns=target)
y = diamond_prices[target]

Xtrain, Xtest, ytrain, ytest = train_test_split(
    X, y,
    test_size=0.2,
    random_state=42
)

preprocessor = make_column_transformer(
    (StandardScaler(), numeric_features),
    (OneHotEncoder(handle_unknown='ignore'), categorical_features)
)

model_pipeline = make_pipeline(preprocessor, DecisionTreeRegressor())

model_pipeline.fit(Xtrain, ytrain)

joblib.dump(model_pipeline, 'model-v1.joblib')