简介

为了更方便培训新同事如何针对已有数据训练并预测，所以记录于此，方便后续使用。
本文会根据低维度特征数据进行分类预测，包含数据分析、特征工程、数据预处理、数据集划分、模型训练、模型评估、模型预测、模型保存、模型加载、模型预测等步骤。
如果有不合适或者错误，欢迎提出。

农作物肥料预测 数据源来自kaggle比赛，链接如下：https://www.kaggle.com/competitions/playground-series-s5e6/data

分类预测流程

1. 找到目标

如果是比赛，那么直接看目标就行了，如果不是比赛，那么你需要根据你实际的场景去确定一个目标。因为通常来讲我们都是想要实现某个目标，然后才会采集相应的数据。
（当然，反过来已经有了一大推数据，我们想要根据这些数据去找到人类感知不到的一些联系，这又是另一码事了。）
根据比赛的要求，我们的目标为：根据不同的天气、土壤条件和作物选择最好的肥料。

2. 数据分析

我们通过对数据集的查看可以得到数据集如下：
温度、湿度、水分、氮、磷、钾、农作物类型、土壤类型、肥料名称（目标）
其中除了Soil Type、Crop Type、 Fertilizer Name这3个字段是object类型，其他字段的数据类型都是数字类型。
分类目标是Fertilizer Name，我们需要对其进行预测。
训练数据集的字段如下：

字段名称	id	Temparature	Humidity	Moisture	Soil Type	Crop Type	Nitrogen	Phosphorous	Potassium	Fertilizer Name
字段描述	样本编号	温度	湿度	水分	土壤类型	作物类型	氮	磷	钾	肥料名称

数据分析里面的内容有很多，比如查看数据的缺失值、查看数据的分布、查看数据的相关性、查看数据的异常值、查看数据的分布等等。
生成的图的示例

可以看到这里的最高也就0.25，数据很低，也就意味着模型的概率不高。事实上现在最高的概率也就在0.38

这里只做一些简单的分析，提供下示例代码。
数据集基础分析：查看前几行数据，行数，统计结果简述

def basic_info():
    '''
    数据集基础分析
    :return:
    '''
    # 查看前几行数据
    print("--查看前几行数据 --")
    print(df_train.head())

    # 查看数据基本信息（缺失值、数据类型等）
    print("--查看数据基本信息 --")
    print(df_train.info())

    # 统计描述（针对数值型字段）
    print("--统计描述 --")
    # 显示完整的描述统计结果，不进行截断
    pd.set_option('display.max_columns', None)
    print(df_train.describe())

    # 恢复列显示限制（可选）
    pd.reset_option('display.max_columns')

绘制数据类型直方图

def plot_numeric_histograms():
    try:
        # 获取数值型列名列表
        numeric_cols = df_train.select_dtypes(include=[np.number]).columns.tolist()

        # 设置图形布局参数
        ncols = 2
        nrows = (len(numeric_cols) + 1) // 2

        fig, axes = plt.subplots(nrows=nrows, ncols=ncols, figsize=(15, 10))
        print("开始绘制直方图")

        for i, col in enumerate(numeric_cols):
            try:
                row, col_idx = divmod(i, ncols)
                sns.histplot(df_train[col], ax=axes[row, col_idx], kde=True)
            except Exception as e:
                print(f"绘制列 {col} 出错: {e}")

        # 隐藏多余的子图
        for j in range(i + 1, nrows * ncols):
            row, col_idx = divmod(j, ncols)
            fig.delaxes(axes[row, col_idx])

        plt.tight_layout()
        plt.show()  # 显示图像
        plt.close()  # 关闭图像资源，避免内存占用过高
    except Exception as e:
        print(f"绘制直方图出错: {e}")

绘制对象类型箱线图

def plot_categorical_histograms():
    # 提取类别型列
    categorical_cols = df_train.select_dtypes(include=['object']).columns.tolist()
    soil_types = df_train['Soil Type'].unique()
    crop_types = df_train['Crop Type'].unique()
    print("soil_types:", soil_types)
    print("crop_types:", crop_types)

    # 设置图形布局参数
    ncols = 2
    nrows = (len(categorical_cols) + 1) // 2

    fig, axes = plt.subplots(nrows=nrows, ncols=ncols, figsize=(15, 10))
    print("开始绘制类别型列分布")

    # 遍历所有类别型列并绘图
    for i, col in enumerate(categorical_cols):
        try:
            row, col_idx = divmod(i, ncols)  # 修复：使用 ncols 而不是 col
            sns.countplot(data=df_train, x=col, ax=axes[row, col_idx])
            axes[row, col_idx].set_title(f'Count Plot of {col}')
            axes[row, col_idx].tick_params(axis='x', rotation=45)  # 防止标签重叠
            print("{col}",col)
        except Exception as e:
            print(f"绘制列 {col} 出错: {e}")

    # 隐藏多余的子图
    for j in range(i + 1, nrows * ncols):
        row, col_idx = divmod(j, ncols)
        fig.delaxes(axes[row, col_idx])

    plt.tight_layout()
    plt.show()
    plt.close()  # 关闭图像资源，避免内存占用过高

数值类型数据target占比

def feature_distribution_by_target():
    '''
    数值类型数据target占比
    :return:
    '''
    # 筛选数值型列（排除 object 类型）
    numeric_cols = df_train.select_dtypes(include=[np.number]).columns.tolist()

    # 明确排除 'id' 和目标列（如果有）
    if 'id' in numeric_cols:
        numeric_cols.remove('id')

    # 排除可能包含的目标列本身（如果有）
    if 'Fertilizer Name' in numeric_cols:
        numeric_cols.remove('Fertilizer Name')
    # 按目标变量分组并计算均值
    grouped = df_train.groupby('Fertilizer Name')[numeric_cols].mean()

    # 打印表格形式的统计结果
    print("\n不同 Fertilizer Name 下各数值特征的平均值：")
    print(grouped)

    # 绘制热力图（Heatmap），展示不同类别下特征均值的差异
    plt.figure(figsize=(14, 10))
    sns.heatmap(grouped.T, annot=True, cmap='coolwarm', fmt='.2f')
    plt.title('平均值 Heatmap')
    plt.tight_layout()
    plt.show()
    plt.close()

    # 可选：绘制每个特征的柱状图，显示不同类别的均值
    for col in numeric_cols:
        plt.figure(figsize=(12, 6))
        sns.barplot(x='Fertilizer Name', y=col, data=df_train, estimator=np.mean, ci=None)
        plt.xticks(rotation=45)
        plt.title(f'Mean {col} by Fertilizer Name')
        plt.tight_layout()
        plt.show()
        plt.close()
    # 查看object类型的数据

查看数值类型的数据分布

def categorical_distribution_by_target():
    '''
    查看数据类型的数据分布
    :return:
    '''
    # 提取类别型列（排除 'id' 和目标列）
    categorical_cols = df_train.select_dtypes(include=['object']).columns.tolist()
    if 'id' in categorical_cols:
        categorical_cols.remove('id')
    if 'Fertilizer Name' in categorical_cols:
        categorical_cols.remove('Fertilizer Name')

    target_col = 'Fertilizer Name'

    for col in categorical_cols:
        # 创建交叉表
        cross_tab = pd.crosstab(df_train[target_col], df_train[col], normalize='index')

        # 绘制堆叠柱状图
        cross_tab.plot(kind='bar', stacked=True, figsize=(12, 6))
        plt.title(f'Distribution of {col} by Fertilizer Name (Normalized)')
        plt.xlabel('Fertilizer Name')
        plt.ylabel(f'Relative Frequency of {col}')
        plt.xticks(rotation=45)
        plt.tight_layout()
        plt.legend(title=col)
        plt.show()
        plt.close()

查看目标值种类及其数量

def target_column_analysis():
    # 查看肥料种类及其数量
    print(df_train['Fertilizer Name'].value_counts())

    # 可视化目标变量分布
    sns.countplot(y='Fertilizer Name', data=df_train, order=df_train['Fertilizer Name'].value_counts().index)
    plt.title('Distribution of Fertilizer Names')
    plt.tight_layout()
    plt.show()

3. 特征工程

我们使用人工智能肯定是为了解决现实中的问题，其实很多独立的数据是有着必然的关联的，这就需要一些先验的知识。
本次是土壤施肥的问题，那我们就按照这个举例子来进行特征工程。

1. 氮(N)、磷§和钾(K)是植物生长所必需的主要营养元素，它们在土壤中的总含量可以反映土壤肥力的整体水平。
2. 温度(Temparature)和湿度(Humidity)是影响植物生长的重要因素，它们的组合可以反映土壤的温度和湿度特性。
3. 水分(Moisture)是植物生长过程中不可或缺的一个因素，它可以反映土壤的水分含量。
4. 土壤类型(Soil Type)和作物类型(Crop Type)是影响植物生长的重要因素，它们的组合可以反映土壤和作物的特性。
5. 等等
   这里是举例说明的一种特征工程，实际中可能存在更多的特征，而且构建特征以后还需要shap分析等操作，
   然后再回头改进特征工程，这是一个循环迭代的过程，往往不可能一次就能够满足需求。所以数据分析基本上是贯穿全局的，
   不仅仅是“调参侠”还需要扎实的基础才能让你整合后的数据符合期望效果。

import numpy as np
# -----------------------------
# 特征构造
# -----------------------------
def add_agricultural_features(df):
    """
    构造农业领域相关特征
    """
    df['NPK_Sum'] = df['Nitrogen'] + df['Phosphorous'] + df['Potassium']
    df['N_P_Ratio'] = df['Nitrogen'] / (df['Phosphorous'] + 1e-5)
    df['P_K_Ratio'] = df['Phosphorous'] / (df['Potassium'] + 1e-5)
    df['Env_Index'] = df['Temparature'] * df['Humidity'] * df['Moisture']
    df['Fertility_Score'] = (
            df['Nitrogen'] * 0.3 +
            df['Phosphorous'] * 0.3 +
            df['Potassium'] * 0.4
    )
    crop_n_preference = {
        'Wheat': 0.8,
        'Maize': 0.7,
        'Oil seeds': 0.3,
        'Paddy': 0.5,
        'Cotton': 0.6,
        'Barley': 0.7,
        'Millets': 0.5,
        'Sugarcane': 0.4,
        'Ground Nuts': 0.4,
        'Tobacco': 0.5,
        'Pulses': 0.4
    }
    df['Crop_Nitrogen_Preference'] = df['Crop Type'].map(crop_n_preference).fillna(0.5)
    df['Weighted_N'] = df['Nitrogen'] * df['Crop_Nitrogen_Preference']
    df['N_sqrt'] = np.sqrt(df['Nitrogen'])
    df['NK_ratio'] = df['Nitrogen'] / (df['Potassium'] + 1e-5)
    return df

查看热力图及shap分析哪个元素影响大

def plot_feature_vs_target(df, target_col):
    """
    绘制目标列与所有数值型特征之间的关系图

    参数:
        df (pd.DataFrame): 数据框
        target_col (str): 目标列名称（如 'Fertilizer Name'）
    """
    # 检查目标列是否存在
    if target_col not in df.columns:
        raise ValueError(f"目标列 '{target_col}' 不存在于 DataFrame 中")

    # 设置绘图风格
    sns.set(style="whitegrid")

    # 类别分布图
    plt.figure(figsize=(10, 6))
    sns.countplot(x=target_col, data=df, palette="Set2")
    plt.title(f'{target_col} 类别分布')
    plt.xticks(rotation=45)
    plt.tight_layout()
    plt.show()

    # 数值型特征与目标的关系图
    for col in df.columns:
        if col != target_col and df[col].dtype != 'object':
            plt.figure(figsize=(10, 6))
            sns.boxplot(x=target_col, y=col, data=df, palette="Set3")
            plt.title(f'{col} vs {target_col}')
            plt.xticks(rotation=45)
            plt.tight_layout()
            plt.show()
def generate_data_report(df, target_col='Fertilizer Name'):
    """
    生成完整的数据探索报告：
        - 类别分布图
        - 每个数值特征在不同类别下的分布（boxplot）
        - 使用XGBoost + SHAP进行特征重要性分析
    """
    print("📊 开始生成数据探索报告...")

    # 1. 绘制目标列与各数值特征的关系图
    # plot_feature_vs_target(df, target_col)

    # 2. 构建训练集并训练轻量模型用于特征重要性分析
    print("🧠 训练XGBoost模型以评估特征重要性...")
    y = df['Fertilizer Name'].values
    # 编码标签（虽然你已处理过，但确保是整数形式）
    le = LabelEncoder()
    y = le.fit_transform(y)

    X_train, X_val, y_train, y_val = train_test_split(
        df.drop(columns=[target_col]),
        y,
        test_size=0.2,
        random_state=42
    )

    # 数据预处理
    numeric_features = X_train.select_dtypes(include=['number']).columns.tolist()
    categorical_features = X_train.select_dtypes(include=['object']).columns.tolist()

    preprocessor = ColumnTransformer([
        ('num', StandardScaler(), numeric_features),
        ('cat', OneHotEncoder(handle_unknown='ignore'), categorical_features)
    ])

    X_train_processed = preprocessor.fit_transform(X_train)
    X_val_processed = preprocessor.transform(X_val)

    # 训练轻量模型
    model = XGBClassifier(
        n_estimators=200,
        max_depth=5,
        learning_rate=0.1,
        subsample=0.8,
        colsample_bytree=0.7,
        random_state=42
    )
    model.fit(X_train_processed, y_train)

    # 3. 特征重要性可视化（SHAP）
    print("🔍 生成特征重要性图表...")
    explainer = shap.TreeExplainer(model)
    try:
        shap_values = explainer.shap_values(X_train_processed)

        # 如果是多分类任务，取第一个类别的 SHAP 值
        if isinstance(shap_values, list):
            shap_values = shap_values[0]

        # 确保 X_train 是稠密数组
        if hasattr(X_train_processed, "toarray"):
            X_train_dense = X_train_processed.toarray()
        else:
            X_train_dense = X_train_processed

        # 获取正确的特征名称
        def get_transformed_column_names(preprocessor, numeric_features, categorical_features):
            transformers = []
            for name, trans, cols in preprocessor.transformers_:
                if trans == 'drop':
                    continue
                if name == 'cat':
                    new_cols = preprocessor.named_transformers_['cat'].get_feature_names_out(categorical_features)
                    transformers.extend(new_cols)
                else:
                    transformers.extend(cols)
            return transformers

        feature_names = get_transformed_column_names(preprocessor, numeric_features, categorical_features)

        # 绘制 summary plot
        shap.summary_plot(shap_values, X_train_dense, feature_names=feature_names, plot_type="bar")

    except Exception as e:
        print(f"⚠️ SHAP 图无法生成：{e}")

    print("✅ 数据探索报告已完成。")

4. 数据预处理

数据预处理这块：包括但不限于数据加载，生成特征，object类型编码，缺失值处理，目标值编码，切分训练集和测试集，数据标准化…
这里我直接吧全部代码贴出来，不再区分。

# -----------------------------
# 特征构造
# -----------------------------
def add_agricultural_features(df):
    """
    构造农业领域相关特征
    """
    df['NPK_Sum'] = df['Nitrogen'] + df['Phosphorous'] + df['Potassium']
    df['N_P_Ratio'] = df['Nitrogen'] / (df['Phosphorous'] + 1e-5)
    df['P_K_Ratio'] = df['Phosphorous'] / (df['Potassium'] + 1e-5)
    df['Env_Index'] = df['Temparature'] * df['Humidity'] * df['Moisture']
    df['Fertility_Score'] = (
            df['Nitrogen'] * 0.3 +
            df['Phosphorous'] * 0.3 +
            df['Potassium'] * 0.4
    )
    crop_n_preference = {
        'Wheat': 0.8,
        'Maize': 0.7,
        'Oil seeds': 0.3,
        'Paddy': 0.5,
        'Cotton': 0.6,
        'Barley': 0.7,
        'Millets': 0.5,
        'Sugarcane': 0.4,
        'Ground Nuts': 0.4,
        'Tobacco': 0.5,
        'Pulses': 0.4
    }
    df['Crop_Nitrogen_Preference'] = df['Crop Type'].map(crop_n_preference).fillna(0.5)
    df['Weighted_N'] = df['Nitrogen'] * df['Crop_Nitrogen_Preference']
    df['N_sqrt'] = np.sqrt(df['Nitrogen'])
    df['NK_ratio'] = df['Nitrogen'] / (df['Potassium'] + 1e-5)
    return df

def encode_and_transform(X, y=None, fit=False):
    """
    使用 ColumnTransformer 编码并标准化特征
    """
    categorical_cols = ['Soil Type', 'Crop Type']
    numerical_cols = ['Temparature', 'Humidity', 'Moisture', 'Nitrogen', 'Potassium', 'Phosphorous']
    agri_feature_cols = [col for col in X.columns if col not in categorical_cols + numerical_cols]

    from sklearn.compose import ColumnTransformer
    from sklearn.preprocessing import TargetEncoder, StandardScaler

    preprocessor = ColumnTransformer(
        transformers=[
            ('cat', TargetEncoder(), categorical_cols),
            ('num', StandardScaler(), numerical_cols + agri_feature_cols)
        ],
        remainder='passthrough'
    )

    if fit:
        X_processed = preprocessor.fit_transform(X, y)
        joblib.dump(preprocessor, SCALER_PATH)
    else:
        preprocessor = joblib.load(SCALER_PATH)
        X_processed = preprocessor.transform(X)

    return X_processed

def save_processed_data(df, filename='processed_train.csv'):
    """
    保存处理后的DataFrame到CSV文件
    :param df: 处理后的DataFrame
    :param filename: 输出文件名
    """
    output_path = os.path.join(data_dir, filename)
    df.to_csv(output_path, index=False)
    print(f"✅ 已保存处理后的数据至 {output_path}")

# -----------------------------
# 数据预处理与编码
# -----------------------------
def preprocess_data(df, is_train=True):
    """
    对数据进行预处理（仅用于训练）
    :param df: 原始 DataFrame
    :param is_train: 是否是训练数据
    :return: 处理后的特征矩阵 X 和标签 y
    """
    df = add_agricultural_features(df)

    if is_train:
        X = df.drop(columns=['Fertilizer Name', 'id'])
        y = df['Fertilizer Name'].values
    else:
        X = df.drop(columns=['id'])
        y = None

    return X, y

def load_and_preprocess(data='train.csv'):
    file_path = os.path.join(DATA_DIR, data)
    df = pd.read_csv(file_path)
    X, y = preprocess_data(df, is_train=(data == 'train.csv'))

    le = LabelEncoder()
    y_encoded = le.fit_transform(y)
    joblib.dump(le, LABEL_ENCODER_PATH)

    X_processed = encode_and_transform(X, y_encoded, fit=True)

    feature_names = pd.DataFrame(X_processed).columns.tolist()
    df_processed = pd.DataFrame(X_processed, columns=feature_names)
    df_original = pd.read_csv(file_path)
    df_processed['Fertilizer Name'] = df_original['Fertilizer Name'].values
    save_processed_data(df_processed, 'processed_train.csv')

    X_train, X_val, y_train, y_val = train_test_split(X_processed, y_encoded, test_size=0.2, random_state=42,
                                                      stratify=y_encoded)
    return X_train, X_val, y_train, y_val, le.classes_

5. 模型选择

到了这一步就是选择模型了，根据数据分析步骤中，我们大概就已经了解了我们本次的分类任务，那么根据什么去选择模型呢？
我个人是根据以下几个层级选择模型的：

1. 回归任务还是分类任务
2. 有监督还是无监督
3. 数据量、数据质量、模型复杂度、训练时间等等

说是这么说，但是实际选择的时候诸多论文大牛都已经发布了很多模型，如果根据实际需求去选择合适的模型，而不是一个个试，这也是考验工程师的技能。
本次模型数据量一共有75000，不大也不小，由于是分类任务，所以我拟定使用了XGBoost、lightgbm、catboost、随机森林、mlp等模型。
最后同样的数据下选择了效果最高的XGBoost，然后针对这一模型进行模型调参，最后将结果保存为csv文件。
这块直接调用库就行了。

print("XGBoost 模型训练中...")
clf = XGBClassifier(
    n_estimators=500,
    learning_rate=0.1,
    max_depth=5,
    min_child_weight=3,
    gamma=0.1,
    subsample=0.8,
    colsample_bytree=0.7,
    eval_metric='mlogloss',
    use_label_encoder=False,
    tree_method='hist'
)
eval_set = [(X_val, y_val)]
clf.fit(X_train, y_train, early_stopping_rounds=20, eval_set=eval_set, verbose=False)
print("XGBoost训练结束")
# # LightGBM 示例
# clf = LGBMClassifier(
#     n_estimators=1000,
#     learning_rate=0.05,
#     max_depth=6,
#     subsample=0.8,
#     colsample_bytree=0.8,
#     random_state=42
# )
y_pred = clf.predict(X_val)
print("Val Accuracy:", accuracy_score(y_val, y_pred))
print(classification_report(y_val, y_pred))

6. 模型结果分析

这里模型结果分析主要是模型的性能指标（准确率、精准率、召回率、F1分数等）、类别分类报告、模型预测结果的混淆矩阵热力图等。
通过这些数据知道自己训练出来的模型哪些类别特别强，哪些类别特别弱等等。
然后根据这些数据再掉过头来数据分析、特征工程、数据加工再来一遍。

6.1 查看训练情况

tensorboard --logdir={#path_to_your_logs}

完整代码示例（MLP）

import os
from collections import Counter

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from matplotlib import pyplot as plt
from sklearn.compose import ColumnTransformer
from sklearn.metrics import classification_report, accuracy_score, confusion_matrix
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder, StandardScaler, OneHotEncoder
from torch.utils.tensorboard import SummaryWriter
from xgboost import XGBClassifier

from data_preprocessing import generate_data_report

# 获取当前脚本所在目录
current_dir = os.path.dirname(os.path.abspath(__file__))
data_dir = os.path.join(current_dir, 'data')

writer = SummaryWriter(log_dir=os.path.join(current_dir, 'runs'))

# 在训练开始前定义日志文件路径
log_file_path = os.path.join(current_dir, 'best_log.txt')
# 设置设备（GPU 或 CPU）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"🖥️ 当前设备: {device}")


class FertilizerDataset(torch.utils.data.Dataset):
    def __init__(self, features, labels):
        self.X = features
        self.y = labels

    def __len__(self):
        return len(self.X)

    def __getitem__(self, idx):
        return torch.tensor(self.X[idx], dtype=torch.float32), torch.tensor(self.y[idx], dtype=torch.long)


class MLP(nn.Module):
    def __init__(self, input_dim, num_classes, hidden_dim=64):
        super(MLP, self).__init__()
        # 初始化神经网络结构
        self.net = nn.Sequential(
            nn.Linear(input_dim, 128),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(128, num_classes)
        )

    def forward(self, x):
        return self.net(x)


def prepare_data():
    '''
    加载数据并进行预处理
    :return:
    '''
    file_path = os.path.join(data_dir, 'train.csv')
    df = pd.read_csv(file_path)
    # 构建领域化特征
    df = add_agricultural_features(df)
    # 特征和标签
    X = df.drop(columns=['Fertilizer Name', 'id'])

    y = df['Fertilizer Name'].values
    '''
    人工校验数据
    '''
    # 检查分类是否均衡
    print(Counter(y))

    print("原始列名:", X.columns.tolist())

    # 定义预处理器：类别型列做 OneHot，数值型列标准化
    # 明确指定类别型和数值型列  Temparature,Humidity,Moisture,Soil Type,Crop Type,Nitrogen,Potassium,Phosphorous
    categorical_cols = ['Soil Type', 'Crop Type']
    numerical_cols = ['Temparature', 'Humidity', 'Moisture', 'Nitrogen', 'Potassium', 'Phosphorous']  # 替换为你的真实数值列

    preprocessor = ColumnTransformer([
        ('cat', OneHotEncoder(handle_unknown='ignore'), categorical_cols),
        ('num', StandardScaler(), numerical_cols)
    ])

    X_processed = preprocessor.fit_transform(X)

    # 检查OneHot 编码后的稀疏矩阵是否被正确转换为稠密矩阵？
    # 特征中是否有大量缺失值或异常值？
    print("X_processed type:", type(X_processed))
    print("X_processed shape:", X_processed.shape)
    print("X_processed sample:\n", X_processed[:5])

    # 编码标签（虽然你已处理过，但确保是整数形式）
    le = LabelEncoder()
    y = le.fit_transform(y)

    # 划分训练集和验证集
    X_train, X_val, y_train, y_val = train_test_split(X_processed, y, test_size=0.2, random_state=42)

    # 构建 Dataset 和 DataLoader
    train_dataset = FertilizerDataset(X_train, y_train)
    val_dataset = FertilizerDataset(X_val, y_val)

    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
    val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=64)

    return train_loader, val_loader, len(le.classes_)


def train_torch_model():
    train_loader, val_loader, num_classes = prepare_data()
    input_dim = train_loader.dataset.X.shape[1]  # 自动获取特征数量
    # 初始化模型、损失函数和优化器
    model = MLP(input_dim=input_dim, num_classes=num_classes).to(device)  # 根据特征数量调整 input_dim
    # 找一个适合多类别分类的损失函数
    criterion = nn.CrossEntropyLoss()
    # 定义优化器
    # optimizer = optim.Adam(model.parameters(), lr=0.0005)
    optimizer = torch.optim.AdamW(model.parameters(), lr=0.0005)
    # 训练循环
    epochs = 50
    print(f"🚀 开始训练，共 {epochs} 个 epoch")
    # 初始化最佳准确率
    best_acc = 0.0

    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, 'min', patience=3)

    for epoch in range(epochs):
        # 将模型设置为训练模式，以便启用dropout、batch normalization等在训练时需要的特性
        model.train()
        # 初始化总损失为0，用于累计整个训练过程中所有批次的损失值
        total_loss = 0
        # 遍历训练数据加载器，获取每个批次的输入数据和目标标签
        for X_batch, y_batch in train_loader:
            # 将当前批次的输入数据和目标标签移动到指定的设备（如GPU）上
            X_batch, y_batch = X_batch.to(device), y_batch.to(device)
            # 清零优化器的梯度，避免梯度累积影响下一次计算
            optimizer.zero_grad()
            # 将当前批次的输入数据传递给模型，获取模型的输出结果
            outputs = model(X_batch)
            # 计算模型输出结果与目标标签之间的损失值
            loss = criterion(outputs, y_batch)
            # 对损失值进行反向传播，计算模型参数的梯度
            loss.backward()
            # 根据计算出的梯度，使用优化器对模型参数进行一次更新
            optimizer.step()
            # 累加当前批次的损失值到总损失中，用于后续的损失值跟踪和可视化
            total_loss += loss.item()

        # 将模型设置为评估模式，以禁用dropout等仅在训练期间启用的功能
        model.eval()
        # 初始化列表，用于存储真实标签和预测标签
        y_true, y_pred = [], []
        # 使用torch.no_grad()上下文管理器禁用梯度计算，以提高推理速度并减少内存消耗
        with torch.no_grad():
            # 遍历验证数据集中的每个批次
            for X_batch, y_batch in val_loader:
                # 将批次数据移动到指定设备（如GPU）
                X_batch, y_batch = X_batch.to(device), y_batch.to(device)
                # 使用模型进行预测
                outputs = model(X_batch)
                # 通过获取每个输出行中最大值的索引来确定预测标签
                preds = torch.argmax(outputs, dim=1)
                # 打印第一个 batch 的真实标签和预测值
                # print("真实标签:", y_batch.cpu().numpy())
                # print("预测标签:", preds.cpu().numpy())
                # 将真实标签从Tensor转换为numpy数组，并添加到y_true列表中
                y_true.extend(y_batch.cpu().numpy())
                # 将预测标签从Tensor转换为numpy数组，并添加到y_pred列表中
                y_pred.extend(preds.cpu().numpy())

        # 看哪些类别的 precision/recall 很低； 看看是否是学不会某一类还是都学不会
        print(classification_report(y_true, y_pred))
        # 计算模型准确率
        acc = accuracy_score(y_true, y_pred)
        # 打印当前轮次的训练信息，包括轮次、损失值和验证集上的准确率
        print(f"Epoch [{epoch + 1}/{epochs}], Loss: {total_loss:.4f}, Val Acc: {acc:.4f}")

        # 保存最佳模型
        if acc > best_acc:
            best_acc = acc
            # 保存模型的状态字典到指定路径
            torch.save(model.state_dict(), os.path.join(current_dir, 'best_torch_model.pth'))
            print("💾 最佳模型已保存")
            # 写入日志文件（覆盖模式，只保留最新一行）
            with open(log_file_path, 'w') as log_file:
                log_file.write(f"Best Accuracy: {best_acc:.4f}\n")

        scheduler.step(total_loss)
        writer.add_scalar('Loss/train', total_loss, epoch)
        writer.add_scalar('Accuracy/val', acc, epoch)

    print("🎉 训练完成！")
    return model


def tradition_model():
    '''
    为了判断是否是深度学习模型的问题，可以快速测试一个传统分类器（如 RandomForestClassifier）：
    如果传统方法表现良好，则说明数据本身是可学的，问题出在神经网络模型的设计或训练上。
    :return:
    '''

    file_path = os.path.join(data_dir, 'train.csv')
    df = pd.read_csv(file_path)
    df = add_agricultural_features(df)
    # 特征和标签
    X = df.drop(columns=['Fertilizer Name', 'id','Soil Type'])
    y = df['Fertilizer Name'].values
    print(df.describe())
    # 定义预处理器：类别型列做 OneHot，数值型列标准化
    # categorical_cols = ['Soil Type', 'Crop Type']
    categorical_cols = [ 'Crop Type','Sugarcane_Clayey']
    # columns_to_exclude = ['Temparature', 'Humidity', 'Moisture', 'Nitrogen', 'Potassium', 'Phosphorous']
    columns_to_exclude = []

    numerical_cols = [col for col in X.columns if col not in categorical_cols and col not in columns_to_exclude]
    preprocessor = ColumnTransformer([
        ('cat', OneHotEncoder(handle_unknown='ignore'), categorical_cols),
        ('num', StandardScaler(), numerical_cols)
    ])

    X_processed = preprocessor.fit_transform(X)


    # 编码标签（虽然你已处理过，但确保是整数形式）
    le = LabelEncoder()
    y = le.fit_transform(y)

    # 划分训练集和验证集
    X_train, X_val, y_train, y_val = train_test_split(X_processed, y, test_size=0.2, random_state=42)

    # XGBoost 示例
    # 初始化XGBClassifier模型，配置特定的参数以优化模型性能
    print("XGBoost 模型训练中...")
    clf = XGBClassifier(
        n_estimators=500,
        learning_rate=0.1,
        max_depth=5,
        min_child_weight=3,
        gamma=0.1,
        subsample=0.8,
        colsample_bytree=0.7,
        eval_metric='mlogloss',
        use_label_encoder=False,
        tree_method='hist'
    )
    eval_set = [(X_val, y_val)]
    clf.fit(X_train, y_train, early_stopping_rounds=20, eval_set=eval_set, verbose=False)
    print("XGBoost训练结束")
    y_pred = clf.predict(X_val)
    print("Val Accuracy:", accuracy_score(y_val, y_pred))
    print(classification_report(y_val, y_pred))
    import seaborn as sns

    sns.heatmap(confusion_matrix(y_val, y_pred), annot=True, fmt='d', cmap='Blues')
    plt.xlabel('Predicted')
    plt.ylabel('True')
    plt.title('Confusion Matrix')
    plt.show()

# -----------------------------
# 特征构造
# -----------------------------
def add_agricultural_features(df):
    df['Moisture_Squared'] = df['Moisture'] ** 3
    df['Phosphorous_Squared'] =df['Phosphorous']**2
    df['Nitrogen_Squared'] =df['Nitrogen']**2
    df['Temparature_Squared'] =df['Temparature']**2
    df['Humidity_Squared'] =df['Humidity']**2
    df['NPK_Sum'] = df['Nitrogen'] + df['Phosphorous'] + df['Potassium']
    df['N_P_Ratio'] = df['Nitrogen'] / (df['Phosphorous'] + 1e-5)
    df['P_K_Ratio'] = df['Phosphorous'] / (df['Potassium'] + 1e-5)
    df['Env_Index'] = df['Temparature'] * df['Humidity'] * df['Moisture']
    df['Crop_Soil_Interaction'] = df['Crop Type'] + '_' + df['Soil Type']
    crop_soil_preference = {
        ('Wheat', 'Clay'): 1.2,
        ('Maize', 'Loam'): 1.3,
        ('Millets', 'Sandy'): 1.5,
    }

    df['Crop_Soil_Preference'] = df.apply(
        lambda row: crop_soil_preference.get((row['Crop Type'], row['Soil Type']), 1.0), axis=1)
    df['Weighted_Nitrogen'] = df['Nitrogen'] * df['Crop_Soil_Preference']

    return df


if __name__ == '__main__':
    # tradition_model()
    trained_model = train_torch_model()