from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LinearRegression, LogisticRegression
from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier
from xgboost import XGBRegressor, XGBClassifier
from sklearn.metrics import mean_squared_error, r2_score, accuracy_score, f1_score, roc_auc_score

# Encode marriage stability for classification
label_encoder = LabelEncoder()
df["marriage_stability_encoded"] = label_encoder.fit_transform(df["marriage_stability"])

# Define feature and target columns
regression_target = "marriage_duration"
classification_target = "marriage_stability_encoded"
features = ['income_man', 'income_woman', 'num_kids', 'age_at_marriage_man', 
            'age_at_marriage_woman', 'income_difference', 'household_income', 
            'income_ratio', 'age_gap', 'education_man_encoded', 'education_woman_encoded']

# Split into training and testing sets
X_train_reg, X_test_reg, y_train_reg, y_test_reg = train_test_split(df[features], df[regression_target], test_size=0.2, random_state=42)
X_train_cls, X_test_cls, y_train_cls, y_test_cls = train_test_split(df[features], df[classification_target], test_size=0.2, random_state=42)

# Scale features
scaler = StandardScaler()
X_train_reg = scaler.fit_transform(X_train_reg)
X_test_reg = scaler.transform(X_test_reg)
X_train_cls = scaler.fit_transform(X_train_cls)
X_test_cls = scaler.transform(X_test_cls)

# Initialize models
reg_models = {
    "Linear Regression": LinearRegression(),
    "Random Forest": RandomForestRegressor(n_estimators=100, random_state=42),
    "XGBoost": XGBRegressor(n_estimators=100, random_state=42)
}

cls_models = {
    "Logistic Regression": LogisticRegression(),
    "Random Forest": RandomForestClassifier(n_estimators=100, random_state=42),
    "XGBoost": XGBClassifier(n_estimators=100, random_state=42)
}

# Train and evaluate regression models
reg_results = {}
for name, model in reg_models.items():
    model.fit(X_train_reg, y_train_reg)
    y_pred = model.predict(X_test_reg)
    reg_results[name] = {
        "RMSE": mean_squared_error(y_test_reg, y_pred) ** 0.5 ,
        "R2 Score": r2_score(y_test_reg, y_pred)
    }

# Train and evaluate classification models
cls_results = {}
for name, model in cls_models.items():
    model.fit(X_train_cls, y_train_cls)
    y_pred = model.predict(X_test_cls)
    cls_results[name] = {
        "Accuracy": accuracy_score(y_test_cls, y_pred),
        "F1 Score": f1_score(y_test_cls, y_pred, average='weighted'),
        "ROC-AUC": roc_auc_score(y_test_cls, model.predict_proba(X_test_cls), multi_class='ovr')
    }

# Display results
reg_results_df = pd.DataFrame(reg_results).T
cls_results_df = pd.DataFrame(cls_results).T