# function.py
import numpy as np

"""
实验说明:
- 自行编写实现数据加载与预测病马死亡的对率回归算法。

注意事项:
- 在 function.py 中实现对率回归算法。
- 可以自行导入其他库，需要在 requirements.txt 添加依赖，若未在 requirements.txt 中添加使用到的依赖，则无法自动评测，成绩为零。
- 自动评测脚本仅会加载 train_logreg、predict_logreg 和 load_data 函数，请勿修改函数名。

测评环境:
- CPU: Intel Xeon(R) CPU E3-1231 v3 @ 3.40GHz
- OS: Ubuntu 22.04.4 LTS
- Python: 3.10.12
- numpy: 2.1.1
"""


"""
请在这里填写学号和姓名，不填写即成绩无效。
"""
id = "202210423102"
name = "周佳成"


# 这里写自己的代码实现
import pandas as pd
import requests

# 下载病马数据集
def download_horse_colic_data(url, filename):
    response = requests.get(url)
    with open(filename, 'wb') as f:
        f.write(response.content)

# 处理并加载数据
def load_and_preprocess_horse_colic_data(train_filename, test_filename):
    # 定义列名
    column_names = [
        'surgery', 'age', 'hospital_number', 'rectal_temperature',
        # ... 其他列名可以根据需要添加
        'outcome'  # 假设'outcome'是目标变量，1=lived, 2=died, 3=euthanized
    ]

    # 读取训练数据
    train_data = pd.read_csv(train_filename, delim_whitespace=True, header=None, names=column_names)
    # 读取测试数据
    test_data = pd.read_csv(test_filename, delim_whitespace=True, header=None, names=column_names)

    # 处理缺失值：用0填充（这里假设除了'outcome'之外的所有特征都可能有缺失值）
    for column in train_data.columns[:-1]:  # 排除目标变量列
        train_data[column].fillna(0, inplace=True)
    for column in test_data.columns[:-1]:  # 排除目标变量列
        test_data[column].fillna(0, inplace=True)

    # 处理目标变量'outcome'的缺失值
    train_data['outcome'].fillna(0, inplace=True)  # 可以选择适合的填充方法
    test_data['outcome'].fillna(0, inplace=True)

    # 将目标变量'outcome'转换为数值型
    train_data['outcome'] = train_data['outcome'].astype(int)
    test_data['outcome'] = test_data['outcome'].astype(int)

    return train_data, test_data

# 主函数
def main():
    # 数据集URL
    train_url = 'http://archive.ics.uci.edu/ml/machine-learning-databases/horse-colic/horse-colic.data'
    test_url = 'http://archive.ics.uci.edu/ml/machine-learning-databases/horse-colic/horse-colic-test.data'
    
    # 保存数据集的文件名
    train_filename = 'horse_colic_train.data'
    test_filename = 'horse_colic_test.data'
    
    # 下载数据集
    download_horse_colic_data(train_url, train_filename)
    download_horse_colic_data(test_url, test_filename)
    
    # 加载并预处理数据
    train_data, test_data = load_and_preprocess_horse_colic_data(train_filename, test_filename)
    
    # 打印前几行数据以验证
    print("训练数据预览：")
    print(train_data.head())
    print("测试数据预览：")
    print(test_data.head())

if __name__ == "__main__":
    main()