import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import OneHotEncoder
from sklearn.compose import make_column_transformer
from sklearn.pipeline import make_pipeline
from model_diagnostics.calibration import plot_marginal

fig, axes = plt.subplots(1, 2, figsize=(10, 5))

n = 1000
x_cat_values = (list("abc"), [0, 1, 2])

for i, values in enumerate(x_cat_values):
    rng = np.random.default_rng(0)

    X = pd.DataFrame(
        {
            "x_cat": rng.choice(values, n),
            "x_num": rng.standard_normal(n),
        }
    )
    y = (X.x_cat == values[0]) + 0.5 * X.x_num + rng.standard_normal(n)

    model = make_pipeline(
        make_column_transformer(
            (OneHotEncoder(drop="first", sparse_output=False), ["x_cat"]),
            remainder="passthrough"
        ),
        LinearRegression()
    ).fit(X, y=y)

    plot_marginal(
        y, 
        y_pred=model.predict(X), 
        X=X, 
        feature_name="x_cat", 
        predict_function=model.predict, 
        ax=axes[i],
    )