### copyright: Takabe, K., Hirayama, T., Yawata, Y., 2020###
##Principal Component Analysis

import os
import glob
import matplotlib.pyplot as plt
from skimage import io
from sklearn.decomposition import PCA
from sklearn.preprocessing import StandardScaler
import numpy as np
import tkinter, tkinter.filedialog, tkinter.messagebox

root = tkinter.Tk()
root.withdraw()
root.attributes('-topmost',True)
iDir = os.path.abspath(os.path.dirname(__file__))
tkinter.messagebox.showinfo('Select','Select target directory')

dirname = tkinter.filedialog.askdirectory(initialdir = iDir)

Path = dirname
files = glob.glob(Path+'/*/*')
I = io.imread(files[0])
Ex, Em = I.shape     ## Ex : The number of excitation laser, Em : The number of detectors

images = np.ndarray((len(files), Ex, Em), dtype = np.uint16)
labels = []
for idx, file in enumerate(files):
   image = io.imread(file)
   images[idx] = image 
   label = os.path.split(os.path.dirname(file))[-1]
   labels.append(label)         
data = images.reshape(-1,Ex*Em)
Labels = list(sorted(set(labels)))

scaler=StandardScaler()
scaler.fit(data)
data2=scaler.transform(data)

pca = PCA(n_components=2)
pca.fit(data2)
data2D = pca.transform(data2)

X=[]
Y=[]
for i, j in enumerate(data2D):
    if labels[i]==Labels[0]:
        X.append(j)
    else:
        Y.append(j)
X=np.array(X)
Y=np.array(Y)

plt.scatter(X[:,0], X[:,1], c='red',label=Labels[0], alpha=0.5)
plt.scatter(Y[:,0], Y[:,1], c='blue',label=Labels[1], alpha=0.5)
plt.xlabel('PC1(%s)'%(round(pca.explained_variance_ratio_[0]*100)))
plt.ylabel('PC2(%s)'%(round(pca.explained_variance_ratio_[1]*100)))
plt.legend()
plt.savefig(Path + '/PCA.png')
plt.show()