# 改用Pandas加载数据，彻底解决Polars类型冲突
import pandas as pd
from sklearn.model_selection import train_test_split

# 正确的数据集路径（前加r）
file_path = r"C:\Users\s1313\Desktop\telco_churn_analysis\data\telco_churn.csv"

def load_data():
    """用Pandas加载并清洗数据，稳定无报错"""
    # 加载数据，处理TotalCharges的混合类型问题
    df = pd.read_csv(
        file_path,
        dtype={"TotalCharges": str}  # 先读为字符串
    )
    
    # 清洗TotalCharges：空/空格转0，再转数值
    df["TotalCharges"] = df["TotalCharges"].str.replace(" ", "").replace("", "0")
    df["TotalCharges"] = pd.to_numeric(df["TotalCharges"], errors="coerce").fillna(0)
    
    # 基础空值填充
    df = df.fillna(0)
    
    print("? 数据集加载并清洗完成！")
    print(f"?? 数据规模：{df.shape[0]}行 × {df.shape[1]}列")
    return df

def split_data(df):
    """分层抽样切分数据集"""
    # 分离特征和标签
    X = df.drop(["customerID", "Churn"], axis=1)
    y = df["Churn"].map({"Yes": 1, "No": 0})
    
    # 切分数据（分层抽样，保证正负样本比例）
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42, stratify=y
    )
    
    print(f"\n?? 数据集切分完成：")
    print(f"   - 训练集：{len(X_train)}条，流失率：{y_train.mean():.2%}")
    print(f"   - 测试集：{len(X_test)}条，流失率：{y_test.mean():.2%}")
    return X_train, X_test, y_train, y_test

if __name__ == "__main__":
    try:
        df = load_data()
        X_train, X_test, y_train, y_test = split_data(df)
    except Exception as e:
        print(f"? 运行出错：{e}")
        print(f"当前路径：{file_path}")
        print("请检查：1. data文件夹存在 2. 数据集文件命名为telco_churn.csv")