import torch
import pickle
import numpy as np
from sklearn.decomposition import PCA
import matplotlib.pyplot as plt
import os

#load the pickle file 
def load_pickle_file(pickle_file):
    with open(pickle_file, 'rb') as f:
        data = pickle.load(f)
    return data

#visualization and folder creation
def visualize_pca(ok_modified_data, ng_modified_data):
    # create a folder to store the visualizations
    if not os.path.exists('pca_visualizations'):
        os.makedirs('pca_visualizations')
    
    plt.scatter(ok_modified_data[:, 0], ok_modified_data[:, 1])
    plt.xlabel('PC1')
    plt.ylabel('PC2')
    plt.title('PCA on data_ok')
    plt.savefig('pca_visualizations/data_ok_pc1_pc2.png')
    plt.show()

    plt.scatter(ng_modified_data[:, 0], ng_modified_data[:, 1], c='red')
    plt.xlabel('PC1')
    plt.ylabel('PC2')
    plt.title('PCA on data_ng')
    plt.savefig('pca_visualizations/data_ng_pc1_pc2.png')
    plt.show()

    # plot PC3 against PC4
    plt.scatter(ok_modified_data[:, 2], ok_modified_data[:, 3])
    plt.xlabel('PC3')
    plt.ylabel('PC4')
    plt.title('PCA on data_ok PC3&PC4')
    plt.savefig('pca_visualizations/data_ok_pc3_pc4.png')
    plt.show()

    # plot PC5 against PC6
    plt.scatter(ok_modified_data[:, 4], ok_modified_data[:, 5])
    plt.xlabel('PC5')
    plt.ylabel('PC6')
    plt.title('PCA on data_ok PC5&PC6')
    plt.savefig('pca_visualizations/data_ok_pc5_pc6.png')
    plt.show()

    # plot PC3 against PC4
    plt.scatter(ng_modified_data[:, 2], ng_modified_data[:, 3], c='red')
    plt.xlabel('PC3')
    plt.ylabel('PC4')
    plt.title('PCA on data_ng PC3&PC4')
    plt.savefig('pca_visualizations/data_ng_pc3_pc4.png')
    plt.show()

    # plot PC5 against PC6
    plt.scatter(ng_modified_data[:, 4], ng_modified_data[:, 5], c='red')
    plt.xlabel('PC5')
    plt.ylabel('PC6')
    plt.title('PCA on data_ng PC5&PC6')
    plt.savefig('pca_visualizations/data_ng_pc5_pc6.png')
    plt.show()

    #combine good and non_good 

    plt.scatter(ok_modified_data[:, 0], ok_modified_data[:, 1],color='blue', label='OK')
    plt.scatter(ng_modified_data[:, 0], ng_modified_data[:, 1],color='red', label=' NG')
    plt.legend(labels=['OK', 'NG'])
    plt.title('')
    plt.savefig('pca_visualizations/data_PC1&PC2_good_nongood.png')
    plt.show()


    plt.scatter(ok_modified_data[:, 2], ok_modified_data[:, 3],color='blue', label='OK')
    plt.scatter(ng_modified_data[:, 2], ng_modified_data[:, 3],color='red', label=' NG')
    plt.legend(labels=['OK', 'NG'])
    plt.title('')
    plt.savefig('pca_visualizations/data_PC3&PC4_good_nongood.png')
    plt.show()

    plt.scatter(ok_modified_data[:, 4], ok_modified_data[:, 5],color='blue', label='OK')
    plt.scatter(ng_modified_data[:, 4], ng_modified_data[:, 5],color='red', label=' NG')
    plt.legend(labels=['OK', 'NG'])
    plt.title('')
    plt.savefig('pca_visualizations/data_PC5&PC6_good_nongood.png')
    plt.show()


#processing with data 

def process_data():
    data_ok = load_pickle_file('../good.pickle')
    print(data_ok.shape)

    data_ng = load_pickle_file('../non_good.pickle')
    print(data_ng.shape)

    # Concatenate the two tensors along the first dimension
    data_0 = torch.cat((data_ok, data_ng), dim=0)
    print(data_0.shape)

    # perform PCA on the data_ng
    pca = PCA(n_components=6)

    # perform PCA fit train on data_0
    pca.fit(data_0)

    # perform PCA transform on data_ok and data_ng
    ok_modified_data = pca.transform(data_ok)
    ng_modified_data = pca.transform(data_ng)

    # print the explained variance ratio of the principal components
    print(pca.explained_variance_ratio_)

    # visualization for each
    visualize_pca(ok_modified_data, ng_modified_data)





def main():
    process_data()

if __name__ == '__main__':
    main()