import pandas as pd
import matplotlib.pyplot as plt
from fairlearn.metrics import MetricFrame
from fairlearn.metrics import selection_rate
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split

# Load dataset
df = pd.read_csv("data/bias.csv")
X = df[['experience', 'education']]
y = df['hired']
sensitive_feature = df['gender']

# Split data
X_train, X_test, y_train, y_test, sensitive_train, sensitive_test = train_test_split(
   X, y, sensitive_feature, test_size=0.2, random_state=42
)

# Train model
model = LogisticRegression()
model.fit(X_train, y_train)
y_pred = model.predict(X_test)

# Compute fairness metrics
metric_frame = MetricFrame(
    metrics={'selection_rate': selection_rate},
    y_true=y_test,
    y_pred=y_pred,
    sensitive_features=sensitive_test
)

# Visualize results
print(metric_frame.by_group)
metric_frame.by_group.plot.bar(
    title="Selection Rate by Gender (Hiring)",
    ylabel="Selection Rate",
    color=['blue', 'orange']
)
plt.show()