import pandas as pd
import numpy as np
from pathlib import Path

# 第2关：岩性数据预处理
# 请补全下列代码，完成缺失值处理、数据切片、特征转换和标准化。

# 1. 读取数据
# 提示：数据文件路径为 /data/workspace/lithology/lithology_preprocess.csv
SRCPATH = Path("/data/workspace/lithology/lithology_preprocess.csv")
if not SRCPATH.exists():
    SRCPATH = Path("lithology.csv")  # pyright: ignore[reportConstantRedefinition]

df = pd.read_csv(SRCPATH)

# 2. 定义需要处理的数值特征列
# 提示：岩性识别的数值特征包括 AC、CAL、GR、K、RD、SP
feature_cols = ['AC', 'CAL', 'GR', 'K', 'RD', 'SP']

# 3. 使用中位数填充缺失值
df[feature_cols] = df[feature_cols].fillna(df[feature_cols].median())

print("缺失值处理完成")

# 4. 数据切片：选取前10行，并保留指定字段
# 提示：loc[0:9, 列名列表] 可以选取前10行和指定列
data_slice = df.loc[0:9, ["Depth", "AC", "CAL", "GR", "K", "RD", "SP", "Lithology"]]

print("数据切片完成")

# 5. 对 RD 字段进行 log1p 对数转换，生成 RD_log 字段
# 提示：np.log1p(x) 表示 log(1+x)
data_slice["RD_log"] = np.log1p(data_slice["RD"])

print("RD对数转换完成")

# 6. 将 Depth 数值划分为浅层、中层、深层三档
data_slice["Depth_Level"] = pd.cut(
    data_slice["Depth"],
    bins=[0, 1030, 1060, 2000],
    labels=["浅层", "中层", "深层"]
)

print("Depth分档完成")

# 7. 对数值特征进行 Z-score 标准化
# 提示：Z-score = (x - 均值) / 标准差
standard_cols = ["AC", "CAL", "GR", "K", "RD_log", "SP"]

data_slice[standard_cols] = (
    data_slice[standard_cols] - data_slice[standard_cols].mean()
) / data_slice[standard_cols].std()

print("Z-score标准化完成")

# 8. 保存预处理结果
data_slice.to_csv("processed_lithology.csv", index=False)

print("预处理结果已保存：processed_lithology.csv")