import os
import pandas as pd
import numpy as np
from util import createXY
from lazypredict.Supervised import LazyClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
import joblib

def train_and_find_best_model(train_folder, dest_folder):
    # 创建特征和标签
    X, y = createXY(train_folder, dest_folder, method='vgg')

    # 划分数据集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    # 将数据转换为DataFrame
    X_train = pd.DataFrame(X_train)
    X_test = pd.DataFrame(X_test)
    y_train = pd.DataFrame(y_train, columns=['label'])
    y_test = pd.DataFrame(y_test, columns=['label'])

    # 使用LazyClassifier自动选择和评估各种分类器
    clf = LazyClassifier(verbose=0, ignore_warnings=True, custom_metric=None)
    models, predictions = clf.fit(X_train, X_test, y_train, y_test)

    # 打印模型性能
    print(models)

    # 获取准确率最高的模型
    best_model_name = models['Accuracy'].idxmax()
    print("\nAccuracy最高的模型是: ", best_model_name)

    # 测试模型
    best_model = clf.models[best_model_name]
    predictions = best_model.predict(X_test)
    # 如果predictions是一维数组，直接使用np.argmax(predictions)
    # 如果predictions是二维数组，使用np.argmax(predictions, axis=1)
    predicted_classes = np.argmax(predictions) if predictions.ndim > 1 else predictions
    accuracy = accuracy_score(y_test, predicted_classes)
    print(f'Best model accuracy: {accuracy:.4f}')

    # 保存最佳模型
    model_save_path = os.path.join(dest_folder, f'best_model_{best_model_name.replace(" ", "_")}.pkl')
    joblib.dump(best_model, model_save_path)
    print(f'Best model saved to: {model_save_path}')


if __name__ == "__main__":
    train_folder = "C:/Users/刘亚翔/Desktop/faiss_dog_cat_question-main/data/train"
    dest_folder = "C:/Users/刘亚翔/Desktop/faiss_dog_cat_question-main"
    train_and_find_best_model(train_folder, dest_folder)