# https://www.globalsino.com/ICs/
# Display the labeled features and a specific pattern

import numpy as np
from ripser import Rips
import pandas as pd
import matplotlib.pyplot as plt

df=pd.read_pickle(r"C:\0Python\wafermap\LSWMD.pkl")
df.info()

df = df.drop(['waferIndex'], axis = 1)
def find_dim(x):
    dim0=np.size(x,axis=0)
    dim1=np.size(x,axis=1)
    return dim0,dim1
df['waferMapDim']=df.waferMap.apply(find_dim)
df.sample(5)
uni_waferDim=np.unique(df.waferMapDim, return_counts=True)
uni_waferDim[0].shape[0]
df['failureNum']=df.failureType
df['trainTestNum']=df.trianTestLabel
mapping_type={'Center':0,'Donut':1,'Edge-Loc':2,'Edge-Ring':3,'Loc':4,'Random':5,'Scratch':6,'Near-full':7,'none':8}
mapping_traintest={'Training':0.7,'Test':0.3}
df=df.replace({'failureNum':mapping_type, 'trainTestNum':mapping_traintest})

tol_wafers = df.shape[0]
tol_wafers

df_withlabel = df[(df['failureNum']>=0) & (df['failureNum']<=8)]
df_withlabel =df_withlabel.reset_index()
df_withpattern = df[(df['failureNum']>=0) & (df['failureNum']<=7)]
df_withpattern = df_withpattern.reset_index()
df_nonpattern = df[(df['failureNum']==8)]
df_withlabel.shape[0], df_withpattern.shape[0], df_nonpattern.shape[0]

x = [9, 340, 3, 16, 0, 25, 84, 37]
labels2 = ['Center','Donut','Edge-Loc','Edge-Ring','Loc','Random','Scratch','Near-full']

fig, ax = plt.subplots(nrows = 2, ncols = 4, figsize=(20, 10))
ax = ax.ravel(order='C')
for i in range(8):
    img = df_withpattern.waferMap[x[i]]
    ax[i].imshow(img)
    ax[i].set_title(df_withpattern.failureType[x[i]][0][0],fontsize=24)
    ax[i].set_xticks([])
    ax[i].set_yticks([])
plt.tight_layout()
plt.show()

data = df_withpattern.waferMap[2]
data = np.where(data == 2, 0, data)
data = np.where(data == 1, 1, data)

for line in data:
    print(line)

fig, ax = plt.subplots()
ax.imshow(data)
plt.show()

rips = Rips()
diagrams = rips.fit_transform(df_withpattern.waferMap[340])
rips.plot(diagrams)
print(diagrams[1])

rips = Rips()
diagrams = rips.fit_transform(data)
rips.plot(diagrams)
print(diagrams[1])

y = np.where(df_withpattern.failureType == "Donut")[0] # 0: image
print(y)

fig, ax = plt.subplots(nrows = 2, ncols = 4, figsize=(10, 10))
ax = ax.ravel(order='C')
for i in range(8):
    img = df_withpattern.waferMap[y[i]]
    ax[i].imshow(img)
    ax[i].set_title(df_withpattern.failureType[y[i]][0][0],fontsize=24)
    ax[i].set_xticks([])
    ax[i].set_yticks([])
plt.tight_layout()
plt.show()