# 0. 引入必要的包
import os                             # os，用于处理文件和目录路径
import yaml                           # yaml，用于读取配置文件
import numpy as np                    # NumPy，用于科学计算
from skimage import io                # skimage的io，用于图像读取
from sklearn.model_selection import train_test_split  # 用于数据集的划分
from util import preprocess_image, save_file, get_config_value  # 引入预处理、保存和配置文件读取的方法
import glob                           # glob，用于文件路径名的匹配
import cv2                            # OpenCV，用于图像处理
from tqdm import tqdm                 # tqdm，用于显示进度条

# 1. 读取配置文件中的信息
train_dir = get_config_value("train")  # 获取 训练数据路径
char_styles = get_config_value("char_styles")  # 获取 字符样式列表，注意: 必须是列标
new_size = get_config_value("new_size")  # 获取 新图像大小元组, 注意: 必须包含h和w

# 2. 生成X,y
print("读取训练数据并进行预处理...")
X = []
y = []

for style in char_styles:
    style_dir = os.path.join(train_dir, f"train_{style}_*")
    image_files = glob.glob(style_dir)

    for file in tqdm(image_files, desc=f"处理 {style} 图像"):
        img = preprocess_image(file, new_size)
        X.append(img)
        y.append(char_styles.index(style))

X = np.array(X)
y = np.array(y)

# 3. 分割测试集和训练集
print("将数据按 80% 和 20% 的比例分割...")
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)

# 4. 打印样本维度和类型信息
print("X_train: ", X_train.shape, X_train.dtype)  # 训练集特征的维度和类型
print("X_test: ", X_test.shape, X_test.dtype)  # 测试集特征的维度和类型
print("y_train: ", y_train.shape, y_train.dtype)  # 训练集标签的维度和类型
print("y_test: ", y_test.shape, y_test.dtype)  # 测试集标签的维度和类型

# 5. 序列化分割后的训练和测试样本
save_file((X_train, X_test, y_train, y_test), "(X_train, X_test, y_train, y_test)", f'{get_config_value("Xy_root")}/Xy')