import numpy as np
import struct
import matplotlib.pyplot as plt
import os
def read_binary_file(file_name, fmt):
    """Read binary file"""
    with open(file_name, "rb") as f:
        content = f.read()
        head = struct.unpack_from(fmt, content, 0)
        offset = struct.calcsize(fmt)
        num = head[1]
        if fmt == '>IIII':
            bits = num * head[2] * head[3]
        else:
            bits = num
        bits_string = '>' + str(bits) + 'B'
        data = struct.unpack_from(bits_string, content, offset)
        data = np.array(data).reshape(num, -1)
    return data

def read_data(image_file, label_file):
    """Read images and labels from MNIST datasets"""
    images = read_binary_file(image_file, '>IIII')
    labels = read_binary_file(label_file, '>II')
    return images, labels

def test_k_values(train_x, train_y, test_x, test_y, algo = 'knn'):
    """Testing the Effect of Different K Values on the Performance of Classifiers"""
    results = []
    k_values = [i for i in range(5, 15)]
    from sklearn import metrics
    if not os.path.exists(f'{algo}/'):
        os.mkdir(f'{algo}')
    for k in k_values:
        # calculate AUC and ROC curve
        plt.title(f'{algo}_params_k_{k}_AOC_Curve')
        plt.legend(loc='lower right')
        plt.xlim([0, 1])
        plt.ylim([0, 1])
        plt.ylabel('True Positive Rate')
        plt.xlabel('False Positive Rate')
        if algo == 'knn':
            from sklearn.neighbors import KNeighborsClassifier
            model = KNeighborsClassifier(n_neighbors=k)
        else:
            from sklearn.ensemble import RandomForestClassifier
            model = RandomForestClassifier(n_estimators=k)
        model.fit(train_x, train_y)
        y_pred = model.predict_proba(test_x)
        y_pred_s = model.predict(test_x)
        # auc = metrics.roc_auc_score(test_y.reshape(-1), y_pred, multi_class = 'ovo')
        accuracy = metrics.accuracy_score(test_y, y_pred_s)
        results.append(accuracy)
        for i in range(y_pred.shape[1]):
            fpr, tpr, thresholds = metrics.roc_curve(test_y, y_pred[:, i], pos_label=i)
            # Draw ROC curve
            plt.plot(fpr, tpr, label=f"ROC Curve {i}", linestyle='solid')
        plt.legend()
        plt.savefig(f"{algo}/{algo}_param_{k}_AOC.png")
        plt.clf()
    plt.title(f'{algo}_accuracy')
    plt.ylim([0, 1])
    plt.ylabel('accuracy')
    plt.xlabel('K')
    plt.plot(k_values, results, marker='D', label="accuracy", linestyle='solid')
    for x1, y1 in zip(k_values, results):
        plt.text(x1, y1, str(y1), ha='center', va='bottom', fontsize=10)
    plt.legend()
    plt.savefig(f"{algo}/{algo}_accuracy.png")
    plt.clf()


if __name__ == "__main__":
    """main function"""
    train_image_file = "dataset/train-images-idx3-ubyte"
    test_image_file = "dataset/t10k-images-idx3-ubyte"
    train_label_file = "dataset/train-labels-idx1-ubyte"
    test_label_file = "dataset/t10k-labels-idx1-ubyte"
    train_x, train_y = read_data(train_image_file, train_label_file)
    test_x, test_y = read_data(test_image_file, test_label_file)
    from sklearn.preprocessing import MinMaxScaler
    scaler_1, scaler_2 = MinMaxScaler(), MinMaxScaler()
    scaler_1, scaler_2 = scaler_1.fit(train_x), scaler_2.fit(test_x)
    train_x, test_x = scaler_1.transform(train_x), scaler_2.transform(test_x)
    test_k_values(train_x, train_y, test_x, test_y, algo='random_forest')
    test_k_values(train_x, train_y, test_x, test_y, algo = 'knn')