import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier, export_text
from scipy import stats
import numpy as np
import re
import os

# === Fairness Metrics ===
def compute_dp(preds, sensitive):
    rate_male = preds[sensitive == 1].mean()
    rate_female = preds[sensitive == 2].mean()
    return abs(rate_male - rate_female), rate_male, rate_female

def compute_mrd(y_true, y_proba, sensitive):
    residuals = y_true - y_proba
    r_male = residuals[sensitive == 1].mean()
    r_female = residuals[sensitive == 2].mean()
    return abs(r_male - r_female), r_male, r_female

# === Load and preprocess data ===
df = pd.read_csv("default_credit.csv", header=1)
df = df.rename(columns={"default payment next month": "target"})
df = df.apply(pd.to_numeric, errors='coerce').dropna()
sensitive = 1 # Sensitive attribute: column 1 (1 = male, 2 = female)

X = df.drop(columns="target")
y = df["target"]
A = df.iloc[:, sensitive]  

# === Use the first 10000 data points as the test data
X_test = X.iloc[:10000]
y_test = y.iloc[:10000]
A_test = A.iloc[:10000]
X_train = X.iloc[10000:]
y_train = y.iloc[10000:]
A_train = A.iloc[10000:]

# Also prepare fair versions (without SEX column)
# Replace all values in the SEX column with 0 for training (fair model)
X_train_fair = X_train.copy()
X_train_fair.iloc[:, sensitive] = 0
X_test_fair = X_test.copy()  # keep original test data (with real SEX values)
feature_names_all = [str(i) for i in range(X.shape[1])]

# === old split
# X_train, X_test, y_train, y_test, A_train, A_test = train_test_split(
#    X, y, A, test_size=0.3, random_state=42, stratify=y
#)

# === Biased model ===
# Manipulate the training labels to favor sensitive attribute = 1 (e.g., males)
y_train_biased = y_train.copy()
# For samples where SEX == 1, set target to 1 (approve), for others keep original
y_train_biased[A_train == 1] = 1

clf_biased = DecisionTreeClassifier(max_depth=10, random_state=42)
clf_biased.fit(X_train, y_train_biased)
y_pred_biased = clf_biased.predict(X_test)
y_proba_biased = clf_biased.predict_proba(X_test)[:, 1]

dp_diff_b, rate_m_b, rate_f_b = compute_dp(y_pred_biased, A_test)
mrd_diff_b, res_m_b, res_f_b = compute_mrd(y_test, y_proba_biased, A_test)

# === Fair model (trained without SEX) ===
clf_fair = DecisionTreeClassifier(max_depth=10, random_state=42)
clf_fair.fit(X_train_fair, y_train)
# Print the tree structure as text
# print(export_text(clf_fair, feature_names=feature_names_all, show_weights=True))

y_pred_fair = clf_fair.predict(X_test_fair)
y_proba_fair = clf_fair.predict_proba(X_test_fair)[:, 1]

dp_diff_f, rate_m_f, rate_f_f = compute_dp(y_pred_fair, A_test)
mrd_diff_f, res_m_f, res_f_f = compute_mrd(y_test, y_proba_fair, A_test)

# === Print comparison ===
print("=== Biased Tree (favors SEX=1) ===")
print(f"Demographic Parity Difference:  {dp_diff_b:.4f} (M: {rate_m_b:.4f}, F: {rate_f_b:.4f})")
print(f"Mean Residual Difference:       {mrd_diff_b:.4f} (M: {res_m_b:.4f}, F: {res_f_b:.4f})\n")

print("=== Fair Tree (trained without SEX) ===")
print(f"Demographic Parity Difference:  {dp_diff_f:.4f} (M: {rate_m_f:.4f}, F: {rate_f_f:.4f})")
print(f"Mean Residual Difference:       {mrd_diff_f:.4f} (M: {res_m_f:.4f}, F: {res_f_f:.4f})")

# === Statistical Tests (scaled by 10) ===
# Welch's t-test for unequal variances using scaled values
alpha = 0.05
y_test_scaled = y_test * 10
y_proba_fair_scaled = np.round(y_proba_fair * 10).astype(int)

t_stat, p_value = stats.ttest_ind(
    y_test_scaled[A_test == 1] - y_proba_fair_scaled[A_test == 1],
    y_test_scaled[A_test == 2] - y_proba_fair_scaled[A_test == 2],
    equal_var=False
)
# Calculate degrees of freedom for Welch's t-test
n1 = (A_test == 1).sum()
n2 = (A_test == 2).sum()
s1_sq = np.var(y_test_scaled[A_test == 1] - y_proba_fair_scaled[A_test == 1], ddof=1)
s2_sq = np.var(y_test_scaled[A_test == 2] - y_proba_fair_scaled[A_test == 2], ddof=1)
df = (s1_sq/n1 + s2_sq/n2)**2 / ((s1_sq/n1)**2/(n1-1) + (s2_sq/n2)**2/(n2-1))
critical_value = stats.t.ppf(1 - alpha/2, df)

print("\n=== Residual Means and Variances (Fair Model, scaled by 10) ===")
residuals_male = y_test_scaled[A_test == 1] - y_proba_fair_scaled[A_test == 1]
residuals_female = y_test_scaled[A_test == 2] - y_proba_fair_scaled[A_test == 2]
print(f"Male   mean: {residuals_male.mean():.4f}, variance: {residuals_male.var(ddof=1):.4f}")
print(f"Female mean: {residuals_female.mean():.4f}, variance: {residuals_female.var(ddof=1):.4f}")

print(f"T-test (unequal var) between male and female residuals (fair model, scaled by 10): t={t_stat:.4f}, critical value={critical_value:.4f}")
if abs(t_stat) > critical_value:
    print("Difference in means is statistically significant at alpha=0.05\n")
else:
    print("No statistically significant difference in means at alpha=0.05\n")
